feat: 替换为 project_7 验证通过的 wudixzy stack

project_7 docker build 已在竞赛平台验证成功。
完整搬运 wudixzy/competition stack:
- qwen3_5.py 2615 行 (12 个 corex .so 调用)
- patch_ops.sh 251 行 (set -eo pipefail + cd dirname)
- 12 prebuilt corex .so (SHA256 verified)
- 13 CUDA .cu 源码 + 11 build scripts
- 9 vendor overrides (block/sampler/scheduler)
- transformers-4.55.3 offline wheel
- computility-run.yaml: 262144 max-model-len, BI100 env vars
- Dockerfile 结构不变 (COPY qwen3_6_scripts + RUN patch_ops.sh)

Dockerfile

@@ -1,12 +1,9 @@
 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 ; \

computility-run.yaml

@@ -8,39 +8,37 @@ command:
     - --served-model-name
     - llm
     - --max-model-len
-    - '100000'
+    - '262144'
     - --gpu-memory-utilization
-    - '0.90'
+    - '0.9'
     - --trust-remote-code
     - -tp
     - '4'
     - --max-num-seqs
-    - '2'
+    - '1'
     - --disable-log-requests
     - --disable-frontend-multiprocessing
-    - --enforce-eager
+    - --max-num-batched-tokens
+    - '8192'
+    - --enable-chunked-prefill
+    - --max-seq-len-to-capture
+    - '32768'
     - --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'
+      value: 3600
+    - name: BI100_MOE_COREX_DIRECT_ROUTED
+      value: 1
+    - name: BI100_GDN_COREX_PACKED_DECODE
+      value: 1
+    - name: BI100_HYBRID_KV_ACCOUNTING
+      value: full_attention
+    - name: BI100_GDN_CACHE_POLICY
+      value: admission64
+    - name: BI100_GDN_RESTORE_MODE
+      value: hybrid64

qwen3_6_scripts/api_server.py

@@ -6,13 +6,365 @@ import os
 import regex as re
 import signal
 import socket
+import sys
 import tempfile
+import time
 from argparse import Namespace
 from contextlib import asynccontextmanager
 from functools import partial
 from http import HTTPStatus
 from typing import AsyncIterator, Set
 
+
+def _bi100_field(value, name):
+    if isinstance(value, dict):
+        return value.get(name)
+    return getattr(value, name, None)
+
+
+def _bi100_scalar(value):
+    return getattr(value, "value", value)
+
+
+def _bi100_tool_choice_kind(value):
+    value = _bi100_scalar(value)
+    if value is None:
+        return "unset"
+    if isinstance(value, str):
+        return value if value in ("none", "auto", "required") else "other"
+    function = _bi100_field(value, "function")
+    if function is not None and isinstance(
+            _bi100_field(function, "name"), str):
+        return "named"
+    return "other"
+
+
+def _bi100_image_source_kind(value):
+    if not isinstance(value, str):
+        return "other"
+    prefix = value[:8].lower()
+    if prefix.startswith("data:"):
+        return "data"
+    if prefix.startswith(("http://", "https://")):
+        return "remote"
+    return "other"
+
+
+def _bi100_chat_4xx_reason(message):
+    if message == "messages must contain at least one message":
+        return "empty_messages"
+    if (isinstance(message, str)
+            and message.startswith("top_p must be in (0, 1], got ")):
+        return "invalid_top_p"
+    if (isinstance(message, str)
+            and message.startswith("max_tokens must be at least 1, got ")):
+        return "invalid_max_tokens"
+    if (isinstance(message, str)
+            and message.startswith("This model's maximum context length is ")
+            and "tokens. However, you requested " in message):
+        return "context_length_exceeded"
+    if (isinstance(message, str) and message.startswith("n=")
+            and " exceeds max_num_seqs=" in message):
+        return "n_exceeds_max_num_seqs"
+    if message == 'tool_choice = "required" is not supported!':
+        return "unsupported_tool_choice_required"
+    if (isinstance(message, str)
+            and message.startswith('"auto" tool choice requires ')):
+        return "tool_parser_unavailable"
+    if message == "Tool call arguments are not valid JSON.":
+        return "invalid_tool_arguments_json"
+    if (isinstance(message, str)
+            and message.startswith("Tool call arguments must ")):
+        return "invalid_tool_arguments_type"
+    if (isinstance(message, str)
+            and (
+                (message.startswith("At most ")
+                 and " image(s) may be provided in one request." in message)
+                or (message.startswith("You set image=")
+                    and "items in the same prompt." in message))):
+        return "image_count_limit"
+    if message == "Unknown model type: qwen3_5_moe":
+        return "image_model_type_unsupported"
+    return "unclassified_chat_error"
+
+
+def _bi100_chat_request_shape(request):
+    messages = _bi100_field(request, "messages")
+    if not isinstance(messages, (list, tuple)):
+        messages = ()
+    tools = _bi100_field(request, "tools")
+    if not isinstance(tools, (list, tuple)):
+        tools = ()
+
+    system_count = 0
+    system_part_message_count = 0
+    system_text_part_count = 0
+    system_other_part_count = 0
+    tool_message_count = 0
+    assistant_tool_message_count = 0
+    image_count = 0
+    image_data_count = 0
+    image_remote_count = 0
+    image_other_count = 0
+    for message in messages:
+        role = _bi100_scalar(_bi100_field(message, "role"))
+        if role == "system":
+            system_count += 1
+        elif role == "tool":
+            tool_message_count += 1
+        elif (role == "assistant"
+              and _bi100_field(message, "tool_calls")):
+            assistant_tool_message_count += 1
+        content = _bi100_field(message, "content")
+        if not isinstance(content, (list, tuple)):
+            continue
+        if role == "system":
+            system_part_message_count += 1
+        for part in content:
+            part_type = _bi100_scalar(_bi100_field(part, "type"))
+            if role == "system":
+                if part_type == "text":
+                    system_text_part_count += 1
+                else:
+                    system_other_part_count += 1
+            if part_type in ("image", "image_url"):
+                image_count += 1
+                image_url = _bi100_field(part, "image_url")
+                source_kind = _bi100_image_source_kind(
+                    _bi100_field(image_url, "url"))
+                if source_kind == "data":
+                    image_data_count += 1
+                elif source_kind == "remote":
+                    image_remote_count += 1
+                else:
+                    image_other_count += 1
+
+    strict_false_count = 0
+    strict_true_count = 0
+    for tool in tools:
+        function = _bi100_field(tool, "function")
+        strict = _bi100_field(function, "strict")
+        if strict is False:
+            strict_false_count += 1
+        elif strict is True:
+            strict_true_count += 1
+
+    n = _bi100_field(request, "n")
+    return {
+        "message_count": len(messages),
+        "system_count": system_count,
+        "system_part_message_count": system_part_message_count,
+        "system_text_part_count": system_text_part_count,
+        "system_other_part_count": system_other_part_count,
+        "tool_count": len(tools),
+        "tool_message_count": tool_message_count,
+        "assistant_tool_message_count": assistant_tool_message_count,
+        "strict_false_count": strict_false_count,
+        "strict_true_count": strict_true_count,
+        "tool_choice_kind": _bi100_tool_choice_kind(
+            _bi100_field(request, "tool_choice")),
+        "image_count": image_count,
+        "image_data_count": image_data_count,
+        "image_remote_count": image_remote_count,
+        "image_other_count": image_other_count,
+        "has_image": image_count > 0,
+        "stream": bool(_bi100_field(request, "stream")),
+        "n": n if isinstance(n, int) else None,
+    }
+
+
+def _bi100_validation_message_reason(error, tool_choice_kind):
+    if not isinstance(error, dict):
+        return None
+
+    messages = []
+    context = error.get("ctx")
+    if isinstance(context, dict):
+        context_error = context.get("error")
+        if isinstance(context_error, ValueError):
+            messages.append(str(context_error))
+
+    message = error.get("msg")
+    if isinstance(message, str):
+        if message.startswith("Value error, "):
+            message = message.removeprefix("Value error, ")
+        messages.append(message)
+
+    for message in messages:
+        if message == "Tool call arguments are not valid JSON.":
+            return "invalid_tool_arguments_json"
+        if message in (
+                "Tool call arguments must decode to a JSON object.",
+                "Tool call arguments must be a JSON object or a "
+                "JSON-encoded object string."):
+            return "invalid_tool_arguments_type"
+        if message == (
+                "`tool_choice` must be a named tool, \"auto\", or \"none\"."):
+            if tool_choice_kind == "required":
+                return "unsupported_tool_choice_required"
+            return "request_validation_tool_choice"
+    return None
+
+
+def _bi100_validation_reason(errors, request_shape=None):
+    categories = set()
+    message_categories = set()
+    tool_choice_kind = (
+        request_shape.get("tool_choice_kind")
+        if isinstance(request_shape, dict) else None
+    )
+    validation_errors = errors if isinstance(errors, (list, tuple)) else ()
+    for error in validation_errors:
+        if not isinstance(error, dict):
+            continue
+        message_category = _bi100_validation_message_reason(
+            error, tool_choice_kind)
+        if message_category is not None:
+            message_categories.add(message_category)
+        location = error.get("loc")
+        if not isinstance(location, (list, tuple)):
+            continue
+        fields = [
+            value for value in location
+            if isinstance(value, str)
+            and value not in ("body", "query", "path")
+        ]
+        if not fields:
+            continue
+        field = fields[0]
+        descendants = set(fields[1:])
+        if field == "messages":
+            if "tool_call_id" in descendants:
+                categories.add("request_validation_message_tool_call_id")
+            elif "tool_calls" in descendants:
+                categories.add("request_validation_message_tool_calls")
+            elif "content" in descendants:
+                categories.add("request_validation_message_content")
+            elif "role" in descendants:
+                categories.add("request_validation_message_role")
+            else:
+                categories.add("request_validation_messages")
+        elif field == "tools":
+            if "strict" in descendants:
+                categories.add("request_validation_tool_strict")
+            elif "parameters" in descendants:
+                categories.add("request_validation_tool_parameters")
+            else:
+                categories.add("request_validation_tools")
+        elif field in ("tool_choice", "parallel_tool_calls"):
+            categories.add("request_validation_tool_choice")
+        elif field == "response_format":
+            categories.add("request_validation_response_format")
+        elif field in ("stream", "stream_options"):
+            categories.add("request_validation_streaming")
+        elif field in ("n", "max_tokens", "min_tokens", "stop"):
+            categories.add("request_validation_generation")
+        elif field in (
+                "temperature", "top_p", "top_k", "frequency_penalty",
+                "presence_penalty", "repetition_penalty", "seed"):
+            categories.add("request_validation_sampling")
+        elif field == "model":
+            categories.add("request_validation_model")
+        else:
+            categories.add("request_validation_other")
+
+    priority = (
+        "request_validation_tool_strict",
+        "request_validation_tool_parameters",
+        "request_validation_tool_choice",
+        "request_validation_message_tool_call_id",
+        "request_validation_message_tool_calls",
+        "request_validation_message_content",
+        "request_validation_message_role",
+        "request_validation_messages",
+        "request_validation_tools",
+        "request_validation_response_format",
+        "request_validation_streaming",
+        "request_validation_generation",
+        "request_validation_sampling",
+        "request_validation_model",
+        "request_validation_other",
+    )
+    for category in priority:
+        if category in categories:
+            return category
+    message_priority = (
+        "invalid_tool_arguments_json",
+        "invalid_tool_arguments_type",
+        "unsupported_tool_choice_required",
+        "request_validation_tool_choice",
+    )
+    for category in message_priority:
+        if category in message_categories:
+            return category
+    return "request_validation_unknown"
+
+
+def _bi100_validation_identifier(value):
+    if not isinstance(value, str) or not value or len(value) > 64:
+        return "unknown"
+    if not value.isascii():
+        return "unknown"
+    if not all(character.isalnum() or character in "._-"
+               for character in value):
+        return "unknown"
+    return value
+
+
+def _bi100_validation_diagnostics(errors):
+    if not isinstance(errors, (list, tuple)):
+        return "unknown", "unknown"
+    try:
+        error_count = len(errors)
+    except Exception:
+        return "unknown", "unknown"
+    if error_count > 1:
+        return "multiple", "multiple"
+    if error_count == 0:
+        return "unknown", "unknown"
+
+    try:
+        error = errors[0]
+        if not isinstance(error, dict):
+            return "unknown", "unknown"
+        location = error.get("loc")
+        validation_type = _bi100_validation_identifier(error.get("type"))
+        if not isinstance(location, (list, tuple)):
+            return "unknown", validation_type
+        if not location:
+            return "root", validation_type
+        index = 0
+        if location[0] in ("body", "query", "path", "header", "cookie"):
+            index = 1
+        if index >= len(location):
+            return "root", validation_type
+        field = location[index]
+        if field in ("__root__", "root"):
+            return "root", validation_type
+        return _bi100_validation_identifier(field), validation_type
+    except Exception:
+        return "unknown", "unknown"
+
+
+def _bi100_safe_validation_errors(exc):
+    try:
+        errors = exc.errors()
+        if not isinstance(errors, (list, tuple)):
+            return ()
+        return tuple(errors)
+    except Exception:
+        return ()
+
+
+def _bi100_startup_trace(message: str) -> None:
+    if os.getenv("BI100_EXECUTOR_STARTUP_DEBUG") == "1":
+        stamp = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
+        print(f"[BI100 STARTUP] {stamp} pid={os.getpid()} {message}",
+              file=sys.stderr, flush=True)
+
+
+_bi100_startup_trace("api_server stdlib imports complete; loading runtime dependencies")
+
 import uvloop
 from fastapi import APIRouter, FastAPI, Request
 from fastapi.exceptions import RequestValidationError
@@ -70,6 +422,124 @@ logger = init_logger('vllm.entrypoints.openai.api_server')
 
 _running_tasks: Set[asyncio.Task] = set()
 
+_bi100_startup_trace("api_server runtime imports complete")
+
+
+def _bi100_log_chat_4xx(request, error) -> None:
+    code = getattr(error, "code", None)
+    if not isinstance(code, int) or not 400 <= code < 500:
+        return
+    shape = _bi100_chat_request_shape(request)
+    reason = _bi100_chat_4xx_reason(getattr(error, "message", None))
+    logger.warning(
+        "[BI100 4XX] endpoint=chat code=%d reason=%s messages=%d "
+        "systems=%d system_part_msgs=%d system_text_parts=%d "
+        "system_other_parts=%d tools=%d tool_msgs=%d "
+        "assistant_tool_msgs=%d strict_false=%d strict_true=%d choice=%s "
+        "images=%d image_data=%d image_remote=%d image_other=%d "
+        "stream=%d n=%s",
+        code,
+        reason,
+        shape["message_count"],
+        shape["system_count"],
+        shape["system_part_message_count"],
+        shape["system_text_part_count"],
+        shape["system_other_part_count"],
+        shape["tool_count"],
+        shape["tool_message_count"],
+        shape["assistant_tool_message_count"],
+        shape["strict_false_count"],
+        shape["strict_true_count"],
+        shape["tool_choice_kind"],
+        shape["image_count"],
+        shape["image_data_count"],
+        shape["image_remote_count"],
+        shape["image_other_count"],
+        int(shape["stream"]),
+        shape["n"] if shape["n"] is not None else "unset",
+    )
+
+
+def _bi100_log_request_validation_4xx(raw_request, exc) -> None:
+    validation_errors = ()
+    validation_field = "unknown"
+    validation_type = "unknown"
+    try:
+        validation_errors = _bi100_safe_validation_errors(exc)
+        validation_field, validation_type = (
+            _bi100_validation_diagnostics(validation_errors)
+        )
+        body = getattr(exc, "body", None)
+        url = getattr(raw_request, "url", None)
+        path = getattr(url, "path", "")
+        is_chat_request = (
+            isinstance(path, str)
+            and path.endswith("/v1/chat/completions")
+            and isinstance(body, dict)
+        )
+        shape = (
+            _bi100_chat_request_shape(body) if is_chat_request else None
+        )
+        reason = _bi100_validation_reason(validation_errors, shape)
+        if shape is not None:
+            if (reason == "request_validation_tools"
+                    and shape["strict_true_count"]):
+                reason = "request_validation_tool_strict"
+            logger.warning(
+                "[BI100 4XX] endpoint=request_validation code=400 reason=%s "
+                "messages=%d systems=%d system_part_msgs=%d "
+                "system_text_parts=%d system_other_parts=%d tools=%d "
+                "tool_msgs=%d assistant_tool_msgs=%d strict_false=%d "
+                "strict_true=%d choice=%s images=%d image_data=%d "
+                "image_remote=%d image_other=%d stream=%d n=%s errors=%d "
+                "validation_field=%s validation_type=%s",
+                reason,
+                shape["message_count"],
+                shape["system_count"],
+                shape["system_part_message_count"],
+                shape["system_text_part_count"],
+                shape["system_other_part_count"],
+                shape["tool_count"],
+                shape["tool_message_count"],
+                shape["assistant_tool_message_count"],
+                shape["strict_false_count"],
+                shape["strict_true_count"],
+                shape["tool_choice_kind"],
+                shape["image_count"],
+                shape["image_data_count"],
+                shape["image_remote_count"],
+                shape["image_other_count"],
+                int(shape["stream"]),
+                shape["n"] if shape["n"] is not None else "unset",
+                len(validation_errors),
+                validation_field,
+                validation_type,
+            )
+        else:
+            logger.warning(
+                "[BI100 4XX] endpoint=request_validation code=400 reason=%s "
+                "errors=%d validation_field=%s validation_type=%s",
+                reason,
+                len(validation_errors),
+                validation_field,
+                validation_type,
+            )
+        return
+    except Exception:
+        pass
+
+    try:
+        logger.warning(
+            "[BI100 4XX] endpoint=request_validation code=400 "
+            "reason=request_validation_unknown errors=%d "
+            "validation_field=%s validation_type=%s",
+            len(validation_errors),
+            validation_field,
+            validation_type,
+        )
+    except Exception:
+        pass
+
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
@@ -101,12 +571,15 @@ async def lifespan(app: FastAPI):
 async def build_async_engine_client(
         args: Namespace) -> AsyncIterator[EngineClient]:
 
+    _bi100_startup_trace("building AsyncEngineArgs")
     # Context manager to handle engine_client lifecycle
     # Ensures everything is shutdown and cleaned up on error/exit
     engine_args = AsyncEngineArgs.from_cli_args(args)
 
+    _bi100_startup_trace("entering engine client construction")
     async with build_async_engine_client_from_engine_args(
             engine_args, args.disable_frontend_multiprocessing) as engine:
+        _bi100_startup_trace("engine client construction completed")
         yield engine
 
 
@@ -317,6 +790,7 @@ async def create_chat_completion(request: ChatCompletionRequest,
         request, raw_request)
 
     if isinstance(generator, ErrorResponse):
+        _bi100_log_chat_4xx(request, generator)
         return JSONResponse(content=generator.model_dump(),
                             status_code=generator.code)
 
@@ -430,7 +904,8 @@ def build_app(args: Namespace) -> FastAPI:
     )
 
     @app.exception_handler(RequestValidationError)
-    async def validation_exception_handler(_, exc):
+    async def validation_exception_handler(raw_request, exc):
+        _bi100_log_request_validation_4xx(raw_request, exc)
         chat = app.state.openai_serving_chat
         err = chat.create_error_response(message=str(exc))
         return JSONResponse(err.model_dump(),
@@ -527,6 +1002,7 @@ def init_app_state(
 
 
 async def run_server(args, **uvicorn_kwargs) -> None:
+    _bi100_startup_trace("run_server entered")
     logger.info("vLLM API server version %s", VLLM_VERSION)
     logger.info("args: %s", args)
 
@@ -559,12 +1035,17 @@ async def run_server(args, **uvicorn_kwargs) -> None:
 
     signal.signal(signal.SIGTERM, signal_handler)
 
+    _bi100_startup_trace("starting engine client context")
     async with build_async_engine_client(args) as engine_client:
+        _bi100_startup_trace("building FastAPI application")
         app = build_app(args)
 
+        _bi100_startup_trace("requesting model config from engine")
         model_config = await engine_client.get_model_config()
+        _bi100_startup_trace("model config received; initializing app state")
         init_app_state(engine_client, model_config, app.state, args)
 
+        _bi100_startup_trace("starting HTTP server")
         shutdown_task = await serve_http(
             app,
             host=args.host,
@@ -584,6 +1065,7 @@ async def run_server(args, **uvicorn_kwargs) -> None:
 
 
 if __name__ == "__main__":
+    _bi100_startup_trace("api_server __main__ entered")
     # NOTE(simon):
     # This section should be in sync with vllm/scripts.py for CLI entrypoints.
     parser = FlexibleArgumentParser(
@@ -591,5 +1073,8 @@ if __name__ == "__main__":
     parser = make_arg_parser(parser)
     args = parser.parse_args()
     validate_parsed_serve_args(args)
+    _bi100_startup_trace(
+        f"arguments parsed model={args.model} tp={args.tensor_parallel_size} "
+        f"max_model_len={args.max_model_len}")
 
     uvloop.run(run_server(args))

qwen3_6_scripts/bi100_env.py

@@ -0,0 +1,26 @@
+import os
+
+
+def env_bool(name: str, default: bool = False) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    if raw in ("1", "true", "True", "yes", "YES", "on", "ON"):
+        return True
+    if raw in ("0", "false", "False", "no", "NO", "off", "OFF"):
+        return False
+    raise RuntimeError(f"{name} must be boolean, got {raw!r}")
+
+
+def env_int(name: str, default: int, min_value: int, max_value: int) -> int:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    try:
+        value = int(raw)
+    except ValueError as exc:
+        raise RuntimeError(f"{name} must be int, got {raw!r}") from exc
+    if not (min_value <= value <= max_value):
+        raise RuntimeError(
+            f"{name}={value} outside [{min_value}, {max_value}]")
+    return value

qwen3_6_scripts/bi100_profile.py

@@ -0,0 +1,237 @@
+import contextlib
+import fnmatch
+import functools
+import json
+import os
+import re
+import threading
+import time
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+_EVENT_SCHEMA = "bi100-profile-event-v1"
+_EVENT_VERSION = 1
+_NAME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_.-]{0,63}$")
+_FILTER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9_.*?-]{0,63}$")
+
+
+def _strict_bool(name: str, default: str = "0") -> bool:
+    value = os.getenv(name, default).strip()
+    if value not in {"0", "1"}:
+        raise RuntimeError(f"{name} must be exactly 0 or 1, got {value!r}")
+    return value == "1"
+
+
+_ENABLED = _strict_bool("BI100_PROFILE")
+_INCLUDE_STARTUP = _strict_bool("BI100_PROFILE_INCLUDE_STARTUP")
+_MODE = os.getenv("BI100_PROFILE_MODE", "sync").strip().lower()
+_FILTERS = tuple(
+    item.strip()
+    for item in os.getenv("BI100_PROFILE_FILTER", "").split(",")
+    if item.strip()
+)
+if _ENABLED and _MODE not in {"sync", "event"}:
+    raise RuntimeError(f"unsupported BI100_PROFILE_MODE={_MODE!r}")
+if _ENABLED and any(_FILTER_RE.fullmatch(pattern) is None
+                    for pattern in _FILTERS):
+    raise RuntimeError("BI100_PROFILE_FILTER contains an invalid pattern")
+
+_EVENT_RECORDS = []
+_COUNTERS = {}
+_LOCK = threading.Lock()
+_FORWARD_INDEX = 0
+_LAST_FLUSH_NS = None
+_ACTIVE_FORWARD_TOKEN = None
+_NEXT_FORWARD_TOKEN = 0
+
+
+def _enabled_for(name: str) -> bool:
+    return (_ENABLED
+            and (not _FILTERS
+                 or any(fnmatch.fnmatchcase(name, pattern)
+                        for pattern in _FILTERS)))
+
+
+def _skip_startup() -> bool:
+    return (not _INCLUDE_STARTUP
+            and os.getenv("BI100_IN_STARTUP_PROFILE") == "1")
+
+
+def bi100_profile_event_enabled() -> bool:
+    return _ENABLED and _MODE == "event" and not _skip_startup()
+
+
+def _begin_profile_forward():
+    global _ACTIVE_FORWARD_TOKEN, _NEXT_FORWARD_TOKEN
+    if not bi100_profile_event_enabled():
+        return None
+    with _LOCK:
+        _EVENT_RECORDS.clear()
+        _COUNTERS.clear()
+        token = _NEXT_FORWARD_TOKEN
+        _NEXT_FORWARD_TOKEN += 1
+        _ACTIVE_FORWARD_TOKEN = token
+    return token
+
+
+def _abort_profile_forward(token) -> None:
+    global _ACTIVE_FORWARD_TOKEN
+    if token is None:
+        return
+    with _LOCK:
+        if _ACTIVE_FORWARD_TOKEN != token:
+            return
+        _EVENT_RECORDS.clear()
+        _COUNTERS.clear()
+        _ACTIVE_FORWARD_TOKEN = None
+
+
+def bi100_profile_transaction(function):
+    """Keep one top-level model forward isolated from failed forwards."""
+    @functools.wraps(function)
+    def wrapped(*args, **kwargs):
+        token = _begin_profile_forward()
+        if token is None:
+            return function(*args, **kwargs)
+        try:
+            result = function(*args, **kwargs)
+        except BaseException:
+            _abort_profile_forward(token)
+            raise
+        with _LOCK:
+            was_flushed = _ACTIVE_FORWARD_TOKEN != token
+        if not was_flushed:
+            _abort_profile_forward(token)
+            raise RuntimeError(
+                "BI100 profile transaction completed without a flush")
+        return result
+
+    return wrapped
+
+
+def _normalize_metadata(metadata):
+    normalized = {}
+    for key, value in metadata.items():
+        if not isinstance(key, str) or _NAME_RE.fullmatch(key) is None:
+            raise TypeError("profile metadata keys must be bounded names")
+        if isinstance(value, bool):
+            normalized[key] = value
+        elif isinstance(value, int) and not isinstance(value, bool):
+            normalized[key] = value
+        elif isinstance(value, str) and len(value) <= 64:
+            normalized[key] = value
+        else:
+            raise TypeError(
+                "profile metadata values must be bool, int, or short strings")
+    return normalized
+
+
+def bi100_profile_count(name: str, **metadata) -> None:
+    """Record privacy-safe path metadata for the current model forward."""
+    if not bi100_profile_event_enabled() or not _enabled_for(name):
+        return
+    if not isinstance(name, str) or _NAME_RE.fullmatch(name) is None:
+        raise TypeError("profile counter name must be a bounded name")
+    normalized = _normalize_metadata(metadata)
+    encoded = json.dumps(
+        {"name": name, **normalized}, sort_keys=True, separators=(",", ":"))
+    with _LOCK:
+        _COUNTERS[encoded] = _COUNTERS.get(encoded, 0) + 1
+
+
+@contextlib.contextmanager
+def bi100_timer(name: str):
+    if not _enabled_for(name) or _skip_startup():
+        yield
+        return
+    import torch
+
+    if _MODE == "event":
+        started = torch.cuda.Event(enable_timing=True)
+        finished = torch.cuda.Event(enable_timing=True)
+        host_started_ns = time.monotonic_ns()
+        started.record()
+        try:
+            yield
+        finally:
+            finished.record()
+            with _LOCK:
+                _EVENT_RECORDS.append(
+                    (name, started, finished, host_started_ns))
+        return
+
+    torch.cuda.synchronize()
+    t0 = time.perf_counter()
+    try:
+        yield
+    finally:
+        torch.cuda.synchronize()
+        logger.info("[BI100_PROFILE] %s %.3f ms", name,
+                    (time.perf_counter() - t0) * 1000)
+
+
+def bi100_profile_flush(*, tp_rank, **metadata):
+    """Synchronize once and emit one aggregate event record per model forward."""
+    global _ACTIVE_FORWARD_TOKEN, _FORWARD_INDEX, _LAST_FLUSH_NS
+    if not bi100_profile_event_enabled():
+        return None
+    if (not isinstance(tp_rank, int) or isinstance(tp_rank, bool)
+            or not 0 <= tp_rank < 256):
+        raise TypeError("profile TP rank must be an integer in [0, 255]")
+    normalized_metadata = _normalize_metadata(metadata)
+
+    with _LOCK:
+        records = list(_EVENT_RECORDS)
+        counters = dict(_COUNTERS)
+        _EVENT_RECORDS.clear()
+        _COUNTERS.clear()
+        _ACTIVE_FORWARD_TOKEN = None
+    if not records:
+        return None
+
+    import torch
+
+    torch.cuda.synchronize()
+    flushed_ns = time.monotonic_ns()
+    regions = {}
+    model_started_ns = []
+    for name, started, finished, host_started_ns in records:
+        stats = regions.setdefault(name, {"count": 0, "total_ms": 0.0})
+        stats["count"] += 1
+        stats["total_ms"] += float(started.elapsed_time(finished))
+        if name == "model.forward":
+            model_started_ns.append(host_started_ns)
+
+    counter_rows = []
+    for encoded, count in sorted(counters.items()):
+        row = json.loads(encoded)
+        row["count"] = count
+        counter_rows.append(row)
+
+    first_model_started_ns = (
+        min(model_started_ns) if model_started_ns else None)
+    payload = {
+        "schema": _EVENT_SCHEMA,
+        "version": _EVENT_VERSION,
+        "tp_rank": tp_rank,
+        "forward_index": _FORWARD_INDEX,
+        "metadata": normalized_metadata,
+        "event_count": len(records),
+        "model_forward_event_count": len(model_started_ns),
+        "regions": regions,
+        "counters": counter_rows,
+        "host_model_start_to_flush_ms": (
+            (flushed_ns - first_model_started_ns) / 1_000_000
+            if first_model_started_ns is not None else None),
+        "host_gap_since_previous_flush_ms": (
+            (first_model_started_ns - _LAST_FLUSH_NS) / 1_000_000
+            if first_model_started_ns is not None
+            and _LAST_FLUSH_NS is not None
+            else None),
+    }
+    _FORWARD_INDEX += 1
+    _LAST_FLUSH_NS = flushed_ns
+    logger.info("[BI100_PROFILE_EVENT] %s",
+                json.dumps(payload, sort_keys=True, separators=(",", ":")))
+    return payload

qwen3_6_scripts/block_major_kv_cache.py

@@ -0,0 +1,398 @@
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Mapping
+
+import torch
+
+from vllm.logger import init_logger
+
+
+logger = init_logger(__name__)
+
+ENABLE_ENV = "BI100_BLOCK_MAJOR_CPU_KV"
+TRACE_ENV = "BI100_BLOCK_MAJOR_CPU_KV_TRACE"
+CPU_OFFLOAD_ENV = "BI100_CPU_KV_OFFLOAD"
+HYBRID_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
+NUM_ATTENTION_LAYERS = 10
+KV_PLANES = 2
+ELEMENTS_PER_PLANE_BLOCK = 4096
+STAGING_BLOCKS = 512
+STAGING_BUFFER_COUNT = 2
+BYTES_PER_BLOCK = (
+    NUM_ATTENTION_LAYERS * KV_PLANES * ELEMENTS_PER_PLANE_BLOCK * 2
+)
+GPU_STAGING_BYTES = STAGING_BLOCKS * STAGING_BUFFER_COUNT * BYTES_PER_BLOCK
+
+
+def _strict_binary_selector(
+    name: str,
+    environ: Mapping[str, str] | None = None,
+) -> bool:
+    source = os.environ if environ is None else environ
+    raw = source.get(name, "0")
+    if raw == "0":
+        return False
+    if raw == "1":
+        return True
+    raise RuntimeError(f"{name} must be exactly '0' or '1', got {raw!r}")
+
+
+def block_major_cpu_kv_enabled(
+    environ: Mapping[str, str] | None = None,
+) -> bool:
+    return _strict_binary_selector(ENABLE_ENV, environ)
+
+
+def block_major_cpu_kv_trace_enabled(
+    environ: Mapping[str, str] | None = None,
+) -> bool:
+    return _strict_binary_selector(TRACE_ENV, environ)
+
+
+def _require_block_major_runtime(
+    environ: Mapping[str, str] | None = None,
+) -> None:
+    source = os.environ if environ is None else environ
+    if source.get(CPU_OFFLOAD_ENV, "0") != "1":
+        raise RuntimeError(
+            f"{ENABLE_ENV}=1 requires {CPU_OFFLOAD_ENV}=1")
+    if source.get(HYBRID_ACCOUNTING_ENV, "legacy40") != "full_attention":
+        raise RuntimeError(
+            f"{ENABLE_ENV}=1 requires "
+            f"{HYBRID_ACCOUNTING_ENV}=full_attention")
+
+
+def reserve_block_major_gpu_blocks(
+    num_gpu_blocks: int,
+    cache_block_size: int,
+    environ: Mapping[str, str] | None = None,
+) -> int:
+    if (not isinstance(num_gpu_blocks, int)
+            or isinstance(num_gpu_blocks, bool)
+            or num_gpu_blocks < 0):
+        raise ValueError("num_gpu_blocks must be a non-negative integer")
+    if not block_major_cpu_kv_enabled(environ):
+        return num_gpu_blocks
+
+    _require_block_major_runtime(environ)
+    if cache_block_size != BYTES_PER_BLOCK:
+        raise RuntimeError(
+            f"{ENABLE_ENV}=1 requires cache block size "
+            f"{BYTES_PER_BLOCK}, got {cache_block_size}")
+    reserved_blocks = (
+        GPU_STAGING_BYTES + cache_block_size - 1
+    ) // cache_block_size
+    remaining_blocks = num_gpu_blocks - reserved_blocks
+    if remaining_blocks <= 0:
+        raise RuntimeError(
+            "block-major GPU staging leaves no usable GPU KV blocks")
+    logger.info(
+        "[BI100 BLOCK KV] capacity reserve blocks=%d bytes=%d "
+        "profiled_blocks=%d usable_blocks=%d",
+        reserved_blocks,
+        GPU_STAGING_BYTES,
+        num_gpu_blocks,
+        remaining_blocks,
+    )
+    return remaining_blocks
+
+
+def validate_block_mapping(
+    mapping: torch.Tensor,
+    source_limit: int,
+    destination_limit: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if not isinstance(mapping, torch.Tensor):
+        raise TypeError("block mapping must be a torch.Tensor")
+    if mapping.device.type != "cpu":
+        raise ValueError("block mapping must be on CPU")
+    if mapping.dtype != torch.int64:
+        raise ValueError("block mapping must use torch.int64")
+    if not mapping.is_contiguous():
+        raise ValueError("block mapping must be contiguous")
+    if mapping.dim() != 2 or mapping.shape[1] != 2:
+        raise ValueError("block mapping must have shape [N, 2]")
+    if source_limit <= 0 or destination_limit <= 0:
+        raise ValueError("block mapping limits must be positive")
+
+    sources: set[int] = set()
+    destinations: set[int] = set()
+    for row, pair in enumerate(mapping.tolist()):
+        source, destination = pair
+        if not 0 <= source < source_limit:
+            raise ValueError(
+                f"source block out of range at row {row}: {source}")
+        if not 0 <= destination < destination_limit:
+            raise ValueError(
+                f"destination block out of range at row {row}: "
+                f"{destination}")
+        if source in sources:
+            raise ValueError(f"duplicate source block: {source}")
+        if destination in destinations:
+            raise ValueError(f"duplicate destination block: {destination}")
+        sources.add(source)
+        destinations.add(destination)
+
+    return mapping[:, 0].contiguous(), mapping[:, 1].contiguous()
+
+
+class BlockMajorCpuKVCache:
+
+    def __init__(
+        self,
+        gpu_cache: list[torch.Tensor],
+        num_cpu_blocks: int,
+        pin_memory: bool,
+    ) -> None:
+        self._validate_gpu_cache(gpu_cache)
+        if block_major_cpu_kv_enabled():
+            _require_block_major_runtime()
+        if num_cpu_blocks <= 0:
+            raise RuntimeError(
+                f"{ENABLE_ENV}=1 requires a positive CPU block count")
+        if not pin_memory:
+            raise RuntimeError(
+                f"{ENABLE_ENV}=1 requires pinned CPU memory")
+
+        try:
+            from vllm import corex_block_major_kv_transfer as extension
+        except ImportError as exc:
+            raise RuntimeError(
+                "block-major CoreX extension is unavailable") from exc
+
+        self.extension = extension
+        self.gpu_cache = gpu_cache
+        self.device = gpu_cache[0].device
+        self.dtype = gpu_cache[0].dtype
+        self.num_gpu_blocks = gpu_cache[0].shape[1]
+        self.num_cpu_blocks = num_cpu_blocks
+        self.trace_enabled = block_major_cpu_kv_trace_enabled()
+
+        self.cpu_pool = torch.zeros(
+            (
+                num_cpu_blocks,
+                NUM_ATTENTION_LAYERS,
+                KV_PLANES,
+                ELEMENTS_PER_PLANE_BLOCK,
+            ),
+            dtype=self.dtype,
+            device="cpu",
+            pin_memory=True,
+        )
+        if not self.cpu_pool.is_pinned():
+            raise RuntimeError("block-major CPU pool is not pinned")
+
+        # Preserve the public CacheEngine shape without allocating a second
+        # layer-major CPU cache. Transfer methods use cpu_pool directly.
+        self.layer_views = [
+            self.cpu_pool[:, layer, :, :].permute(1, 0, 2)
+            for layer in range(NUM_ATTENTION_LAYERS)
+        ]
+        self.cpu_staging = [
+            torch.empty(
+                (
+                    STAGING_BLOCKS,
+                    NUM_ATTENTION_LAYERS,
+                    KV_PLANES,
+                    ELEMENTS_PER_PLANE_BLOCK,
+                ),
+                dtype=self.dtype,
+                device="cpu",
+                pin_memory=True,
+            )
+            for _ in range(STAGING_BUFFER_COUNT)
+        ]
+        if not all(staging.is_pinned() for staging in self.cpu_staging):
+            raise RuntimeError("block-major CPU staging is not pinned")
+
+        with torch.cuda.device(self.device):
+            self.gpu_staging = [
+                torch.empty_like(staging, device=self.device)
+                for staging in self.cpu_staging
+            ]
+            self.events = [
+                torch.cuda.Event(enable_timing=False)
+                for _ in range(STAGING_BUFFER_COUNT)
+            ]
+            self.error_flag = torch.zeros(
+                1, dtype=torch.int32, device=self.device)
+
+        logger.info(
+            "[BI100 BLOCK KV] enabled device=%s gpu_blocks=%d cpu_blocks=%d "
+            "layers=%d block_bytes=%d staging_blocks=%d staging_buffers=%d",
+            self.device,
+            self.num_gpu_blocks,
+            self.num_cpu_blocks,
+            NUM_ATTENTION_LAYERS,
+            BYTES_PER_BLOCK,
+            STAGING_BLOCKS,
+            STAGING_BUFFER_COUNT,
+        )
+
+    @staticmethod
+    def _validate_gpu_cache(gpu_cache: list[torch.Tensor]) -> None:
+        if len(gpu_cache) != NUM_ATTENTION_LAYERS:
+            raise RuntimeError(
+                f"{ENABLE_ENV}=1 requires exactly "
+                f"{NUM_ATTENTION_LAYERS} GPU attention caches, got "
+                f"{len(gpu_cache)}")
+        first = gpu_cache[0]
+        if first.device.type != "cuda":
+            raise RuntimeError("block-major GPU cache must be on CUDA")
+        if first.dtype != torch.float16:
+            raise RuntimeError("block-major GPU cache must use float16")
+        if (first.dim() != 3 or first.shape[0] != KV_PLANES
+                or first.shape[2] != ELEMENTS_PER_PLANE_BLOCK):
+            raise RuntimeError(
+                "block-major GPU cache must have shape [2, blocks, 4096]")
+        if not first.is_contiguous():
+            raise RuntimeError("block-major GPU cache must be contiguous")
+
+        for layer, tensor in enumerate(gpu_cache):
+            if tensor.device != first.device:
+                raise RuntimeError(
+                    f"GPU cache layer {layer} is on a different device")
+            if tensor.dtype != first.dtype or tensor.shape != first.shape:
+                raise RuntimeError(
+                    f"GPU cache layer {layer} has inconsistent geometry")
+            if not tensor.is_contiguous():
+                raise RuntimeError(
+                    f"GPU cache layer {layer} is not contiguous")
+
+    def _to_gpu_ids(self, block_ids: torch.Tensor) -> torch.Tensor:
+        return block_ids.to(
+            device=self.device,
+            dtype=torch.int32,
+            non_blocking=False,
+        )
+
+    @staticmethod
+    def _chunks(
+        source: torch.Tensor,
+        destination: torch.Tensor,
+        gpu_ids: torch.Tensor,
+    ):
+        for start in range(0, source.numel(), STAGING_BLOCKS):
+            end = min(start + STAGING_BLOCKS, source.numel())
+            yield (
+                source[start:end],
+                destination[start:end],
+                gpu_ids[start:end],
+                end - start,
+            )
+
+    def _begin(self) -> None:
+        self.error_flag.zero_()
+
+    def _finish(
+        self,
+        direction: str,
+        block_count: int,
+        started: float | None,
+    ) -> None:
+        # check_error performs the final stream synchronization. This also
+        # makes every staging slot safe to reuse in the next CacheEngine call.
+        self.extension.check_error(self.error_flag)
+        if started is not None:
+            elapsed_ms = (time.perf_counter() - started) * 1000.0
+            logger.info(
+                "[BI100 BLOCK KV TRACE] direction=%s blocks=%d bytes=%d "
+                "elapsed_ms=%.3f",
+                direction,
+                block_count,
+                block_count * BYTES_PER_BLOCK,
+                elapsed_ms,
+            )
+
+    def swap_out(self, mapping: torch.Tensor) -> None:
+        started = time.perf_counter() if self.trace_enabled else None
+        source_gpu, destination_cpu = validate_block_mapping(
+            mapping,
+            source_limit=self.num_gpu_blocks,
+            destination_limit=self.num_cpu_blocks,
+        )
+        block_count = source_gpu.numel()
+        if block_count == 0:
+            return
+        source_gpu_ids = self._to_gpu_ids(source_gpu)
+
+        self._begin()
+        pending: tuple[int, torch.Tensor, int] | None = None
+        for index, (_, destination, gpu_ids, count) in enumerate(
+                self._chunks(
+                    source_gpu, destination_cpu, source_gpu_ids)):
+            slot = index % STAGING_BUFFER_COUNT
+            self.extension.pack(
+                self.gpu_cache,
+                gpu_ids,
+                self.gpu_staging[slot],
+                self.error_flag,
+                count,
+            )
+            self.cpu_staging[slot][:count].copy_(
+                self.gpu_staging[slot][:count],
+                non_blocking=True,
+            )
+            self.events[slot].record()
+            if pending is not None:
+                pending_slot, pending_destination, pending_count = pending
+                self.events[pending_slot].synchronize()
+                self.extension.cpu_scatter(
+                    self.cpu_staging[pending_slot],
+                    self.cpu_pool,
+                    pending_destination,
+                    pending_count,
+                )
+            pending = (slot, destination, count)
+
+        if pending is not None:
+            pending_slot, pending_destination, pending_count = pending
+            self.events[pending_slot].synchronize()
+            self.extension.cpu_scatter(
+                self.cpu_staging[pending_slot],
+                self.cpu_pool,
+                pending_destination,
+                pending_count,
+            )
+        self._finish("d2h", block_count, started)
+
+    def swap_in(self, mapping: torch.Tensor) -> None:
+        started = time.perf_counter() if self.trace_enabled else None
+        source_cpu, destination_gpu = validate_block_mapping(
+            mapping,
+            source_limit=self.num_cpu_blocks,
+            destination_limit=self.num_gpu_blocks,
+        )
+        block_count = source_cpu.numel()
+        if block_count == 0:
+            return
+        destination_gpu_ids = self._to_gpu_ids(destination_gpu)
+
+        self._begin()
+        for index, (source, _, gpu_ids, count) in enumerate(
+                self._chunks(
+                    source_cpu, destination_gpu, destination_gpu_ids)):
+            slot = index % STAGING_BUFFER_COUNT
+            if index >= STAGING_BUFFER_COUNT:
+                self.events[slot].synchronize()
+            self.extension.cpu_gather(
+                self.cpu_pool,
+                source,
+                self.cpu_staging[slot],
+                count,
+            )
+            self.gpu_staging[slot][:count].copy_(
+                self.cpu_staging[slot][:count],
+                non_blocking=True,
+            )
+            self.extension.scatter(
+                self.gpu_staging[slot],
+                gpu_ids,
+                self.gpu_cache,
+                self.error_flag,
+                count,
+            )
+            self.events[slot].record()
+        self._finish("h2d", block_count, started)

qwen3_6_scripts/build_corex_attn_head_rms_norm.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_attn_head_rms_norm.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_attn_head_rms_norm.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" \
+    --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check \
+    -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_attn_head_rms_norm \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" \
+    -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_attn_head_rms_norm.cu" \
+    -L"${TORCH_ROOT}/lib" \
+    -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" \
+    -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart \
+    -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX attention head RMSNorm extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_fused_paged_prefill_split4.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_fused_paged_prefill_split4.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_fused_paged_prefill_split4.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_fused_paged_prefill \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_fused_paged_prefill_split4.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcublas -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX split4 fused paged-prefill extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_gdn_beta_decay.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_beta_decay.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_beta_decay.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" \
+    --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check \
+    -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_beta_decay \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" \
+    -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_gdn_beta_decay.cu" \
+    -L"${TORCH_ROOT}/lib" \
+    -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" \
+    -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart \
+    -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN beta/decay extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_gdn_causal_conv.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_causal_conv.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_causal_conv.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" \
+    --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check \
+    -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_causal_conv \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" \
+    -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_gdn_causal_conv.cu" \
+    -L"${TORCH_ROOT}/lib" \
+    -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" \
+    -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart \
+    -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN causal conv extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_gdn_gated_norm.sh

@@ -0,0 +1,33 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_gated_norm.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_gated_norm.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" \
+    --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check \
+    -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_gated_norm \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" \
+    -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_gdn_gated_norm.cu" \
+    -L"${TORCH_ROOT}/lib" \
+    -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" \
+    -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart \
+    -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN gated norm extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_gdn_packed_decode.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_packed_decode.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_packed_decode.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_packed_decode \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_gdn_packed_decode.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN packed decode extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_gdn_qk_map.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_gdn_qk_map.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_gdn_qk_map.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_gdn_qk_map \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_gdn_qk_map.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX GDN q/k map extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_moe_direct_routed.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_moe_direct_routed.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_moe_direct_routed.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_moe_direct_routed \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_moe_direct_routed.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX direct routed-expert extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_moe_exact_reduce.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_moe_exact_reduce.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_moe_exact_reduce.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_moe_exact_reduce \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_moe_exact_reduce.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX MoE exact reduce extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_moe_weight_gather.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_moe_weight_gather.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_moe_weight_gather.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_moe_weight_gather \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_moe_weight_gather.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX MoE selected-weight gather extension %s\n' "${OUTPUT}"

qwen3_6_scripts/build_corex_paged_kv_gather.sh

@@ -0,0 +1,27 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: build_corex_paged_kv_gather.sh VLLM_ROOT}
+COREX_ROOT=${COREX_ROOT:-/usr/local/corex-3.2.3}
+TORCH_ROOT=${TORCH_ROOT:-${COREX_ROOT}/lib64/python3/dist-packages/torch}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+OUTPUT=${VLLM_ROOT}/corex_paged_kv_gather.so
+
+"${COREX_ROOT}/bin/clang++" \
+    -std=c++17 -O3 -shared -fPIC \
+    --cuda-path="${COREX_ROOT}" --cuda-gpu-arch=ivcore10 \
+    --no-cuda-version-check -D_GLIBCXX_USE_CXX11_ABI=0 \
+    -DTORCH_EXTENSION_NAME=corex_paged_kv_gather \
+    -DTORCH_API_INCLUDE_EXTENSION_H \
+    -I"${TORCH_ROOT}/include" \
+    -I"${TORCH_ROOT}/include/torch/csrc/api/include" \
+    -I"${TORCH_ROOT}/include/TH" -I"${TORCH_ROOT}/include/THC" \
+    -I/usr/local/include/python3.10 \
+    "${SCRIPT_DIR}/corex_paged_kv_gather.cu" \
+    -L"${TORCH_ROOT}/lib" -L"${COREX_ROOT}/lib64" \
+    -Wl,-rpath,"${TORCH_ROOT}/lib" -Wl,-rpath,"${COREX_ROOT}/lib64" \
+    -ltorch_python -ltorch_cuda -ltorch_cpu -ltorch \
+    -lc10_cuda -lc10 -lcudart -o "${OUTPUT}"
+
+test -s "${OUTPUT}"
+printf '[ok] CoreX paged K/V gather extension %s\n' "${OUTPUT}"

qwen3_6_scripts/chat_utils.py

@@ -172,8 +172,8 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
                 return "<image>"
             if model_type == "mllama":
                 return "<|image|>"
-            if model_type in ("qwen2_vl", "qwen2_5_vl",
-                              "qwen3_5", "qwen3_5_moe"):
+            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 ""
@@ -184,8 +184,7 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):
                 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"):
+            if model_type in ("qwen2_vl","qwen2_5_vl"):
                 return "<|vision_start|><|video_pad|><|vision_end|>"
             raise TypeError(f"Unknown model type: {model_type}")
         else:
@@ -514,11 +513,26 @@ def _postprocess_messages(messages: List[ConversationMessage]) -> None:
     # from openAI format) to dict
     for message in messages:
         if (message["role"] == "assistant" and "tool_calls" in message
-                and isinstance(message["tool_calls"], list)):
+                and message["tool_calls"] is not None):
+            if not isinstance(message["tool_calls"], list):
+                message["tool_calls"] = list(message["tool_calls"])
 
             for item in message["tool_calls"]:
-                item["function"]["arguments"] = json.loads(
-                    item["function"]["arguments"])
+                arguments = item["function"]["arguments"]
+                if isinstance(arguments, str):
+                    try:
+                        arguments = json.loads(arguments)
+                    except json.JSONDecodeError as exc:
+                        raise ValueError(
+                            "Tool call arguments are not valid JSON.") from exc
+                elif not isinstance(arguments, dict):
+                    raise TypeError(
+                        "Tool call arguments must be a JSON object or a "
+                        "JSON-encoded object string.")
+                if not isinstance(arguments, dict):
+                    raise TypeError(
+                        "Tool call arguments must decode to a JSON object.")
+                item["function"]["arguments"] = arguments
 
 
 def parse_chat_messages(

qwen3_6_scripts/corex_attn_head_rms_norm.cu

@@ -0,0 +1,102 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+#include <vector>
+
+namespace {
+
+constexpr int kHeadDim = 256;
+constexpr int kThreads = 256;
+
+void check_half_matrix(const torch::Tensor& input, const char* name) {
+  TORCH_CHECK(input.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(input.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(input.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(input.dim() == 2 && input.size(1) == kHeadDim,
+              name, " must have shape (rows, 256)");
+}
+
+__global__ void prepare_kernel(const __half* input, float* converted,
+                               float* squares, int rows) {
+  const int row = blockIdx.x;
+  const int column = threadIdx.x;
+  if (row >= rows || column >= kHeadDim) {
+    return;
+  }
+  const int offset = row * kHeadDim + column;
+  const float value = __half2float(input[offset]);
+  converted[offset] = value;
+  squares[offset] = __fmul_rn(value, value);
+}
+
+__global__ void apply_inverse_kernel(
+    const float* input, const __half* weight, const float* inverse,
+    __half* output, int rows) {
+  const int row = blockIdx.x;
+  const int column = threadIdx.x;
+  if (row >= rows || column >= kHeadDim) {
+    return;
+  }
+  const int offset = row * kHeadDim + column;
+  const float scaled = __fmul_rn(input[offset], inverse[row]);
+  const float factor = __fadd_rn(1.0f, __half2float(weight[column]));
+  output[offset] = __float2half_rn(__fmul_rn(scaled, factor));
+}
+
+}  // namespace
+
+std::vector<torch::Tensor> prepare(const torch::Tensor& input) {
+  check_half_matrix(input, "input");
+  auto float_options = input.options().dtype(torch::kFloat32);
+  auto converted = torch::empty(input.sizes(), float_options);
+  auto squares = torch::empty(input.sizes(), float_options);
+  const int rows = static_cast<int>(input.size(0));
+  prepare_kernel<<<rows, kThreads, 0, at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
+      converted.data_ptr<float>(), squares.data_ptr<float>(), rows);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {converted, squares};
+}
+
+torch::Tensor apply_inverse(const torch::Tensor& input,
+                            const torch::Tensor& weight,
+                            const torch::Tensor& inverse) {
+  TORCH_CHECK(input.is_cuda() && weight.is_cuda() && inverse.is_cuda(),
+              "all tensors must be CUDA tensors");
+  TORCH_CHECK(input.scalar_type() == torch::kFloat32,
+              "input must have dtype float32");
+  TORCH_CHECK(weight.scalar_type() == torch::kFloat16,
+              "weight must have dtype float16");
+  TORCH_CHECK(inverse.scalar_type() == torch::kFloat32,
+              "inverse must have dtype float32");
+  TORCH_CHECK(input.is_contiguous() && weight.is_contiguous()
+                  && inverse.is_contiguous(),
+              "all tensors must be contiguous");
+  TORCH_CHECK(input.dim() == 2 && input.size(1) == kHeadDim,
+              "input must have shape (rows, 256)");
+  TORCH_CHECK(weight.dim() == 1 && weight.size(0) == kHeadDim,
+              "weight must have shape (256,)");
+  TORCH_CHECK(inverse.numel() == input.size(0),
+              "inverse must contain one value per row");
+  auto output = torch::empty(
+      input.sizes(), input.options().dtype(torch::kFloat16));
+  const int rows = static_cast<int>(input.size(0));
+  apply_inverse_kernel<<<rows, kThreads, 0,
+                         at::cuda::getCurrentCUDAStream()>>>(
+      input.data_ptr<float>(),
+      reinterpret_cast<const __half*>(weight.data_ptr<at::Half>()),
+      inverse.data_ptr<float>(),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()), rows);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("prepare", &prepare,
+             "Convert FP16 attention heads and compute exact squares");
+  module.def("apply_inverse", &apply_inverse,
+             "Apply PyTorch-computed attention head RMSNorm inverse");
+}

qwen3_6_scripts/corex_block_major_kv_transfer.cu

@@ -0,0 +1,402 @@
+#include <ATen/ATen.h>
+#include <ATen/Parallel.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstring>
+#include <vector>
+
+namespace {
+
+constexpr int kAttentionLayers = 10;
+constexpr int kKvPlanes = 2;
+constexpr int kElementsPerPlaneBlock = 4096;
+constexpr int kElementsPerVector = 8;
+constexpr int kVectorsPerPlaneBlock =
+    kElementsPerPlaneBlock / kElementsPerVector;
+constexpr int kVectorsPerBlockMajorRow =
+    kAttentionLayers * kKvPlanes * kVectorsPerPlaneBlock;
+constexpr int kThreads = 256;
+constexpr int kMaxGridBlocks = 65535;
+
+using PackedVector = uint4;
+
+__device__ __forceinline__ const PackedVector* select_const_layer(
+    int layer, const PackedVector* layer0, const PackedVector* layer1,
+    const PackedVector* layer2, const PackedVector* layer3,
+    const PackedVector* layer4, const PackedVector* layer5,
+    const PackedVector* layer6, const PackedVector* layer7,
+    const PackedVector* layer8, const PackedVector* layer9) {
+  switch (layer) {
+    case 0:
+      return layer0;
+    case 1:
+      return layer1;
+    case 2:
+      return layer2;
+    case 3:
+      return layer3;
+    case 4:
+      return layer4;
+    case 5:
+      return layer5;
+    case 6:
+      return layer6;
+    case 7:
+      return layer7;
+    case 8:
+      return layer8;
+    default:
+      return layer9;
+  }
+}
+
+__device__ __forceinline__ PackedVector* select_mutable_layer(
+    int layer, PackedVector* layer0, PackedVector* layer1,
+    PackedVector* layer2, PackedVector* layer3, PackedVector* layer4,
+    PackedVector* layer5, PackedVector* layer6, PackedVector* layer7,
+    PackedVector* layer8, PackedVector* layer9) {
+  switch (layer) {
+    case 0:
+      return layer0;
+    case 1:
+      return layer1;
+    case 2:
+      return layer2;
+    case 3:
+      return layer3;
+    case 4:
+      return layer4;
+    case 5:
+      return layer5;
+    case 6:
+      return layer6;
+    case 7:
+      return layer7;
+    case 8:
+      return layer8;
+    default:
+      return layer9;
+  }
+}
+
+__global__ void pack_block_major_kernel(
+    const PackedVector* layer0, const PackedVector* layer1,
+    const PackedVector* layer2, const PackedVector* layer3,
+    const PackedVector* layer4, const PackedVector* layer5,
+    const PackedVector* layer6, const PackedVector* layer7,
+    const PackedVector* layer8, const PackedVector* layer9,
+    const int* source_blocks, PackedVector* staging, int* error_flag,
+    int count, int gpu_blocks) {
+  const int64_t total =
+      static_cast<int64_t>(count) * kVectorsPerBlockMajorRow;
+  for (int64_t linear =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       linear < total;
+       linear += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    int64_t cursor = linear;
+    const int feature_vector = cursor % kVectorsPerPlaneBlock;
+    cursor /= kVectorsPerPlaneBlock;
+    const int kv_plane = cursor % kKvPlanes;
+    cursor /= kKvPlanes;
+    const int layer = cursor % kAttentionLayers;
+    const int row = cursor / kAttentionLayers;
+    const int source_block = source_blocks[row];
+    if (static_cast<unsigned int>(source_block) >=
+        static_cast<unsigned int>(gpu_blocks)) {
+      atomicExch(error_flag, 1);
+      continue;
+    }
+    const PackedVector* source = select_const_layer(
+        layer, layer0, layer1, layer2, layer3, layer4, layer5, layer6,
+        layer7, layer8, layer9);
+    const int64_t source_index =
+        ((static_cast<int64_t>(kv_plane) * gpu_blocks + source_block)
+         * kVectorsPerPlaneBlock) +
+        feature_vector;
+    staging[linear] = source[source_index];
+  }
+}
+
+__global__ void scatter_block_major_kernel(
+    const PackedVector* staging, const int* destination_blocks,
+    PackedVector* layer0, PackedVector* layer1, PackedVector* layer2,
+    PackedVector* layer3, PackedVector* layer4, PackedVector* layer5,
+    PackedVector* layer6, PackedVector* layer7, PackedVector* layer8,
+    PackedVector* layer9, int* error_flag, int count, int gpu_blocks) {
+  const int64_t total =
+      static_cast<int64_t>(count) * kVectorsPerBlockMajorRow;
+  for (int64_t linear =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       linear < total;
+       linear += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    int64_t cursor = linear;
+    const int feature_vector = cursor % kVectorsPerPlaneBlock;
+    cursor /= kVectorsPerPlaneBlock;
+    const int kv_plane = cursor % kKvPlanes;
+    cursor /= kKvPlanes;
+    const int layer = cursor % kAttentionLayers;
+    const int row = cursor / kAttentionLayers;
+    const int destination_block = destination_blocks[row];
+    if (static_cast<unsigned int>(destination_block) >=
+        static_cast<unsigned int>(gpu_blocks)) {
+      atomicExch(error_flag, 1);
+      continue;
+    }
+    PackedVector* destination = select_mutable_layer(
+        layer, layer0, layer1, layer2, layer3, layer4, layer5, layer6,
+        layer7, layer8, layer9);
+    const int64_t destination_index =
+        ((static_cast<int64_t>(kv_plane) * gpu_blocks + destination_block)
+         * kVectorsPerPlaneBlock) +
+        feature_vector;
+    destination[destination_index] = staging[linear];
+  }
+}
+
+void check_gpu_layers(const std::vector<torch::Tensor>& layers) {
+  TORCH_CHECK(layers.size() == kAttentionLayers, "expected exactly ",
+              kAttentionLayers, " GPU attention-layer tensors");
+  const auto device = layers.front().device();
+  const int64_t blocks = layers.front().size(1);
+  for (int layer = 0; layer < kAttentionLayers; ++layer) {
+    const auto& tensor = layers[layer];
+    TORCH_CHECK(tensor.is_cuda(), "GPU layer ", layer,
+                " must be a CUDA tensor");
+    TORCH_CHECK(tensor.device() == device, "GPU layer ", layer,
+                " is on a different device");
+    TORCH_CHECK(tensor.scalar_type() == torch::kFloat16, "GPU layer ",
+                layer, " must use float16");
+    TORCH_CHECK(tensor.is_contiguous(), "GPU layer ", layer,
+                " must be contiguous");
+    TORCH_CHECK(tensor.dim() == 3 && tensor.size(0) == kKvPlanes &&
+                    tensor.size(1) == blocks &&
+                    tensor.size(2) == kElementsPerPlaneBlock,
+                "GPU layer ", layer, " must have shape [2, blocks, 4096]");
+    TORCH_CHECK(
+        reinterpret_cast<uintptr_t>(tensor.data_ptr<at::Half>()) %
+                alignof(PackedVector) ==
+            0,
+        "GPU layer ", layer, " is not 16-byte aligned");
+  }
+}
+
+void check_gpu_transfer_args(const std::vector<torch::Tensor>& layers,
+                             const torch::Tensor& block_ids,
+                             const torch::Tensor& staging,
+                             const torch::Tensor& error_flag,
+                             int64_t count) {
+  check_gpu_layers(layers);
+  TORCH_CHECK(block_ids.is_cuda(), "block_ids must be a CUDA tensor");
+  TORCH_CHECK(block_ids.device() == layers.front().device(),
+              "block_ids must be on the cache device");
+  TORCH_CHECK(block_ids.scalar_type() == torch::kInt32,
+              "block_ids must use int32");
+  TORCH_CHECK(block_ids.dim() == 1 && block_ids.is_contiguous(),
+              "block_ids must be a contiguous one-dimensional tensor");
+  TORCH_CHECK(count > 0 && count <= block_ids.numel(),
+              "count must be in [1, block_ids.numel()]");
+  TORCH_CHECK(staging.is_cuda(), "staging must be a CUDA tensor");
+  TORCH_CHECK(staging.device() == layers.front().device(),
+              "staging must be on the cache device");
+  TORCH_CHECK(staging.scalar_type() == torch::kFloat16,
+              "staging must use float16");
+  TORCH_CHECK(staging.is_contiguous(), "staging must be contiguous");
+  TORCH_CHECK(
+      staging.dim() == 4 && staging.size(0) >= count &&
+          staging.size(1) == kAttentionLayers &&
+          staging.size(2) == kKvPlanes &&
+          staging.size(3) == kElementsPerPlaneBlock,
+      "staging must have shape [capacity>=count, 10, 2, 4096]");
+  TORCH_CHECK(
+      reinterpret_cast<uintptr_t>(staging.data_ptr<at::Half>()) %
+              alignof(PackedVector) ==
+          0,
+      "staging is not 16-byte aligned");
+  TORCH_CHECK(error_flag.is_cuda(),
+              "error_flag must be a CUDA tensor");
+  TORCH_CHECK(error_flag.device() == layers.front().device(),
+              "error_flag must be on the cache device");
+  TORCH_CHECK(error_flag.scalar_type() == torch::kInt32,
+              "error_flag must use int32");
+  TORCH_CHECK(error_flag.is_contiguous() && error_flag.numel() == 1,
+              "error_flag must be one contiguous int32 value");
+}
+
+int launch_blocks(int64_t count) {
+  const int64_t total = count * kVectorsPerBlockMajorRow;
+  return static_cast<int>(std::min<int64_t>(
+      (total + kThreads - 1) / kThreads, kMaxGridBlocks));
+}
+
+void pack_block_major(const std::vector<torch::Tensor>& layers,
+                      const torch::Tensor& source_blocks,
+                      torch::Tensor staging, torch::Tensor error_flag,
+                      int64_t count) {
+  check_gpu_transfer_args(
+      layers, source_blocks, staging, error_flag, count);
+  const int blocks = static_cast<int>(layers.front().size(1));
+  pack_block_major_kernel<<<launch_blocks(count), kThreads, 0,
+                            at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const PackedVector*>(
+          layers[0].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[1].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[2].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[3].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[4].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[5].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[6].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[7].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[8].data_ptr<at::Half>()),
+      reinterpret_cast<const PackedVector*>(
+          layers[9].data_ptr<at::Half>()),
+      source_blocks.data_ptr<int>(),
+      reinterpret_cast<PackedVector*>(staging.data_ptr<at::Half>()),
+      error_flag.data_ptr<int>(), static_cast<int>(count), blocks);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void scatter_block_major(const torch::Tensor& staging,
+                         const torch::Tensor& destination_blocks,
+                         const std::vector<torch::Tensor>& layers,
+                         torch::Tensor error_flag,
+                         int64_t count) {
+  check_gpu_transfer_args(
+      layers, destination_blocks, staging, error_flag, count);
+  const int blocks = static_cast<int>(layers.front().size(1));
+  scatter_block_major_kernel<<<launch_blocks(count), kThreads, 0,
+                               at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const PackedVector*>(
+          staging.data_ptr<at::Half>()),
+      destination_blocks.data_ptr<int>(),
+      reinterpret_cast<PackedVector*>(layers[0].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[1].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[2].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[3].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[4].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[5].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[6].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[7].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[8].data_ptr<at::Half>()),
+      reinterpret_cast<PackedVector*>(layers[9].data_ptr<at::Half>()),
+      error_flag.data_ptr<int>(), static_cast<int>(count), blocks);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+}
+
+void check_transfer_error(const torch::Tensor& error_flag) {
+  TORCH_CHECK(error_flag.is_cuda(),
+              "error_flag must be a CUDA tensor");
+  TORCH_CHECK(error_flag.scalar_type() == torch::kInt32,
+              "error_flag must use int32");
+  TORCH_CHECK(error_flag.is_contiguous() && error_flag.numel() == 1,
+              "error_flag must be one contiguous int32 value");
+  TORCH_CHECK(error_flag.item<int>() == 0,
+              "GPU block mapping contains an out-of-range id");
+}
+
+void check_cpu_transfer_args(const torch::Tensor& pool,
+                             const torch::Tensor& block_ids,
+                             const torch::Tensor& staging, int64_t count) {
+  TORCH_CHECK(!pool.is_cuda() && !staging.is_cuda() &&
+                  !block_ids.is_cuda(),
+              "CPU gather/scatter tensors must be on CPU");
+  TORCH_CHECK(pool.scalar_type() == torch::kFloat16 &&
+                  staging.scalar_type() == torch::kFloat16,
+              "CPU pool and staging must use float16");
+  TORCH_CHECK(pool.is_contiguous() && staging.is_contiguous(),
+              "CPU pool and staging must be contiguous");
+  TORCH_CHECK(
+      pool.dim() == 4 && pool.size(1) == kAttentionLayers &&
+          pool.size(2) == kKvPlanes &&
+          pool.size(3) == kElementsPerPlaneBlock,
+      "CPU pool must have shape [slots, 10, 2, 4096]");
+  TORCH_CHECK(
+      staging.dim() == 4 && staging.size(0) >= count &&
+          staging.size(1) == kAttentionLayers &&
+          staging.size(2) == kKvPlanes &&
+          staging.size(3) == kElementsPerPlaneBlock,
+      "CPU staging must have shape [capacity>=count, 10, 2, 4096]");
+  TORCH_CHECK(block_ids.scalar_type() == torch::kInt64,
+              "CPU block_ids must use int64");
+  TORCH_CHECK(block_ids.dim() == 1 && block_ids.is_contiguous(),
+              "CPU block_ids must be contiguous and one-dimensional");
+  TORCH_CHECK(count > 0 && count <= block_ids.numel(),
+              "count must be in [1, block_ids.numel()]");
+
+  const int64_t* ids = block_ids.data_ptr<int64_t>();
+  for (int64_t row = 0; row < count; ++row) {
+    TORCH_CHECK(ids[row] >= 0 && ids[row] < pool.size(0),
+                "CPU block id out of range at row ", row, ": ", ids[row]);
+  }
+}
+
+void cpu_gather_rows(const torch::Tensor& pool,
+                     const torch::Tensor& source_blocks,
+                     torch::Tensor staging, int64_t count) {
+  check_cpu_transfer_args(pool, source_blocks, staging, count);
+  const int64_t row_elements =
+      kAttentionLayers * kKvPlanes * kElementsPerPlaneBlock;
+  const size_t row_bytes =
+      static_cast<size_t>(row_elements) * sizeof(at::Half);
+  const char* source = reinterpret_cast<const char*>(
+      pool.data_ptr<at::Half>());
+  char* destination =
+      reinterpret_cast<char*>(staging.data_ptr<at::Half>());
+  const int64_t* ids = source_blocks.data_ptr<int64_t>();
+  at::parallel_for(0, count, 8, [&](int64_t begin, int64_t end) {
+    for (int64_t row = begin; row < end; ++row) {
+      std::memcpy(destination + row * row_bytes,
+                  source + ids[row] * row_bytes, row_bytes);
+    }
+  });
+}
+
+void cpu_scatter_rows(const torch::Tensor& staging,
+                      torch::Tensor pool,
+                      const torch::Tensor& destination_blocks,
+                      int64_t count) {
+  check_cpu_transfer_args(pool, destination_blocks, staging, count);
+  const int64_t row_elements =
+      kAttentionLayers * kKvPlanes * kElementsPerPlaneBlock;
+  const size_t row_bytes =
+      static_cast<size_t>(row_elements) * sizeof(at::Half);
+  const char* source = reinterpret_cast<const char*>(
+      staging.data_ptr<at::Half>());
+  char* destination =
+      reinterpret_cast<char*>(pool.data_ptr<at::Half>());
+  const int64_t* ids = destination_blocks.data_ptr<int64_t>();
+  at::parallel_for(0, count, 8, [&](int64_t begin, int64_t end) {
+    for (int64_t row = begin; row < end; ++row) {
+      std::memcpy(destination + ids[row] * row_bytes,
+                  source + row * row_bytes, row_bytes);
+    }
+  });
+}
+
+}  // namespace
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("pack", &pack_block_major,
+             "Pack ten layer-major FP16 KV caches into block-major staging");
+  module.def("scatter", &scatter_block_major,
+             "Scatter block-major FP16 staging into ten layer-major caches");
+  module.def("check_error", &check_transfer_error,
+             "Fail fast after a bounds-safe asynchronous transfer");
+  module.def("cpu_gather", &cpu_gather_rows,
+             "Gather block-major CPU pool rows into bounded staging");
+  module.def("cpu_scatter", &cpu_scatter_rows,
+             "Scatter bounded staging rows into the block-major CPU pool");
+}

qwen3_6_scripts/corex_fused_paged_prefill_split4.cu

@@ -0,0 +1,494 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cublas_v2.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+namespace {
+
+constexpr int kBlockSize = 16;
+constexpr int kHeadDim = 256;
+constexpr int kKeyPack = 8;
+constexpr int kNumQueryHeads = 4;
+constexpr int kNumKvHeads = 1;
+constexpr int kTileTokens = 512;
+constexpr int kSplitCount = 4;
+constexpr int kGroupTokens = kSplitCount * kTileTokens;
+constexpr int kThreads = 256;
+constexpr int kMaxQueryTokens = 8192;
+constexpr int kMaxSequenceTokens = 262144;
+
+void check_half_cuda_contiguous(const torch::Tensor& tensor,
+                                const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+__global__ void convert_query_kernel(const __half* query, float* converted,
+                                     int query_len, float scale) {
+  const int64_t total = static_cast<int64_t>(query_len)
+      * kNumQueryHeads * kHeadDim;
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < total;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int dim = index % kHeadDim;
+    const int query_index =
+        (index / kHeadDim) % query_len;
+    const int head =
+        index / (static_cast<int64_t>(kHeadDim) * query_len);
+    const int64_t source =
+        (static_cast<int64_t>(query_index) * kNumQueryHeads + head)
+        * kHeadDim + dim;
+    converted[index] = __half2float(query[source]) * scale;
+  }
+}
+
+__global__ void gather_kv_group_kernel(
+    const __half* key_new, const __half* value_new,
+    const __half* key_cache, const __half* value_cache,
+    const int* block_table, float* key_tiles, float* value_tiles,
+    int context_len, int query_len, int group_start, int group_tokens,
+    int active_splits) {
+  constexpr int kElements = kTileTokens * kHeadDim;
+  const int64_t total = static_cast<int64_t>(active_splits) * kElements;
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < total;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int split = index / kElements;
+    const int element = index - static_cast<int64_t>(split) * kElements;
+    const int token_offset = element / kHeadDim;
+    const int dim = element - token_offset * kHeadDim;
+    const int remaining_tokens = group_tokens - split * kTileTokens;
+    const int split_tokens =
+        remaining_tokens < kTileTokens ? remaining_tokens : kTileTokens;
+    const int logical_token =
+        group_start + split * kTileTokens + token_offset;
+    float key_value = 0.0f;
+    float value_value = 0.0f;
+    if (token_offset >= split_tokens) {
+      // The fixed 512-column GEMMs require zero-filled tail columns.
+    } else if (logical_token < context_len) {
+      const int logical_block = logical_token / kBlockSize;
+      const int block_offset = logical_token % kBlockSize;
+      const int physical_block = block_table[logical_block];
+      const int64_t key_index =
+          (((static_cast<int64_t>(physical_block) * kNumKvHeads)
+             * (kHeadDim / kKeyPack) + dim / kKeyPack)
+            * kBlockSize + block_offset) * kKeyPack + dim % kKeyPack;
+      const int64_t value_index =
+          ((static_cast<int64_t>(physical_block) * kNumKvHeads)
+             * kHeadDim + dim) * kBlockSize + block_offset;
+      key_value = __half2float(key_cache[key_index]);
+      value_value = __half2float(value_cache[value_index]);
+    } else if (logical_token < context_len + query_len) {
+      const int query_index = logical_token - context_len;
+      const int64_t source =
+          static_cast<int64_t>(query_index) * kHeadDim + dim;
+      key_value = __half2float(key_new[source]);
+      value_value = __half2float(value_new[source]);
+    }
+    key_tiles[index] = key_value;
+    value_tiles[index] = value_value;
+  }
+}
+
+__global__ void mask_group_scores_kernel(
+    float* scores, int query_len, int context_len,
+    int group_start, int group_tokens, int active_splits,
+    int rows, bool causal) {
+  const int64_t split_elements =
+      static_cast<int64_t>(rows) * kTileTokens;
+  const int64_t elements = active_splits * split_elements;
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < elements;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int split = index / split_elements;
+    const int split_index = index - split * split_elements;
+    const int column = split_index % kTileTokens;
+    const int row = split_index / kTileTokens;
+    const int query_index = row % query_len;
+    const int remaining_tokens = group_tokens - split * kTileTokens;
+    const int split_tokens =
+        remaining_tokens < kTileTokens ? remaining_tokens : kTileTokens;
+    const int logical_token =
+        group_start + split * kTileTokens + column;
+    if (column >= split_tokens
+        || (causal && logical_token > context_len + query_index)) {
+      scores[index] = -std::numeric_limits<float>::infinity();
+    }
+  }
+}
+
+__global__ void normalize_split_scores_kernel(
+    float* scores, float* corrections, float* running_max,
+    float* running_sum, int active_splits, int rows) {
+  const int row = blockIdx.x;
+  if (row >= rows) {
+    return;
+  }
+
+  __shared__ float reduction[kThreads];
+  __shared__ float state_max;
+  __shared__ float state_sum;
+  __shared__ float next_max;
+  __shared__ float correction;
+
+  if (threadIdx.x == 0) {
+    state_max = running_max[row];
+    state_sum = running_sum[row];
+  }
+  __syncthreads();
+
+  for (int split = 0; split < active_splits; ++split) {
+    float* row_scores =
+        scores + (static_cast<int64_t>(split) * rows + row) * kTileTokens;
+    float local_max = -std::numeric_limits<float>::infinity();
+    for (int column = threadIdx.x; column < kTileTokens;
+         column += blockDim.x) {
+      local_max = fmaxf(local_max, row_scores[column]);
+    }
+    reduction[threadIdx.x] = local_max;
+    __syncthreads();
+    for (int stride = kThreads / 2; stride > 0; stride /= 2) {
+      if (threadIdx.x < stride) {
+        reduction[threadIdx.x] = fmaxf(
+            reduction[threadIdx.x], reduction[threadIdx.x + stride]);
+      }
+      __syncthreads();
+    }
+    if (threadIdx.x == 0) {
+      next_max = fmaxf(state_max, reduction[0]);
+      correction =
+          (state_max == -std::numeric_limits<float>::infinity()
+           && next_max == -std::numeric_limits<float>::infinity())
+              ? 1.0f
+              : expf(state_max - next_max);
+      corrections[static_cast<int64_t>(split) * rows + row] = correction;
+    }
+    __syncthreads();
+
+    float local_sum = 0.0f;
+    for (int column = threadIdx.x; column < kTileTokens;
+         column += blockDim.x) {
+      const float score = row_scores[column];
+      const float probability =
+          (score == -std::numeric_limits<float>::infinity()
+           && next_max == -std::numeric_limits<float>::infinity())
+              ? 0.0f
+              : expf(score - next_max);
+      row_scores[column] = probability;
+      local_sum = __fadd_rn(local_sum, probability);
+    }
+    reduction[threadIdx.x] = local_sum;
+    __syncthreads();
+    for (int stride = kThreads / 2; stride > 0; stride /= 2) {
+      if (threadIdx.x < stride) {
+        reduction[threadIdx.x] = __fadd_rn(
+            reduction[threadIdx.x], reduction[threadIdx.x + stride]);
+      }
+      __syncthreads();
+    }
+    if (threadIdx.x == 0) {
+      state_sum = __fadd_rn(
+          __fmul_rn(state_sum, correction), reduction[0]);
+      state_max = next_max;
+    }
+    __syncthreads();
+  }
+
+  if (threadIdx.x == 0) {
+    running_max[row] = state_max;
+    running_sum[row] = state_sum;
+  }
+}
+
+__global__ void merge_split_output_kernel(
+    float* running_output, const float* split_output,
+    const float* corrections, int active_splits,
+    int rows, int64_t output_elements) {
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < output_elements;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int row = index / kHeadDim;
+    float value = running_output[index];
+    for (int split = 0; split < active_splits; ++split) {
+      const int64_t row_index =
+          static_cast<int64_t>(split) * rows + row;
+      const int64_t output_index =
+          static_cast<int64_t>(split) * output_elements + index;
+      value = __fadd_rn(
+          __fmul_rn(value, corrections[row_index]),
+          split_output[output_index]);
+    }
+    running_output[index] = value;
+  }
+}
+
+__global__ void accumulate_output_kernel(
+    float* running_output, const float* tile_output,
+    const float* correction, int64_t elements) {
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < elements;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int row = index / kHeadDim;
+    const float scaled =
+        __fmul_rn(running_output[index], correction[row]);
+    running_output[index] = __fadd_rn(scaled, tile_output[index]);
+  }
+}
+
+int launch_blocks(int64_t elements) {
+  const int64_t needed = (elements + kThreads - 1) / kThreads;
+  return static_cast<int>(std::min<int64_t>(needed, 65535));
+}
+
+void check_cublas(cublasStatus_t status, const char* operation) {
+  TORCH_CHECK(status == CUBLAS_STATUS_SUCCESS, operation,
+              " failed with cuBLAS status ", static_cast<int>(status));
+}
+
+cublasStatus_t qk_batched(
+    cublasHandle_t handle, const float* key_tile, const float* query,
+    float* scores, int query_len) {
+  const float alpha = 1.0f;
+  const float beta = 0.0f;
+  return cublasSgemmStridedBatched(
+      handle, CUBLAS_OP_T, CUBLAS_OP_N,
+      kTileTokens, query_len, kHeadDim,
+      &alpha, key_tile, kHeadDim, 0,
+      query, kHeadDim, static_cast<long long>(query_len) * kHeadDim,
+      &beta, scores, kTileTokens,
+      static_cast<long long>(query_len) * kTileTokens,
+      kNumQueryHeads);
+}
+
+cublasStatus_t pv_batched(
+    cublasHandle_t handle, const float* value_tile, const float* scores,
+    float* output, int query_len) {
+  const float alpha = 1.0f;
+  const float beta = 0.0f;
+  return cublasSgemmStridedBatched(
+      handle, CUBLAS_OP_N, CUBLAS_OP_N,
+      kHeadDim, query_len, kTileTokens,
+      &alpha, value_tile, kHeadDim, 0,
+      scores, kTileTokens,
+      static_cast<long long>(query_len) * kTileTokens,
+      &beta, output, kHeadDim,
+      static_cast<long long>(query_len) * kHeadDim,
+      kNumQueryHeads);
+}
+
+}  // namespace
+
+std::vector<torch::Tensor> fused_paged_prefill_forward(
+    const torch::Tensor& query, const torch::Tensor& key_new,
+    const torch::Tensor& value_new, const torch::Tensor& key_cache,
+    const torch::Tensor& value_cache, const torch::Tensor& block_table,
+    int64_t context_len_arg, double scale_arg) {
+  check_half_cuda_contiguous(query, "query");
+  check_half_cuda_contiguous(key_new, "key_new");
+  check_half_cuda_contiguous(value_new, "value_new");
+  check_half_cuda_contiguous(key_cache, "key_cache");
+  check_half_cuda_contiguous(value_cache, "value_cache");
+  TORCH_CHECK(block_table.is_cuda(),
+              "block_table must be a CUDA tensor");
+  TORCH_CHECK(block_table.scalar_type() == torch::kInt32,
+              "block_table must have dtype int32");
+  TORCH_CHECK(block_table.is_contiguous(),
+              "block_table must be contiguous");
+  TORCH_CHECK(block_table.dim() == 1,
+              "block_table must be one-dimensional");
+  TORCH_CHECK(query.dim() == 3 && query.size(1) == kNumQueryHeads
+                  && query.size(2) == kHeadDim,
+              "query must have shape (Q, 4, 256)");
+  TORCH_CHECK(key_new.dim() == 3 && key_new.size(1) == kNumKvHeads
+                  && key_new.size(2) == kHeadDim,
+              "key_new must have shape (Q, 1, 256)");
+  TORCH_CHECK(value_new.sizes() == key_new.sizes(),
+              "value_new must match key_new");
+  TORCH_CHECK(key_new.size(0) == query.size(0),
+              "query, key_new, and value_new lengths must match");
+  TORCH_CHECK(key_cache.dim() == 5
+                  && key_cache.size(1) == kNumKvHeads
+                  && key_cache.size(2) == kHeadDim / kKeyPack
+                  && key_cache.size(3) == kBlockSize
+                  && key_cache.size(4) == kKeyPack,
+              "key_cache must have shape (N, 1, 32, 16, 8)");
+  TORCH_CHECK(value_cache.dim() == 4
+                  && value_cache.size(1) == kNumKvHeads
+                  && value_cache.size(2) == kHeadDim
+                  && value_cache.size(3) == kBlockSize,
+              "value_cache must have shape (N, 1, 256, 16)");
+  TORCH_CHECK(key_cache.size(0) == value_cache.size(0),
+              "key/value cache block counts must match");
+  TORCH_CHECK(query.device() == key_new.device()
+                  && query.device() == value_new.device()
+                  && query.device() == key_cache.device()
+                  && query.device() == value_cache.device()
+                  && query.device() == block_table.device(),
+              "all tensors must use the same device");
+  TORCH_CHECK(context_len_arg >= 0
+                  && context_len_arg <= kMaxSequenceTokens,
+              "context_len is out of range");
+  TORCH_CHECK(context_len_arg % kBlockSize == 0,
+              "context_len must be block aligned");
+  const int query_len = static_cast<int>(query.size(0));
+  const int context_len = static_cast<int>(context_len_arg);
+  TORCH_CHECK(query_len > 0 && query_len <= kMaxQueryTokens,
+              "query length must be in [1, 8192]");
+  TORCH_CHECK(context_len + query_len <= kMaxSequenceTokens,
+              "context_len + query_len exceeds 262144");
+  const int required_blocks =
+      (context_len + kBlockSize - 1) / kBlockSize;
+  TORCH_CHECK(block_table.numel() >= required_blocks,
+              "block_table is too short for context_len");
+  if (required_blocks > 0) {
+    auto active_blocks = block_table.narrow(0, 0, required_blocks);
+    const int minimum_block = active_blocks.min().item<int>();
+    const int maximum_block = active_blocks.max().item<int>();
+    TORCH_CHECK(minimum_block >= 0
+                    && maximum_block < key_cache.size(0),
+                "block_table contains an out-of-range physical block ID");
+  }
+  TORCH_CHECK(std::isfinite(scale_arg) && scale_arg > 0.0,
+              "scale must be finite and positive");
+  TORCH_CHECK(query_len <= std::numeric_limits<int>::max() / kNumQueryHeads,
+              "query length overflows row count");
+
+  const int rows = kNumQueryHeads * query_len;
+  const int64_t output_elements =
+      static_cast<int64_t>(rows) * kHeadDim;
+  auto float_options = query.options().dtype(torch::kFloat32);
+  auto converted_query = torch::empty(
+      {kNumQueryHeads, query_len, kHeadDim}, float_options);
+  auto key_tiles = torch::empty(
+      {kSplitCount, kTileTokens, kHeadDim}, float_options);
+  auto value_tiles = torch::empty(
+      {kSplitCount, kTileTokens, kHeadDim}, float_options);
+  auto scores = torch::empty(
+      {kSplitCount, kNumQueryHeads, query_len, kTileTokens},
+      float_options);
+  auto split_output = torch::empty(
+      {kSplitCount, kNumQueryHeads, query_len, kHeadDim},
+      float_options);
+  auto running_max = torch::full(
+      {kNumQueryHeads, query_len},
+      -std::numeric_limits<float>::infinity(), float_options);
+  auto running_sum = torch::zeros(
+      {kNumQueryHeads, query_len}, float_options);
+  auto running_output = torch::zeros(
+      {kNumQueryHeads, query_len, kHeadDim}, float_options);
+  auto corrections = torch::empty(
+      {kSplitCount, kNumQueryHeads, query_len}, float_options);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+  convert_query_kernel<<<launch_blocks(output_elements), kThreads, 0, stream>>>(
+      reinterpret_cast<const __half*>(query.data_ptr<at::Half>()),
+      converted_query.data_ptr<float>(), query_len,
+      static_cast<float>(scale_arg));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  check_cublas(cublasSetStream(handle, stream), "cublasSetStream");
+  const int64_t key_split_stride =
+      static_cast<int64_t>(kTileTokens) * kHeadDim;
+  const int64_t score_split_stride =
+      static_cast<int64_t>(rows) * kTileTokens;
+  const int64_t output_split_stride = output_elements;
+  const auto run_group = [&](int group_start, int group_tokens,
+                             bool causal) {
+    const int active_splits =
+        (group_tokens + kTileTokens - 1) / kTileTokens;
+    TORCH_CHECK(active_splits > 0 && active_splits <= kSplitCount,
+                "invalid split count for paged-prefill group");
+    constexpr int kGatherBlocks = 512;
+    gather_kv_group_kernel<<<kGatherBlocks, kThreads, 0, stream>>>(
+        reinterpret_cast<const __half*>(key_new.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(value_new.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(key_cache.data_ptr<at::Half>()),
+        reinterpret_cast<const __half*>(value_cache.data_ptr<at::Half>()),
+        block_table.data_ptr<int>(), key_tiles.data_ptr<float>(),
+        value_tiles.data_ptr<float>(), context_len, query_len, group_start,
+        group_tokens, active_splits);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    for (int split = 0; split < active_splits; ++split) {
+      check_cublas(qk_batched(
+          handle,
+          key_tiles.data_ptr<float>() + split * key_split_stride,
+          converted_query.data_ptr<float>(),
+          scores.data_ptr<float>() + split * score_split_stride,
+          query_len), "split4 paged prefill QK");
+    }
+
+    const bool needs_mask =
+        causal || group_tokens != active_splits * kTileTokens;
+    if (needs_mask) {
+      const int64_t score_elements =
+          static_cast<int64_t>(active_splits) * score_split_stride;
+      mask_group_scores_kernel<<<
+          launch_blocks(score_elements), kThreads, 0, stream>>>(
+          scores.data_ptr<float>(), query_len, context_len, group_start,
+          group_tokens, active_splits, rows, causal);
+      C10_CUDA_KERNEL_LAUNCH_CHECK();
+    }
+
+    normalize_split_scores_kernel<<<rows, kThreads, 0, stream>>>(
+        scores.data_ptr<float>(), corrections.data_ptr<float>(),
+        running_max.data_ptr<float>(), running_sum.data_ptr<float>(),
+        active_splits, rows);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    for (int split = 0; split < active_splits; ++split) {
+      check_cublas(pv_batched(
+          handle,
+          value_tiles.data_ptr<float>() + split * key_split_stride,
+          scores.data_ptr<float>() + split * score_split_stride,
+          split_output.data_ptr<float>() + split * output_split_stride,
+          query_len), "split4 paged prefill PV");
+    }
+
+    merge_split_output_kernel<<<
+        launch_blocks(output_elements), kThreads, 0, stream>>>(
+        running_output.data_ptr<float>(), split_output.data_ptr<float>(),
+        corrections.data_ptr<float>(), active_splits, rows,
+        output_elements);
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+  };
+  for (int group_start = 0; group_start < context_len;
+       group_start += kGroupTokens) {
+    run_group(group_start,
+              std::min(kGroupTokens, context_len - group_start), false);
+  }
+  for (int key_start = 0; key_start < query_len;
+       key_start += kGroupTokens) {
+    run_group(context_len + key_start,
+              std::min(kGroupTokens, query_len - key_start), true);
+  }
+
+  running_output.div_(running_sum.unsqueeze(-1));
+  auto output = running_output.permute({1, 0, 2})
+                    .to(query.scalar_type()).contiguous();
+  auto lse = (running_max + at::log(running_sum))
+                 .transpose(0, 1).contiguous();
+  return {output, lse};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("forward", &fused_paged_prefill_forward,
+             "Fixed-shape FP32 paged-prefill pipeline for cache-only context");
+}

qwen3_6_scripts/corex_gdn_beta_decay.cu

@@ -0,0 +1,84 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+__global__ void beta_decay_kernel(const half* beta_input,
+                                  const half* decay_input,
+                                  const half* a_log,
+                                  const half* dt_bias,
+                                  float* output, int elements,
+                                  int heads) {
+  const int index = blockIdx.x * blockDim.x + threadIdx.x;
+  if (index >= elements) {
+    return;
+  }
+  const int head = index % heads;
+
+  const float beta_value = __half2float(beta_input[index]);
+  const float beta_fp32 = 1.0f / (1.0f + expf(-beta_value));
+  output[index] = __half2float(__float2half(beta_fp32));
+
+  const float x = (__half2float(decay_input[index])
+                   + __half2float(dt_bias[head]));
+  const float softplus = x > 20.0f ? x : log1pf(expf(x));
+  output[elements + index] = expf(
+      -expf(__half2float(a_log[head])) * softplus);
+}
+
+void check_half(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 2, name, " must have shape (batch, heads)");
+}
+
+void check_half_vector(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 1, name, " must have shape (heads)");
+}
+
+}  // namespace
+
+torch::Tensor beta_decay(const torch::Tensor& beta_input,
+                         const torch::Tensor& decay_input,
+                         const torch::Tensor& a_log,
+                         const torch::Tensor& dt_bias) {
+  check_half(beta_input, "beta_input");
+  check_half(decay_input, "decay_input");
+  check_half_vector(a_log, "a_log");
+  check_half_vector(dt_bias, "dt_bias");
+  TORCH_CHECK(beta_input.sizes() == decay_input.sizes(),
+              "beta_input and decay_input shapes must match");
+  TORCH_CHECK(beta_input.size(1) == a_log.size(0) &&
+                  a_log.sizes() == dt_bias.sizes(),
+              "parameter heads must match input heads");
+
+  const int elements = static_cast<int>(beta_input.numel());
+  const int heads = static_cast<int>(beta_input.size(1));
+  torch::Tensor output = torch::empty(
+      {2, beta_input.size(0), beta_input.size(1)},
+      beta_input.options().dtype(torch::kFloat32));
+  constexpr int threads = 128;
+  const int blocks = (elements + threads - 1) / threads;
+  beta_decay_kernel<<<blocks, threads, 0,
+                      at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const half*>(beta_input.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(decay_input.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(a_log.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(dt_bias.data_ptr<at::Half>()),
+      output.data_ptr<float>(), elements, heads);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("beta_decay", &beta_decay,
+             "Fused GDN beta sigmoid and decay factor");
+}

qwen3_6_scripts/corex_gdn_causal_conv.cu

@@ -0,0 +1,89 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kStateLen = 3;
+constexpr int kKernelSize = kStateLen + 1;
+constexpr int kThreads = 256;
+
+__global__ void causal_conv_update_kernel(
+    float* state, const __half* hidden, const __half* weight,
+    __half* output, int channels) {
+  const int channel = blockIdx.x * blockDim.x + threadIdx.x;
+  const int batch = blockIdx.y;
+  if (channel >= channels) {
+    return;
+  }
+
+  const int state_offset = (batch * channels + channel) * kStateLen;
+  const int vector_offset = batch * channels + channel;
+  const int weight_offset = channel * kKernelSize;
+  const __half current = hidden[vector_offset];
+  const __half state0 = __float2half_rn(state[state_offset]);
+  const __half state1 = __float2half_rn(state[state_offset + 1]);
+  const __half state2 = __float2half_rn(state[state_offset + 2]);
+
+  float value = __half2float(state0) * __half2float(weight[weight_offset]);
+  value += __half2float(state1) * __half2float(weight[weight_offset + 1]);
+  value += __half2float(state2) * __half2float(weight[weight_offset + 2]);
+  value += __half2float(current) * __half2float(weight[weight_offset + 3]);
+
+  state[state_offset] = __half2float(state1);
+  state[state_offset + 1] = __half2float(state2);
+  state[state_offset + 2] = __half2float(current);
+  const __half convolved = __float2half_rn(value);
+  const float activation_input = __half2float(convolved);
+  output[vector_offset] = __float2half_rn(
+      activation_input / (1.0f + expf(-activation_input)));
+}
+
+void check_half_cuda_contiguous(const torch::Tensor& tensor,
+                                const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+}  // namespace
+
+torch::Tensor causal_conv_update(torch::Tensor state,
+                                 const torch::Tensor& hidden,
+                                 const torch::Tensor& weight) {
+  TORCH_CHECK(state.is_cuda(), "state must be a CUDA tensor");
+  TORCH_CHECK(state.scalar_type() == torch::kFloat32,
+              "state must have dtype float32");
+  TORCH_CHECK(state.is_contiguous(), "state must be contiguous");
+  check_half_cuda_contiguous(hidden, "hidden");
+  check_half_cuda_contiguous(weight, "weight");
+  TORCH_CHECK(state.dim() == 3 && state.size(2) == kStateLen,
+              "state must have shape (batch, channels, 3)");
+  TORCH_CHECK(hidden.dim() == 3 && hidden.size(2) == 1 &&
+                  hidden.size(0) == state.size(0) &&
+                  hidden.size(1) == state.size(1),
+              "hidden must have shape (batch, channels, 1)");
+  TORCH_CHECK(weight.dim() == 2 && weight.size(0) == state.size(1) &&
+                  weight.size(1) == kKernelSize,
+              "weight must have shape (channels, 4)");
+
+  auto output = torch::empty_like(hidden);
+  const int channels = static_cast<int>(state.size(1));
+  const dim3 blocks((channels + kThreads - 1) / kThreads,
+                    static_cast<unsigned int>(state.size(0)));
+  causal_conv_update_kernel<<<blocks, kThreads, 0,
+                              at::cuda::getCurrentCUDAStream()>>>(
+      state.data_ptr<float>(),
+      reinterpret_cast<const __half*>(hidden.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(weight.data_ptr<at::Half>()),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()), channels);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("causal_conv_update", &causal_conv_update,
+             "Fused CoreX Gated DeltaNet causal convolution update");
+}

qwen3_6_scripts/corex_gdn_gated_norm.cu

@@ -0,0 +1,80 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kHeadDim = 128;
+
+__device__ __forceinline__ float silu(float value) {
+  return value / (1.0f + expf(-value));
+}
+
+__global__ void gated_rms_norm_inverse_kernel(
+    const float* input, const __half* gate, const __half* weight,
+    const float* inverse, __half* output, int rows) {
+  const int row = blockIdx.x;
+  const int column = threadIdx.x;
+  if (row >= rows || column >= kHeadDim) {
+    return;
+  }
+  const int offset = row * kHeadDim + column;
+  const float scaled = __fmul_rn(input[offset], inverse[row]);
+  const float normalized = __fmul_rn(
+      __half2float(weight[column]), scaled);
+  const float activated = silu(__half2float(gate[offset]));
+  output[offset] = __float2half_rn(__fmul_rn(normalized, activated));
+}
+
+void check_input(const torch::Tensor& input, const torch::Tensor& gate,
+                 const torch::Tensor& weight,
+                 const torch::Tensor& inverse) {
+  TORCH_CHECK(input.is_cuda() && gate.is_cuda() && weight.is_cuda()
+                  && inverse.is_cuda(),
+              "all tensors must be CUDA tensors");
+  TORCH_CHECK(input.scalar_type() == torch::kFloat32,
+              "input must have dtype float32");
+  TORCH_CHECK(gate.scalar_type() == torch::kFloat16,
+              "gate must have dtype float16");
+  TORCH_CHECK(weight.scalar_type() == torch::kFloat16,
+              "weight must have dtype float16");
+  TORCH_CHECK(inverse.scalar_type() == torch::kFloat32,
+              "inverse must have dtype float32");
+  TORCH_CHECK(input.is_contiguous() && gate.is_contiguous()
+                  && weight.is_contiguous() && inverse.is_contiguous(),
+              "all tensors must be contiguous");
+  TORCH_CHECK(input.dim() == 2 && input.size(1) == kHeadDim,
+              "input must have shape (rows, 128)");
+  TORCH_CHECK(gate.sizes() == input.sizes(),
+              "gate must match input shape");
+  TORCH_CHECK(weight.dim() == 1 && weight.size(0) == kHeadDim,
+              "weight must have shape (128,)");
+  TORCH_CHECK(inverse.numel() == input.size(0),
+              "inverse must contain one value per row");
+}
+
+}  // namespace
+
+torch::Tensor apply_inverse(const torch::Tensor& input,
+                            const torch::Tensor& gate,
+                            const torch::Tensor& weight,
+                            const torch::Tensor& inverse) {
+  check_input(input, gate, weight, inverse);
+  auto output = torch::empty_like(gate);
+  const int rows = static_cast<int>(input.size(0));
+  gated_rms_norm_inverse_kernel<<<
+      rows, kHeadDim, 0, at::cuda::getCurrentCUDAStream()>>>(
+      input.data_ptr<float>(),
+      reinterpret_cast<const __half*>(gate.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(weight.data_ptr<at::Half>()),
+      inverse.data_ptr<float>(),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()), rows);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("apply_inverse", &apply_inverse,
+             "CoreX gated RMSNorm using a PyTorch-computed inverse");
+}

qwen3_6_scripts/corex_gdn_packed_decode.cu

@@ -0,0 +1,165 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kKeyHeads = 4;
+constexpr int kValueHeads = 8;
+constexpr int kHeadDim = 128;
+constexpr int kMixedDim =
+    (2 * kKeyHeads + kValueHeads) * kHeadDim;
+constexpr float kQueryScale = 0.08838834764831845f;
+
+__global__ void gdn_packed_decode_kernel(
+    float* state, const half* mixed_qkv, const half* beta_input,
+    const half* decay_input, const half* a_log, const half* dt_bias,
+    float* output) {
+  const int batch_head = blockIdx.x;
+  const int column = threadIdx.x;
+  const int batch = batch_head / kValueHeads;
+  const int value_head = batch_head % kValueHeads;
+  const int key_head = value_head / (kValueHeads / kKeyHeads);
+  const int mixed_offset = batch * kMixedDim;
+  const int query_offset = mixed_offset + key_head * kHeadDim;
+  const int key_offset =
+      mixed_offset + kKeyHeads * kHeadDim + key_head * kHeadDim;
+  const int value_offset = mixed_offset + 2 * kKeyHeads * kHeadDim
+                           + value_head * kHeadDim;
+  const int vector_offset = batch_head * kHeadDim;
+  const int state_offset = batch_head * kHeadDim * kHeadDim;
+
+  __shared__ half norm_squares[kHeadDim * 2];
+  __shared__ float normalized_query[kHeadDim];
+  __shared__ float normalized_key[kHeadDim];
+  const half raw_query = mixed_qkv[query_offset + column];
+  const half raw_key = mixed_qkv[key_offset + column];
+  norm_squares[column] = __hmul(raw_query, raw_query);
+  norm_squares[kHeadDim + column] = __hmul(raw_key, raw_key);
+  __syncthreads();
+
+  for (int stride = kHeadDim / 2; stride > 0; stride >>= 1) {
+    if (column < stride) {
+      norm_squares[column] = __hadd(
+          norm_squares[column], norm_squares[column + stride]);
+      norm_squares[kHeadDim + column] = __hadd(
+          norm_squares[kHeadDim + column],
+          norm_squares[kHeadDim + column + stride]);
+    }
+    __syncthreads();
+  }
+
+  const half epsilon = __float2half(1e-6f);
+  const half query_inverse = __float2half(rsqrtf(__half2float(
+      __hadd(norm_squares[0], epsilon))));
+  const half key_inverse = __float2half(rsqrtf(__half2float(
+      __hadd(norm_squares[kHeadDim], epsilon))));
+  normalized_query[column] = __half2float(
+      __hmul(raw_query, query_inverse)) * kQueryScale;
+  normalized_key[column] = __half2float(__hmul(raw_key, key_inverse));
+  __syncthreads();
+
+  const int coefficient_offset = batch * kValueHeads + value_head;
+  const float beta_value = __half2float(beta_input[coefficient_offset]);
+  const float beta = __half2float(__float2half(
+      1.0f / (1.0f + expf(-beta_value))));
+  const float decay_x = __half2float(decay_input[coefficient_offset])
+                        + __half2float(dt_bias[value_head]);
+  const float softplus =
+      decay_x > 20.0f ? decay_x : log1pf(expf(decay_x));
+  const float decay = expf(
+      -expf(__half2float(a_log[value_head])) * softplus);
+
+  float memory = 0.0f;
+#pragma unroll
+  for (int row = 0; row < kHeadDim; ++row) {
+    const int index = state_offset + row * kHeadDim + column;
+    const float decayed = state[index] * decay;
+    memory += normalized_key[row] * decayed;
+  }
+
+  const float value = __half2float(mixed_qkv[value_offset + column]);
+  const float delta = (value - memory) * beta;
+  float result = 0.0f;
+#pragma unroll
+  for (int row = 0; row < kHeadDim; ++row) {
+    const int index = state_offset + row * kHeadDim + column;
+    const float decayed = state[index] * decay;
+    const float updated = decayed + normalized_key[row] * delta;
+    state[index] = updated;
+    result += normalized_query[row] * updated;
+  }
+  output[vector_offset + column] = result;
+}
+
+void check_half_matrix(const torch::Tensor& tensor, const char* name,
+                       int64_t width) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 2 && tensor.size(1) == width,
+              name, " must have shape (batch, ", width, ")");
+}
+
+void check_half_vector(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 1 && tensor.size(0) == kValueHeads,
+              name, " must have shape (", kValueHeads, ")");
+}
+
+}  // namespace
+
+torch::Tensor packed_decode(torch::Tensor state,
+                            const torch::Tensor& mixed_qkv,
+                            const torch::Tensor& beta_input,
+                            const torch::Tensor& decay_input,
+                            const torch::Tensor& a_log,
+                            const torch::Tensor& dt_bias) {
+  TORCH_CHECK(state.is_cuda(), "state must be a CUDA tensor");
+  TORCH_CHECK(state.scalar_type() == torch::kFloat32,
+              "state must have dtype float32");
+  TORCH_CHECK(state.is_contiguous(), "state must be contiguous");
+  TORCH_CHECK(state.dim() == 4 && state.size(0) == 1
+                  && state.size(1) == kValueHeads
+                  && state.size(2) == kHeadDim
+                  && state.size(3) == kHeadDim,
+              "state must have shape (1, 8, 128, 128)");
+  check_half_matrix(mixed_qkv, "mixed_qkv", kMixedDim);
+  check_half_matrix(beta_input, "beta_input", kValueHeads);
+  check_half_matrix(decay_input, "decay_input", kValueHeads);
+  check_half_vector(a_log, "a_log");
+  check_half_vector(dt_bias, "dt_bias");
+  TORCH_CHECK(mixed_qkv.size(0) == 1 && beta_input.size(0) == 1
+                  && decay_input.size(0) == 1,
+              "packed decode only supports one sequence");
+  TORCH_CHECK(state.device() == mixed_qkv.device()
+                  && state.device() == beta_input.device()
+                  && state.device() == decay_input.device()
+                  && state.device() == a_log.device()
+                  && state.device() == dt_bias.device(),
+              "all inputs must be on the same device");
+
+  torch::Tensor output = torch::empty(
+      {1, kValueHeads, kHeadDim}, state.options());
+  gdn_packed_decode_kernel<<<kValueHeads, kHeadDim, 0,
+                             at::cuda::getCurrentCUDAStream()>>>(
+      state.data_ptr<float>(),
+      reinterpret_cast<const half*>(mixed_qkv.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(beta_input.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(decay_input.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(a_log.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(dt_bias.data_ptr<at::Half>()),
+      output.data_ptr<float>());
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("packed_decode", &packed_decode,
+             "Packed Qwen3.6 GDN single-token decode");
+}

qwen3_6_scripts/corex_gdn_qk_map.cu

@@ -0,0 +1,72 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kHeadDim = 128;
+constexpr float kQueryScale = 0.08838834764831845f;
+
+__global__ void qk_map_kernel(const half* query, const half* key,
+                              float* output, int batch, int key_heads,
+                              int value_heads, int expand_ratio) {
+  const int elements = batch * value_heads * kHeadDim;
+  const int index = blockIdx.x * blockDim.x + threadIdx.x;
+  if (index >= elements) {
+    return;
+  }
+  const int dim = index % kHeadDim;
+  const int value_head_index = index / kHeadDim;
+  const int value_head = value_head_index % value_heads;
+  const int batch_index = value_head_index / value_heads;
+  const int key_head = value_head / expand_ratio;
+  const int source = ((batch_index * key_heads + key_head) * kHeadDim + dim);
+  output[index] = __half2float(query[source]) * kQueryScale;
+  output[elements + index] = __half2float(key[source]);
+}
+
+void check_input(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 3 && tensor.size(2) == kHeadDim,
+              name, " must have shape (batch, key_heads, 128)");
+}
+
+}  // namespace
+
+torch::Tensor qk_map(const torch::Tensor& query,
+                     const torch::Tensor& key,
+                     int64_t value_heads_arg) {
+  check_input(query, "query");
+  check_input(key, "key");
+  TORCH_CHECK(query.sizes() == key.sizes(),
+              "query and key shapes must match");
+  const int batch = static_cast<int>(query.size(0));
+  const int key_heads = static_cast<int>(query.size(1));
+  const int value_heads = static_cast<int>(value_heads_arg);
+  TORCH_CHECK(value_heads > 0 && value_heads % key_heads == 0,
+              "value_heads must be divisible by key_heads");
+
+  torch::Tensor output = torch::empty(
+      {2, batch, value_heads, kHeadDim},
+      query.options().dtype(torch::kFloat32));
+  const int elements = batch * value_heads * kHeadDim;
+  constexpr int threads = 256;
+  const int blocks = (elements + threads - 1) / threads;
+  qk_map_kernel<<<blocks, threads, 0,
+                  at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const half*>(query.data_ptr<at::Half>()),
+      reinterpret_cast<const half*>(key.data_ptr<at::Half>()),
+      output.data_ptr<float>(), batch, key_heads, value_heads,
+      value_heads / key_heads);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("qk_map", &qk_map,
+             "Map normalized FP16 key heads to FP32 value heads");
+}

qwen3_6_scripts/corex_moe_direct_routed.cu

@@ -0,0 +1,181 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kExperts = 256;
+constexpr int kTopK = 8;
+constexpr int kHidden = 2048;
+constexpr int kIntermediate = 128;
+constexpr int kW13Rows = 2 * kIntermediate;
+constexpr int kThreads = 256;
+constexpr int kWarpSize = 32;
+
+__device__ inline float warp_sum(float value) {
+#pragma unroll
+  for (int offset = kWarpSize / 2; offset > 0; offset /= 2) {
+    value += __shfl_down_sync(0xffffffff, value, offset);
+  }
+  return value;
+}
+
+__global__ void direct_w13_kernel(
+    const __half* input, const __half* w13, const int64_t* expert_ids,
+    __half* gate_up) {
+  const int warp =
+      (static_cast<int>(blockIdx.x) * blockDim.x + threadIdx.x) / kWarpSize;
+  const int lane = threadIdx.x & (kWarpSize - 1);
+  if (warp >= kTopK * kW13Rows) {
+    return;
+  }
+
+  const int slot = warp / kW13Rows;
+  const int local_row = warp - slot * kW13Rows;
+  const int64_t expert = expert_ids[slot];
+  const int64_t weight_row =
+      (expert * kW13Rows + local_row) * static_cast<int64_t>(kHidden);
+  const __half2* input2 = reinterpret_cast<const __half2*>(input);
+  const __half2* weight2 =
+      reinterpret_cast<const __half2*>(w13 + weight_row);
+  float sum = 0.0f;
+  for (int index = lane; index < kHidden / 2; index += kWarpSize) {
+    const __half2 x = input2[index];
+    const __half2 weight = weight2[index];
+    sum = fmaf(__half2float(weight.x), __half2float(x.x), sum);
+    sum = fmaf(__half2float(weight.y), __half2float(x.y), sum);
+  }
+  sum = warp_sum(sum);
+  if (lane == 0) {
+    gate_up[warp] = __float2half_rn(sum);
+  }
+}
+
+__global__ void direct_w2_reduce_kernel(
+    const __half* activated, const __half* w2, const int64_t* expert_ids,
+    const __half* weights, __half* output) {
+  const int warp =
+      (static_cast<int>(blockIdx.x) * blockDim.x + threadIdx.x) / kWarpSize;
+  const int lane = threadIdx.x & (kWarpSize - 1);
+  if (warp >= kHidden) {
+    return;
+  }
+
+  float weighted_sum = 0.0f;
+#pragma unroll
+  for (int slot = 0; slot < kTopK; ++slot) {
+    const int64_t expert = expert_ids[slot];
+    const int64_t weight_row =
+        (expert * kHidden + warp) * static_cast<int64_t>(kIntermediate);
+    const __half2* activation2 = reinterpret_cast<const __half2*>(
+        activated + slot * kIntermediate);
+    const __half2* weight2 =
+        reinterpret_cast<const __half2*>(w2 + weight_row);
+    float expert_sum = 0.0f;
+    for (int index = lane; index < kIntermediate / 2;
+         index += kWarpSize) {
+      const __half2 x = activation2[index];
+      const __half2 weight = weight2[index];
+      expert_sum = fmaf(
+          __half2float(weight.x), __half2float(x.x), expert_sum);
+      expert_sum = fmaf(
+          __half2float(weight.y), __half2float(x.y), expert_sum);
+    }
+    expert_sum = warp_sum(expert_sum);
+    if (lane == 0) {
+      const __half expert_half = __float2half_rn(expert_sum);
+      const __half product = __hmul(expert_half, weights[slot]);
+      weighted_sum += __half2float(product);
+    }
+  }
+  if (lane == 0) {
+    output[warp] = __float2half_rn(weighted_sum);
+  }
+}
+
+void check_half_cuda(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+void check_ids(const torch::Tensor& expert_ids) {
+  TORCH_CHECK(expert_ids.is_cuda() && expert_ids.is_contiguous(),
+              "expert_ids must be a contiguous CUDA tensor");
+  TORCH_CHECK(expert_ids.scalar_type() == torch::kInt64,
+              "expert_ids must have dtype int64");
+  TORCH_CHECK(expert_ids.dim() == 1 && expert_ids.numel() == kTopK,
+              "expert_ids must have shape (8,)");
+}
+
+}  // namespace
+
+torch::Tensor direct_w13(const torch::Tensor& input,
+                         const torch::Tensor& w13,
+                         const torch::Tensor& expert_ids) {
+  check_half_cuda(input, "input");
+  check_half_cuda(w13, "w13");
+  check_ids(expert_ids);
+  TORCH_CHECK(input.dim() == 2 && input.size(0) == 1
+                  && input.size(1) == kHidden,
+              "input must have shape (1, 2048)");
+  TORCH_CHECK(w13.dim() == 3 && w13.size(0) == kExperts
+                  && w13.size(1) == kW13Rows
+                  && w13.size(2) == kHidden,
+              "w13 must have shape (256, 256, 2048)");
+
+  auto output = torch::empty({kTopK, kW13Rows}, input.options());
+  constexpr int kWarpsPerBlock = kThreads / kWarpSize;
+  constexpr int kBlocks =
+      (kTopK * kW13Rows + kWarpsPerBlock - 1) / kWarpsPerBlock;
+  direct_w13_kernel<<<kBlocks, kThreads, 0,
+                      at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const __half*>(input.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(w13.data_ptr<at::Half>()),
+      expert_ids.data_ptr<int64_t>(),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+torch::Tensor direct_w2_reduce(const torch::Tensor& activated,
+                               const torch::Tensor& w2,
+                               const torch::Tensor& expert_ids,
+                               const torch::Tensor& weights) {
+  check_half_cuda(activated, "activated");
+  check_half_cuda(w2, "w2");
+  check_half_cuda(weights, "weights");
+  check_ids(expert_ids);
+  TORCH_CHECK(activated.dim() == 2 && activated.size(0) == kTopK
+                  && activated.size(1) == kIntermediate,
+              "activated must have shape (8, 128)");
+  TORCH_CHECK(w2.dim() == 3 && w2.size(0) == kExperts
+                  && w2.size(1) == kHidden
+                  && w2.size(2) == kIntermediate,
+              "w2 must have shape (256, 2048, 128)");
+  TORCH_CHECK(weights.dim() == 1 && weights.numel() == kTopK,
+              "weights must have shape (8,)");
+
+  auto output = torch::empty({1, kHidden}, activated.options());
+  constexpr int kWarpsPerBlock = kThreads / kWarpSize;
+  constexpr int kBlocks =
+      (kHidden + kWarpsPerBlock - 1) / kWarpsPerBlock;
+  direct_w2_reduce_kernel<<<kBlocks, kThreads, 0,
+                            at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const __half*>(activated.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(w2.data_ptr<at::Half>()),
+      expert_ids.data_ptr<int64_t>(),
+      reinterpret_cast<const __half*>(weights.data_ptr<at::Half>()),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("w13", &direct_w13,
+             "Direct selected-expert FP16 W13 matvec");
+  module.def("w2_reduce", &direct_w2_reduce,
+             "Direct selected-expert W2 matvec and routed reduction");
+}

qwen3_6_scripts/corex_moe_exact_reduce.cu

@@ -0,0 +1,107 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+namespace {
+
+constexpr int kTopK = 8;
+constexpr int kThreads = 256;
+
+enum class Mode { kSerialFloat, kTreeFloat, kSerialHalf };
+
+__global__ void exact_reduce_kernel(const __half* expert_output,
+                                    const __half* weights,
+                                    __half* output, int hidden,
+                                    Mode mode) {
+  const int column = blockIdx.x * blockDim.x + threadIdx.x;
+  if (column >= hidden) {
+    return;
+  }
+  __half products[kTopK];
+#pragma unroll
+  for (int expert = 0; expert < kTopK; ++expert) {
+    products[expert] = __hmul(
+        expert_output[expert * hidden + column], weights[expert]);
+  }
+  if (mode == Mode::kSerialHalf) {
+    __half sum = products[0];
+#pragma unroll
+    for (int expert = 1; expert < kTopK; ++expert) {
+      sum = __hadd(sum, products[expert]);
+    }
+    output[column] = sum;
+    return;
+  }
+
+  float sum;
+  if (mode == Mode::kSerialFloat) {
+    sum = __half2float(products[0]);
+#pragma unroll
+    for (int expert = 1; expert < kTopK; ++expert) {
+      sum += __half2float(products[expert]);
+    }
+  } else {
+    const float sum01 = __half2float(products[0]) + __half2float(products[1]);
+    const float sum23 = __half2float(products[2]) + __half2float(products[3]);
+    const float sum45 = __half2float(products[4]) + __half2float(products[5]);
+    const float sum67 = __half2float(products[6]) + __half2float(products[7]);
+    sum = (sum01 + sum23) + (sum45 + sum67);
+  }
+  output[column] = __float2half_rn(sum);
+}
+
+void check_input(const torch::Tensor& expert_output,
+                 const torch::Tensor& weights) {
+  TORCH_CHECK(expert_output.is_cuda() && weights.is_cuda(),
+              "inputs must be CUDA tensors");
+  TORCH_CHECK(expert_output.scalar_type() == torch::kFloat16
+                  && weights.scalar_type() == torch::kFloat16,
+              "inputs must have dtype float16");
+  TORCH_CHECK(expert_output.is_contiguous() && weights.is_contiguous(),
+              "inputs must be contiguous");
+  TORCH_CHECK(expert_output.dim() == 2
+                  && expert_output.size(0) == kTopK,
+              "expert_output must have shape (8, hidden)");
+  TORCH_CHECK(weights.dim() == 1 && weights.size(0) == kTopK,
+              "weights must have shape (8,)");
+}
+
+torch::Tensor launch(const torch::Tensor& expert_output,
+                     const torch::Tensor& weights, Mode mode) {
+  check_input(expert_output, weights);
+  auto output = torch::empty(
+      {1, expert_output.size(1)}, expert_output.options());
+  const int hidden = static_cast<int>(expert_output.size(1));
+  const int blocks = (hidden + kThreads - 1) / kThreads;
+  exact_reduce_kernel<<<blocks, kThreads, 0,
+                        at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const __half*>(expert_output.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(weights.data_ptr<at::Half>()),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()), hidden, mode);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return output;
+}
+
+}  // namespace
+
+torch::Tensor serial_float(const torch::Tensor& expert_output,
+                           const torch::Tensor& weights) {
+  return launch(expert_output, weights, Mode::kSerialFloat);
+}
+
+torch::Tensor tree_float(const torch::Tensor& expert_output,
+                         const torch::Tensor& weights) {
+  return launch(expert_output, weights, Mode::kTreeFloat);
+}
+
+torch::Tensor serial_half(const torch::Tensor& expert_output,
+                          const torch::Tensor& weights) {
+  return launch(expert_output, weights, Mode::kSerialHalf);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("serial_float", &serial_float);
+  module.def("tree_float", &tree_float);
+  module.def("serial_half", &serial_half);
+}

qwen3_6_scripts/corex_moe_weight_gather.cu

@@ -0,0 +1,92 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <cstdint>
+#include <vector>
+
+namespace {
+
+constexpr int kTopK = 8;
+constexpr int kThreads = 256;
+constexpr int kGridX = 8;
+
+__global__ void selected_weight_gather_vec16_kernel(
+    const uint4* w13, const uint4* w2, const int64_t* expert_ids,
+    uint4* selected_w13, uint4* selected_w2,
+    int64_t w13_vecs_per_expert, int64_t w2_vecs_per_expert) {
+  const int segment = blockIdx.y;
+  const int slot = segment & (kTopK - 1);
+  const bool copy_w2 = segment >= kTopK;
+  const int64_t count =
+      copy_w2 ? w2_vecs_per_expert : w13_vecs_per_expert;
+  const uint4* source = copy_w2 ? w2 : w13;
+  uint4* output = copy_w2 ? selected_w2 : selected_w13;
+  const int64_t source_offset = expert_ids[slot] * count;
+  const int64_t output_offset = static_cast<int64_t>(slot) * count;
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < count;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    output[output_offset + index] = source[source_offset + index];
+  }
+}
+
+void check_weight(const torch::Tensor& tensor, const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+  TORCH_CHECK(tensor.dim() == 3, name, " must be rank three");
+  TORCH_CHECK(tensor.size(1) * tensor.size(2) % 8 == 0,
+              name, " expert slices must be divisible by 16 bytes");
+}
+
+}  // namespace
+
+std::vector<torch::Tensor> gather_selected_weights(
+    const torch::Tensor& w13, const torch::Tensor& w2,
+    const torch::Tensor& expert_ids) {
+  check_weight(w13, "w13");
+  check_weight(w2, "w2");
+  TORCH_CHECK(w13.device() == w2.device(),
+              "W13/W2 must be on the same device");
+  TORCH_CHECK(w13.size(0) == w2.size(0),
+              "W13/W2 expert counts differ");
+  TORCH_CHECK(w13.size(2) == w2.size(1),
+              "W13/W2 hidden dimensions differ");
+  TORCH_CHECK(w13.size(1) == 2 * w2.size(2),
+              "W13/W2 intermediate dimensions differ");
+  TORCH_CHECK(expert_ids.is_cuda() && expert_ids.is_contiguous(),
+              "expert_ids must be a contiguous CUDA tensor");
+  TORCH_CHECK(expert_ids.device() == w13.device(),
+              "weights and expert_ids must be on the same device");
+  TORCH_CHECK(expert_ids.scalar_type() == torch::kInt64,
+              "expert_ids must have dtype int64");
+  TORCH_CHECK(expert_ids.dim() == 1 && expert_ids.numel() == kTopK,
+              "expert_ids must have shape (8,)");
+
+  auto selected_w13 = torch::empty(
+      {kTopK, w13.size(1), w13.size(2)}, w13.options());
+  auto selected_w2 = torch::empty(
+      {kTopK, w2.size(1), w2.size(2)}, w2.options());
+  const int64_t w13_vecs_per_expert = w13.size(1) * w13.size(2) / 8;
+  const int64_t w2_vecs_per_expert = w2.size(1) * w2.size(2) / 8;
+  const dim3 grid(kGridX, 2 * kTopK);
+  selected_weight_gather_vec16_kernel<<<
+      grid, kThreads, 0, at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const uint4*>(w13.data_ptr<at::Half>()),
+      reinterpret_cast<const uint4*>(w2.data_ptr<at::Half>()),
+      expert_ids.data_ptr<int64_t>(),
+      reinterpret_cast<uint4*>(selected_w13.data_ptr<at::Half>()),
+      reinterpret_cast<uint4*>(selected_w2.data_ptr<at::Half>()),
+      w13_vecs_per_expert, w2_vecs_per_expert);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {selected_w13, selected_w2};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("gather", &gather_selected_weights,
+             "Gather selected FP16 top-8 MoE weights with 16-byte loads");
+}

qwen3_6_scripts/corex_paged_kv_gather.cu

@@ -0,0 +1,118 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <vector>
+
+namespace {
+
+constexpr int kThreads = 256;
+constexpr int kSmallGridBlocks = 256;
+constexpr int kSmallGridMaxSeqLen = 96 * 1024;
+
+__global__ void paged_kv_gather_kernel(
+    const __half* key_cache, const __half* value_cache,
+    const int* block_table, float* key_output, float* value_output,
+    int seq_len, int num_kv_heads, int head_size, int block_size,
+    int key_pack) {
+  const int64_t total =
+      static_cast<int64_t>(seq_len) * num_kv_heads * head_size;
+  for (int64_t index =
+           static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
+       index < total;
+       index += static_cast<int64_t>(blockDim.x) * gridDim.x) {
+    const int dim = index % head_size;
+    const int token = (index / head_size) % seq_len;
+    const int kv_head = index / (static_cast<int64_t>(head_size) * seq_len);
+    const int logical_block = token / block_size;
+    const int block_offset = token % block_size;
+    const int physical_block = block_table[logical_block];
+
+    const int64_t key_index =
+        (((static_cast<int64_t>(physical_block) * num_kv_heads + kv_head)
+           * (head_size / key_pack) + dim / key_pack)
+          * block_size + block_offset) * key_pack + dim % key_pack;
+    const int64_t value_index =
+        ((static_cast<int64_t>(physical_block) * num_kv_heads + kv_head)
+         * head_size + dim) * block_size + block_offset;
+    const int64_t key_output_index =
+        (static_cast<int64_t>(kv_head) * head_size + dim) * seq_len + token;
+    const int64_t value_output_index =
+        (static_cast<int64_t>(kv_head) * seq_len + token) * head_size + dim;
+
+    key_output[key_output_index] = __half2float(key_cache[key_index]);
+    value_output[value_output_index] = __half2float(value_cache[value_index]);
+  }
+}
+
+void check_half_cuda_contiguous(const torch::Tensor& tensor,
+                                const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+}  // namespace
+
+std::vector<torch::Tensor> gather_paged_kv(
+    const torch::Tensor& key_cache, const torch::Tensor& value_cache,
+    const torch::Tensor& block_table, int64_t seq_len) {
+  check_half_cuda_contiguous(key_cache, "key_cache");
+  check_half_cuda_contiguous(value_cache, "value_cache");
+  TORCH_CHECK(block_table.is_cuda(), "block_table must be a CUDA tensor");
+  TORCH_CHECK(block_table.scalar_type() == torch::kInt32,
+              "block_table must have dtype int32");
+  TORCH_CHECK(block_table.is_contiguous(), "block_table must be contiguous");
+  TORCH_CHECK(key_cache.dim() == 5,
+              "key_cache must have shape (blocks, kv_heads, d/x, block, x)");
+  TORCH_CHECK(value_cache.dim() == 4,
+              "value_cache must have shape (blocks, kv_heads, d, block)");
+  TORCH_CHECK(block_table.dim() == 1,
+              "block_table must be a one-dimensional row");
+  TORCH_CHECK(key_cache.size(0) == value_cache.size(0),
+              "key/value block counts differ");
+  TORCH_CHECK(key_cache.size(1) == value_cache.size(1),
+              "key/value KV-head counts differ");
+  TORCH_CHECK(key_cache.size(3) == value_cache.size(3),
+              "key/value block sizes differ");
+  TORCH_CHECK(key_cache.size(2) * key_cache.size(4) == value_cache.size(2),
+              "key/value head sizes differ");
+  TORCH_CHECK(seq_len > 0, "seq_len must be positive");
+
+  const int block_size = static_cast<int>(value_cache.size(3));
+  const int64_t required_blocks = (seq_len + block_size - 1) / block_size;
+  TORCH_CHECK(required_blocks <= block_table.numel(),
+              "block_table is too short for seq_len");
+  const int num_kv_heads = static_cast<int>(value_cache.size(1));
+  const int head_size = static_cast<int>(value_cache.size(2));
+  const int key_pack = static_cast<int>(key_cache.size(4));
+
+  auto output_options = key_cache.options().dtype(torch::kFloat32);
+  auto key_output = torch::empty(
+      {num_kv_heads, head_size, seq_len}, output_options);
+  auto value_output = torch::empty(
+      {num_kv_heads, seq_len, head_size}, output_options);
+  const int64_t total = seq_len * num_kv_heads * head_size;
+  const int grid_cap =
+      seq_len <= kSmallGridMaxSeqLen ? kSmallGridBlocks : 65535;
+  const int blocks = static_cast<int>(std::min<int64_t>(
+      (total + kThreads - 1) / kThreads, grid_cap));
+  paged_kv_gather_kernel<<<blocks, kThreads, 0,
+                           at::cuda::getCurrentCUDAStream()>>>(
+      reinterpret_cast<const __half*>(key_cache.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(value_cache.data_ptr<at::Half>()),
+      block_table.data_ptr<int>(), key_output.data_ptr<float>(),
+      value_output.data_ptr<float>(), static_cast<int>(seq_len),
+      num_kv_heads, head_size, block_size, key_pack);
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {key_output, value_output};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("gather", &gather_paged_kv,
+             "Gather paged FP16 K/V directly into FP32 attention layouts");
+}

qwen3_6_scripts/corex_query_tiled_paged_prefill.cu

@@ -0,0 +1,503 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAException.h>
+#include <cuda_fp16.h>
+#include <mma.h>
+#include <torch/extension.h>
+
+#include <algorithm>
+#include <cmath>
+#include <cstdint>
+#include <limits>
+#include <vector>
+
+namespace {
+
+constexpr int kBlockSize = 16;
+constexpr int kHeadDim = 256;
+constexpr int kKeyPack = 8;
+constexpr int kNumQueryHeads = 4;
+constexpr int kNumKvHeads = 1;
+constexpr int kQueryTile = 16;
+constexpr int kKeyTile = 16;
+constexpr int kReductionTokens = 512;
+constexpr int kKeyTilesPerReduction = kReductionTokens / kKeyTile;
+constexpr int kPvReductionSplits = 4;
+constexpr int kKeyTilesPerPvSplit =
+    kKeyTilesPerReduction / kPvReductionSplits;
+constexpr int kMmaK = 16;
+constexpr int kDimTiles = kHeadDim / kMmaK;
+constexpr int kWarpSize = 64;
+constexpr int kMaxQueryTokens = 8192;
+constexpr int kMaxSequenceTokens = 262144;
+
+using namespace nvcuda;
+
+struct __align__(128) SharedStorage {
+  float matrix_tile[kQueryTile * kKeyTile];
+  float scores[
+      kKeyTilesPerReduction * kQueryTile * kKeyTile];
+  float running_output[kQueryTile * kHeadDim];
+  float partial_output[
+      kPvReductionSplits * kQueryTile * kMmaK];
+  float running_max[kQueryTile];
+  float running_sum[kQueryTile];
+  float correction[kQueryTile];
+};
+
+void check_half_cuda_contiguous(const torch::Tensor& tensor,
+                                const char* name) {
+  TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor");
+  TORCH_CHECK(tensor.scalar_type() == torch::kFloat16,
+              name, " must have dtype float16");
+  TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous");
+}
+
+__device__ __forceinline__ float load_key(
+    const __half* key_new, const __half* key_cache,
+    const int* block_table, int logical_token, int context_len, int dim) {
+  if (logical_token < context_len) {
+    const int logical_block = logical_token / kBlockSize;
+    const int block_offset = logical_token % kBlockSize;
+    const int physical_block = block_table[logical_block];
+    const int64_t index =
+        ((((static_cast<int64_t>(physical_block) * kNumKvHeads)
+            * (kHeadDim / kKeyPack) + dim / kKeyPack)
+           * kBlockSize + block_offset) * kKeyPack + dim % kKeyPack);
+    return __half2float(key_cache[index]);
+  }
+  const int query_index = logical_token - context_len;
+  return __half2float(
+      key_new[static_cast<int64_t>(query_index) * kHeadDim + dim]);
+}
+
+__device__ __forceinline__ float load_value(
+    const __half* value_new, const __half* value_cache,
+    const int* block_table, int logical_token, int context_len, int dim) {
+  if (logical_token < context_len) {
+    const int logical_block = logical_token / kBlockSize;
+    const int block_offset = logical_token % kBlockSize;
+    const int physical_block = block_table[logical_block];
+    const int64_t index =
+        ((static_cast<int64_t>(physical_block) * kNumKvHeads)
+          * kHeadDim + dim) * kBlockSize + block_offset;
+    return __half2float(value_cache[index]);
+  }
+  const int query_index = logical_token - context_len;
+  return __half2float(
+      value_new[static_cast<int64_t>(query_index) * kHeadDim + dim]);
+}
+
+__global__ void query_tiled_paged_prefill_kernel(
+    const __half* query, const __half* key_new, const __half* value_new,
+    const __half* key_cache, const __half* value_cache,
+    const int* block_table, __half* output, float* lse,
+    int context_len, int query_len, float scale) {
+  __shared__ SharedStorage shared;
+
+  const int lane = threadIdx.x;
+  const int query_tile_index = blockIdx.x / kNumQueryHeads;
+  const int query_head = blockIdx.x % kNumQueryHeads;
+  const int query_start = query_tile_index * kQueryTile;
+  const int active_rows = min(kQueryTile, query_len - query_start);
+  if (active_rows <= 0) {
+    return;
+  }
+
+  wmma::fragment<wmma::matrix_a, 16, 16, 16, float,
+                 wmma::row_major> query_fragments[kDimTiles];
+
+#pragma unroll
+  for (int dim_tile = 0; dim_tile < kDimTiles; ++dim_tile) {
+#pragma unroll
+    for (int quarter = 0; quarter < 4; ++quarter) {
+      const int row = lane / 16 + quarter * 4;
+      const int column = lane % 16;
+      float value = 0.0f;
+      if (row < active_rows) {
+        const int query_index = query_start + row;
+        const int dim = dim_tile * kMmaK + column;
+        const int64_t source =
+            (static_cast<int64_t>(query_index) * kNumQueryHeads
+              + query_head) * kHeadDim + dim;
+        value = __half2float(query[source]) * scale;
+      }
+      const int offset =
+          wmma::CoordToOffset<32, wmma::layout_t::mem_row_major>(
+              row, column);
+      shared.matrix_tile[offset] = value;
+    }
+    __syncthreads();
+    wmma::load_matrix_sync(
+        query_fragments[dim_tile], shared.matrix_tile, 0);
+    __syncthreads();
+  }
+
+  for (int index = lane; index < kQueryTile * kHeadDim;
+       index += kWarpSize) {
+    shared.running_output[index] = 0.0f;
+  }
+  if (lane < kQueryTile) {
+    shared.running_max[lane] = -std::numeric_limits<float>::infinity();
+    shared.running_sum[lane] = 0.0f;
+    shared.correction[lane] = 1.0f;
+  }
+  __syncthreads();
+
+  const int last_query = min(query_start + kQueryTile, query_len);
+
+  // Preserve the installed reference's 512-token reduction boundaries:
+  // paged context and current causal K/V are separate phases.
+  for (int phase = 0; phase < 2; ++phase) {
+    const int phase_base = phase == 0 ? 0 : context_len;
+    const int phase_tokens = phase == 0 ? context_len : last_query;
+    for (int group_start = 0; group_start < phase_tokens;
+         group_start += kReductionTokens) {
+      const int group_tokens =
+          min(kReductionTokens, phase_tokens - group_start);
+      const int group_key_tiles =
+          (group_tokens + kKeyTile - 1) / kKeyTile;
+
+      for (int key_tile_in_group = 0;
+           key_tile_in_group < group_key_tiles;
+           ++key_tile_in_group) {
+        const int local_key_start =
+            group_start + key_tile_in_group * kKeyTile;
+        const int logical_key_start = phase_base + local_key_start;
+        wmma::fragment<wmma::accumulator, 16, 16, 16, float>
+            score_fragment;
+        wmma::fill_fragment(score_fragment, 0.0f);
+
+#pragma unroll
+        for (int dim_tile = 0; dim_tile < kDimTiles; ++dim_tile) {
+#pragma unroll
+          for (int quarter = 0; quarter < 4; ++quarter) {
+            const int row = lane / 16 + quarter * 4;
+            const int column = lane % 16;
+            const int logical_token = logical_key_start + column;
+            const int dim = dim_tile * kMmaK + row;
+            const float value =
+                local_key_start + column < phase_tokens
+                    ? load_key(key_new, key_cache, block_table,
+                               logical_token, context_len, dim)
+                    : 0.0f;
+            const int offset =
+                wmma::CoordToOffset<
+                    32, wmma::layout_t::mem_col_major>(
+                    row, column);
+            shared.matrix_tile[offset] = value;
+          }
+          __syncthreads();
+          wmma::fragment<wmma::matrix_b, 16, 16, 16, float,
+                         wmma::col_major> key_fragment;
+          wmma::load_matrix_sync(
+              key_fragment, shared.matrix_tile, 0);
+          wmma::mma_sync(
+              score_fragment,
+              query_fragments[dim_tile],
+              key_fragment,
+              score_fragment);
+          __syncthreads();
+        }
+
+        float* score_tile =
+            shared.scores
+            + key_tile_in_group * kQueryTile * kKeyTile;
+        wmma::store_matrix_sync(
+            score_tile, score_fragment, 0, wmma::mem_row_major);
+        __syncthreads();
+      }
+
+      if (lane < kQueryTile) {
+        const int row = lane;
+        if (row >= active_rows) {
+          shared.correction[row] = 1.0f;
+          for (int key_offset = 0; key_offset < group_tokens;
+               ++key_offset) {
+            const int key_tile = key_offset / kKeyTile;
+            const int column = key_offset % kKeyTile;
+            shared.scores[
+                key_tile * kQueryTile * kKeyTile
+                + row * kKeyTile + column] = 0.0f;
+          }
+        } else {
+          const int absolute_query =
+              context_len + query_start + row;
+          float block_max =
+              -std::numeric_limits<float>::infinity();
+          for (int key_offset = 0; key_offset < group_tokens;
+               ++key_offset) {
+            const int key_tile = key_offset / kKeyTile;
+            const int column = key_offset % kKeyTile;
+            const int score_index =
+                key_tile * kQueryTile * kKeyTile
+                + row * kKeyTile + column;
+            const int logical_key =
+                phase_base + group_start + key_offset;
+            if (logical_key <= absolute_query) {
+              block_max = fmaxf(
+                  block_max, shared.scores[score_index]);
+            } else {
+              shared.scores[score_index] =
+                  -std::numeric_limits<float>::infinity();
+            }
+          }
+          const float old_max = shared.running_max[row];
+          const float new_max = fmaxf(old_max, block_max);
+          const float correction =
+              old_max == -std::numeric_limits<float>::infinity()
+                  ? 0.0f
+                  : expf(old_max - new_max);
+          float group_sum = 0.0f;
+          for (int key_offset = 0; key_offset < group_tokens;
+               ++key_offset) {
+            const int key_tile = key_offset / kKeyTile;
+            const int column = key_offset % kKeyTile;
+            const int score_index =
+                key_tile * kQueryTile * kKeyTile
+                + row * kKeyTile + column;
+            const float score = shared.scores[score_index];
+            const float probability =
+                score == -std::numeric_limits<float>::infinity()
+                    ? 0.0f
+                    : expf(score - new_max);
+            shared.scores[score_index] = probability;
+            group_sum += probability;
+          }
+          shared.running_sum[row] =
+              shared.running_sum[row] * correction + group_sum;
+          shared.running_max[row] = new_max;
+          shared.correction[row] = correction;
+        }
+        for (int key_offset = group_tokens;
+             key_offset < group_key_tiles * kKeyTile;
+             ++key_offset) {
+          const int key_tile = key_offset / kKeyTile;
+          const int column = key_offset % kKeyTile;
+          shared.scores[
+              key_tile * kQueryTile * kKeyTile
+              + row * kKeyTile + column] = 0.0f;
+        }
+      }
+      __syncthreads();
+
+      for (int index = lane;
+           index < active_rows * kHeadDim;
+           index += kWarpSize) {
+        const int row = index / kHeadDim;
+        shared.running_output[index] *= shared.correction[row];
+      }
+      __syncthreads();
+
+#pragma unroll
+      for (int dim_tile = 0; dim_tile < kDimTiles; ++dim_tile) {
+        // CoreX's reference matmul reduces a 512-token K dimension
+        // hierarchically. Preserve that numerical shape with four fixed,
+        // contiguous 128-token partials and a deterministic binary merge.
+#pragma unroll
+        for (int split = 0; split < kPvReductionSplits; ++split) {
+          wmma::fragment<wmma::accumulator, 16, 16, 16, float>
+              output_fragment;
+          wmma::fill_fragment(output_fragment, 0.0f);
+          const int split_start = split * kKeyTilesPerPvSplit;
+          const int split_end =
+              min(group_key_tiles, split_start + kKeyTilesPerPvSplit);
+          for (int key_tile_in_group = split_start;
+               key_tile_in_group < split_end;
+               ++key_tile_in_group) {
+            const int local_key_start =
+                group_start + key_tile_in_group * kKeyTile;
+            const int logical_key_start = phase_base + local_key_start;
+            const float* score_tile =
+                shared.scores
+                + key_tile_in_group * kQueryTile * kKeyTile;
+            wmma::fragment<wmma::matrix_a, 16, 16, 16, float,
+                           wmma::row_major> probability_fragment;
+            wmma::load_matrix_sync(
+                probability_fragment, score_tile, 0);
+
+#pragma unroll
+            for (int quarter = 0; quarter < 4; ++quarter) {
+              const int row = lane / 16 + quarter * 4;
+              const int column = lane % 16;
+              const int logical_token = logical_key_start + row;
+              const int dim = dim_tile * kMmaK + column;
+              const float value =
+                  local_key_start + row < phase_tokens
+                      ? load_value(value_new, value_cache, block_table,
+                                   logical_token, context_len, dim)
+                      : 0.0f;
+              const int offset =
+                  wmma::CoordToOffset<
+                      32, wmma::layout_t::mem_row_major>(
+                      row, column);
+              shared.matrix_tile[offset] = value;
+            }
+            __syncthreads();
+
+            wmma::fragment<wmma::matrix_b, 16, 16, 16, float,
+                           wmma::row_major> value_fragment;
+            wmma::load_matrix_sync(
+                value_fragment, shared.matrix_tile, 0);
+            wmma::mma_sync(
+                output_fragment,
+                probability_fragment,
+                value_fragment,
+                output_fragment);
+            __syncthreads();
+          }
+
+          wmma::store_matrix_sync(
+              shared.partial_output
+                  + split * kQueryTile * kMmaK,
+              output_fragment,
+              0,
+              wmma::mem_row_major);
+          __syncthreads();
+        }
+
+#pragma unroll
+        for (int quarter = 0; quarter < 4; ++quarter) {
+          const int row = lane / 16 + quarter * 4;
+          const int column = lane % 16;
+          if (row < active_rows) {
+            const int output_index =
+                row * kHeadDim + dim_tile * kMmaK + column;
+            const int tile_index = row * kMmaK + column;
+            const int partial_stride = kQueryTile * kMmaK;
+            const float left = __fadd_rn(
+                shared.partial_output[tile_index],
+                shared.partial_output[partial_stride + tile_index]);
+            const float right = __fadd_rn(
+                shared.partial_output[2 * partial_stride + tile_index],
+                shared.partial_output[3 * partial_stride + tile_index]);
+            shared.running_output[output_index] = __fadd_rn(
+                shared.running_output[output_index],
+                __fadd_rn(left, right));
+          }
+        }
+        __syncthreads();
+      }
+    }
+  }
+
+  for (int index = lane; index < active_rows * kHeadDim;
+       index += kWarpSize) {
+    const int row = index / kHeadDim;
+    const int dim = index % kHeadDim;
+    const int query_index = query_start + row;
+    const int64_t destination =
+        (static_cast<int64_t>(query_index) * kNumQueryHeads
+          + query_head) * kHeadDim + dim;
+    output[destination] = __float2half_rn(
+        shared.running_output[index] / shared.running_sum[row]);
+  }
+  if (lane < active_rows) {
+    const int query_index = query_start + lane;
+    lse[static_cast<int64_t>(query_index) * kNumQueryHeads
+        + query_head] =
+        shared.running_max[lane] + logf(shared.running_sum[lane]);
+  }
+}
+
+}  // namespace
+
+std::vector<torch::Tensor> query_tiled_paged_prefill_forward(
+    const torch::Tensor& query, const torch::Tensor& key_new,
+    const torch::Tensor& value_new, const torch::Tensor& key_cache,
+    const torch::Tensor& value_cache, const torch::Tensor& block_table,
+    int64_t context_len_arg, double scale_arg) {
+  check_half_cuda_contiguous(query, "query");
+  check_half_cuda_contiguous(key_new, "key_new");
+  check_half_cuda_contiguous(value_new, "value_new");
+  check_half_cuda_contiguous(key_cache, "key_cache");
+  check_half_cuda_contiguous(value_cache, "value_cache");
+  TORCH_CHECK(block_table.is_cuda(),
+              "block_table must be a CUDA tensor");
+  TORCH_CHECK(block_table.scalar_type() == torch::kInt32,
+              "block_table must have dtype int32");
+  TORCH_CHECK(block_table.is_contiguous(),
+              "block_table must be contiguous");
+  TORCH_CHECK(block_table.dim() == 1,
+              "block_table must be one-dimensional");
+  TORCH_CHECK(query.dim() == 3 && query.size(1) == kNumQueryHeads
+                  && query.size(2) == kHeadDim,
+              "query must have shape (Q, 4, 256)");
+  TORCH_CHECK(key_new.dim() == 3 && key_new.size(1) == kNumKvHeads
+                  && key_new.size(2) == kHeadDim,
+              "key_new must have shape (Q, 1, 256)");
+  TORCH_CHECK(value_new.sizes() == key_new.sizes(),
+              "value_new must match key_new");
+  TORCH_CHECK(key_new.size(0) == query.size(0),
+              "query, key_new, and value_new lengths must match");
+  TORCH_CHECK(key_cache.dim() == 5
+                  && key_cache.size(1) == kNumKvHeads
+                  && key_cache.size(2) == kHeadDim / kKeyPack
+                  && key_cache.size(3) == kBlockSize
+                  && key_cache.size(4) == kKeyPack,
+              "key_cache must have shape (N, 1, 32, 16, 8)");
+  TORCH_CHECK(value_cache.dim() == 4
+                  && value_cache.size(1) == kNumKvHeads
+                  && value_cache.size(2) == kHeadDim
+                  && value_cache.size(3) == kBlockSize,
+              "value_cache must have shape (N, 1, 256, 16)");
+  TORCH_CHECK(key_cache.size(0) == value_cache.size(0),
+              "key/value cache block counts must match");
+  TORCH_CHECK(query.device() == key_new.device()
+                  && query.device() == value_new.device()
+                  && query.device() == key_cache.device()
+                  && query.device() == value_cache.device()
+                  && query.device() == block_table.device(),
+              "all tensors must use the same device");
+  TORCH_CHECK(context_len_arg >= 0
+                  && context_len_arg <= kMaxSequenceTokens,
+              "context_len is out of range");
+  TORCH_CHECK(context_len_arg % kBlockSize == 0,
+              "context_len must be block aligned");
+  const int query_len = static_cast<int>(query.size(0));
+  const int context_len = static_cast<int>(context_len_arg);
+  TORCH_CHECK(query_len > 0 && query_len <= kMaxQueryTokens,
+              "query length must be in [1, 8192]");
+  TORCH_CHECK(context_len + query_len <= kMaxSequenceTokens,
+              "context_len + query_len exceeds 262144");
+  const int required_blocks = context_len / kBlockSize;
+  TORCH_CHECK(block_table.numel() >= required_blocks,
+              "block_table is too short for context_len");
+  if (required_blocks > 0) {
+    auto active_blocks = block_table.narrow(0, 0, required_blocks);
+    const int minimum_block = active_blocks.min().item<int>();
+    const int maximum_block = active_blocks.max().item<int>();
+    TORCH_CHECK(minimum_block >= 0
+                    && maximum_block < key_cache.size(0),
+                "block_table contains an out-of-range physical block ID");
+  }
+  TORCH_CHECK(std::isfinite(scale_arg) && scale_arg > 0.0,
+              "scale must be finite and positive");
+
+  auto output = torch::empty_like(query);
+  auto lse = torch::empty(
+      {query_len, kNumQueryHeads},
+      query.options().dtype(torch::kFloat32));
+  const int query_tiles =
+      (query_len + kQueryTile - 1) / kQueryTile;
+  const int blocks = query_tiles * kNumQueryHeads;
+  auto stream = at::cuda::getCurrentCUDAStream();
+  query_tiled_paged_prefill_kernel<<<
+      blocks, kWarpSize, 0, stream>>>(
+      reinterpret_cast<const __half*>(query.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(key_new.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(value_new.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(key_cache.data_ptr<at::Half>()),
+      reinterpret_cast<const __half*>(value_cache.data_ptr<at::Half>()),
+      block_table.data_ptr<int>(),
+      reinterpret_cast<__half*>(output.data_ptr<at::Half>()),
+      lse.data_ptr<float>(), context_len, query_len,
+      static_cast<float>(scale_arg));
+  C10_CUDA_KERNEL_LAUNCH_CHECK();
+  return {output, lse};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) {
+  module.def("forward", &query_tiled_paged_prefill_forward,
+             "Fixed BI100 query-tiled paged-prefill forward");
+}

qwen3_6_scripts/gdn_prefix.py

@@ -0,0 +1,291 @@
+"""Shared GDN prefix-state cache contracts for the BI100 runtime."""
+
+from __future__ import annotations
+
+import os
+from collections import OrderedDict
+from dataclasses import dataclass
+from typing import Iterable, List, Optional, Sequence, Tuple
+
+
+GdnPrefixKey = Tuple[int, bytes]
+GdnCapturePoint = Tuple[int, GdnPrefixKey]
+
+_VALID_POLICIES = {"fine32", "admission64", "off"}
+GDN_KERNEL_CHUNK_TOKENS = 64
+GDN_DIRECT_MIN_REPLAY_TOKENS = 2
+
+_VALID_RESTORE_MODES = {"direct", "hybrid64", "chunk64", "aligned"}
+
+
+def _env_choice(name: str, default: str, choices: set[str]) -> str:
+    value = os.getenv(name, default).strip().lower()
+    if value not in choices:
+        allowed = ", ".join(sorted(choices))
+        raise RuntimeError(f"invalid {name}={value!r}; expected one of: {allowed}")
+    return value
+
+
+def gdn_cache_policy_from_env() -> str:
+    return _env_choice("BI100_GDN_CACHE_POLICY", "fine32", _VALID_POLICIES)
+
+
+def gdn_restore_mode_from_env() -> str:
+    return _env_choice(
+        "BI100_GDN_RESTORE_MODE", "direct", _VALID_RESTORE_MODES)
+
+
+def gdn_restore_alignment(restore_mode: str, block_size: int,
+                          scheduler_chunk_tokens: int) -> int:
+    """Return the content boundary required by a restore mode."""
+    if block_size <= 0:
+        raise ValueError("block_size must be positive")
+    if restore_mode == "direct":
+        return block_size
+    if restore_mode in {"hybrid64", "chunk64"}:
+        alignment = GDN_KERNEL_CHUNK_TOKENS
+    elif restore_mode == "aligned":
+        alignment = scheduler_chunk_tokens
+    else:
+        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
+    if alignment <= 0 or alignment % block_size != 0:
+        raise ValueError(
+            f"{restore_mode} GDN restore requires a positive alignment "
+            f"divisible by block_size={block_size}; got {alignment}")
+    return alignment
+
+
+def make_prefix_key(block_count: int, digest: bytes) -> GdnPrefixKey:
+    if block_count <= 0:
+        raise ValueError("GDN prefix key requires at least one complete block")
+    if not isinstance(digest, bytes) or len(digest) != 32:
+        raise ValueError("GDN prefix digest must be exactly 32 bytes")
+    return block_count, digest
+
+
+def keys_from_block_hashes(block_hashes: Sequence[bytes]) -> List[GdnPrefixKey]:
+    return [make_prefix_key(i + 1, digest)
+            for i, digest in enumerate(block_hashes)]
+
+
+def strict_prefix_block_count(token_count: int, block_size: int) -> int:
+    if block_size <= 0:
+        raise ValueError("block_size must be positive")
+    if token_count <= 1:
+        return 0
+    return (token_count - 1) // block_size
+
+
+def key_at_strict_boundary(block_hashes: Sequence[bytes], token_count: int,
+                           block_size: int) -> Optional[GdnPrefixKey]:
+    block_count = min(
+        len(block_hashes), strict_prefix_block_count(token_count, block_size))
+    if block_count <= 0:
+        return None
+    return make_prefix_key(block_count, block_hashes[block_count - 1])
+
+
+def final_capture_key(
+        block_hashes: Sequence[bytes], prompt_tokens: int, block_size: int,
+        restore_mode: str, replay_alignment: int) -> Optional[GdnPrefixKey]:
+    if restore_mode in {"direct", "hybrid64"}:
+        block_count = min(
+            len(block_hashes), strict_prefix_block_count(
+                prompt_tokens, block_size))
+        if (block_count > 0
+                and prompt_tokens - block_count * block_size
+                < GDN_DIRECT_MIN_REPLAY_TOKENS):
+            block_count -= 1
+        if block_count <= 0:
+            return None
+        return make_prefix_key(block_count, block_hashes[block_count - 1])
+    if restore_mode not in {"chunk64", "aligned"}:
+        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
+    if (replay_alignment <= 0 or replay_alignment % block_size != 0
+            or prompt_tokens <= 1):
+        return None
+    boundary_tokens = ((prompt_tokens - 1) // replay_alignment
+                       * replay_alignment)
+    block_count = min(len(block_hashes), boundary_tokens // block_size)
+    if block_count <= 0:
+        return None
+    return make_prefix_key(block_count, block_hashes[block_count - 1])
+
+
+def restore_key_is_eligible(
+        key: GdnPrefixKey, prompt_tokens: int, block_size: int,
+        restore_mode: str, replay_alignment: int,
+        direct_final_key: Optional[GdnPrefixKey] = None) -> bool:
+    """Return whether restoring ``key`` preserves the execution contract."""
+    make_prefix_key(*key)
+    if block_size <= 0:
+        raise ValueError("block_size must be positive")
+    boundary_tokens = key[0] * block_size
+    remaining_tokens = prompt_tokens - boundary_tokens
+    if remaining_tokens <= 0:
+        return False
+    if restore_mode == "direct":
+        return remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
+    if restore_mode == "hybrid64":
+        if direct_final_key is not None:
+            make_prefix_key(*direct_final_key)
+        return (remaining_tokens >= GDN_DIRECT_MIN_REPLAY_TOKENS
+                and replay_alignment > 0
+                and (boundary_tokens % replay_alignment == 0
+                     or key == direct_final_key))
+    if restore_mode not in {"chunk64", "aligned"}:
+        raise ValueError(f"unknown GDN restore mode: {restore_mode}")
+    return (replay_alignment > 0
+            and boundary_tokens % replay_alignment == 0)
+
+
+def capture_points_for_step(
+        targets: Iterable[GdnPrefixKey], physical_context_tokens: int,
+        logical_end_tokens: int, block_size: int) -> Tuple[GdnCapturePoint, ...]:
+    if physical_context_tokens < 0 or logical_end_tokens < 0:
+        raise ValueError("token positions must be non-negative")
+    if logical_end_tokens <= physical_context_tokens:
+        return ()
+    selected = {}
+    for key in targets:
+        make_prefix_key(*key)
+        boundary_tokens = key[0] * block_size
+        if physical_context_tokens < boundary_tokens <= logical_end_tokens:
+            selected[boundary_tokens - physical_context_tokens] = key
+    points = tuple(sorted(selected.items()))
+    if len(points) > 2:
+        raise ValueError("at most two GDN capture points are allowed per step")
+    return points
+
+
+def cap_prefill_end_at_capture_boundary(
+        logical_start_tokens: int, logical_end_tokens: int,
+        targets: Iterable[GdnPrefixKey], block_size: int) -> int:
+    """Stop a physical prefill step at its earliest pending capture boundary."""
+    if logical_start_tokens < 0 or logical_end_tokens < 0:
+        raise ValueError("token positions must be non-negative")
+    if logical_end_tokens < logical_start_tokens:
+        raise ValueError("logical end must not precede logical start")
+    if block_size <= 0:
+        raise ValueError("block_size must be positive")
+
+    capped_end = logical_end_tokens
+    for key in targets:
+        make_prefix_key(*key)
+        boundary_tokens = key[0] * block_size
+        if logical_start_tokens < boundary_tokens < capped_end:
+            capped_end = boundary_tokens
+    return capped_end
+
+
+def canonical_direct_segment_offsets(
+        block_hashes: Sequence[bytes], physical_context_tokens: int,
+        logical_end_tokens: int, block_size: int,
+        scheduler_chunk_tokens: int) -> Tuple[int, ...]:
+    """Reproduce cold fine32/direct segment boundaries after fast-forward."""
+    if physical_context_tokens < 0 or logical_end_tokens < 0:
+        raise ValueError("token positions must be non-negative")
+    if block_size <= 0 or scheduler_chunk_tokens <= 0:
+        raise ValueError("block and scheduler chunk sizes must be positive")
+    if scheduler_chunk_tokens % block_size != 0:
+        raise ValueError("scheduler chunk size must be divisible by block size")
+    if logical_end_tokens <= physical_context_tokens:
+        return ()
+
+    boundaries = set()
+    step_ends = list(range(scheduler_chunk_tokens, logical_end_tokens,
+                           scheduler_chunk_tokens))
+    for step_end in (*step_ends, logical_end_tokens):
+        key = final_capture_key(block_hashes, step_end, block_size,
+                                "direct", block_size)
+        if key is not None:
+            boundaries.add(key[0] * block_size)
+    boundaries.update(step_ends)
+    return tuple(
+        boundary - physical_context_tokens
+        for boundary in sorted(boundaries)
+        if physical_context_tokens < boundary < logical_end_tokens)
+
+
+@dataclass(frozen=True)
+class GdnCachePlan:
+    restore_key: Optional[GdnPrefixKey] = None
+    capture_points: Tuple[GdnCapturePoint, ...] = ()
+    evict_keys: Tuple[GdnPrefixKey, ...] = ()
+
+
+class GdnPrefixStatePolicy:
+    """Scheduler-owned state index with deterministic worker actions."""
+
+    def __init__(self, policy: str) -> None:
+        if policy not in _VALID_POLICIES:
+            raise ValueError(f"unknown GDN cache policy: {policy}")
+        self.policy = policy
+        self.capacity = {"fine32": 32, "admission64": 64, "off": 0}[policy]
+        self._resident: OrderedDict[GdnPrefixKey, None] = OrderedDict()
+
+    def __len__(self) -> int:
+        return len(self._resident)
+
+    def resident_keys(self) -> Tuple[GdnPrefixKey, ...]:
+        return tuple(self._resident)
+
+    def contains(self, key: GdnPrefixKey) -> bool:
+        return key in self._resident
+
+    def should_capture_final(self, key: GdnPrefixKey) -> bool:
+        """Return whether a final state must be materialized on this request."""
+        make_prefix_key(*key)
+        if self.policy == "off":
+            return False
+        if self.policy == "admission64":
+            return key not in self._resident
+        return True
+
+    def select_restore(
+            self, live_prefix_keys: Sequence[GdnPrefixKey],
+            max_blocks: int) -> Optional[GdnPrefixKey]:
+        if self.capacity == 0 or max_blocks <= 0:
+            return None
+        best = None
+        for key in live_prefix_keys[:max_blocks]:
+            if key in self._resident:
+                best = key
+        if best is not None:
+            self._resident.move_to_end(best)
+        return best
+
+    def repeated_branch_candidate(
+            self, live_prefix_keys: Sequence[GdnPrefixKey],
+            max_blocks: int) -> Optional[GdnPrefixKey]:
+        """Return a repeated raw-KV branch that lacks recurrent state.
+
+        A live KV hit proves that the content occurred in an earlier request;
+        the current request is therefore the second or later occurrence.
+        """
+        if (self.policy != "admission64" or max_blocks <= 0
+                or not live_prefix_keys):
+            return None
+        candidate = live_prefix_keys[min(len(live_prefix_keys), max_blocks) - 1]
+        if candidate in self._resident:
+            return None
+        return candidate
+
+    def admit(self, keys: Iterable[GdnPrefixKey]) -> Tuple[GdnPrefixKey, ...]:
+        evicted: List[GdnPrefixKey] = []
+        if self.capacity == 0:
+            return ()
+        for key in keys:
+            make_prefix_key(*key)
+            if key in self._resident:
+                self._resident.move_to_end(key)
+            else:
+                self._resident[key] = None
+            while len(self._resident) > self.capacity:
+                evicted_key, _ = self._resident.popitem(last=False)
+                evicted.append(evicted_key)
+        return tuple(evicted)
+
+    def forget(self, keys: Iterable[GdnPrefixKey]) -> None:
+        for key in keys:
+            self._resident.pop(key, None)

qwen3_6_scripts/install_prebuilt_corex.sh

@@ -0,0 +1,62 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+VLLM_ROOT=${1:?usage: install_prebuilt_corex.sh VLLM_ROOT}
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+BUNDLE_DIR=${SCRIPT_DIR}/prebuilt/corex-3.2.3-ivcore10
+MANIFEST=${BUNDLE_DIR}/SHA256SUMS
+
+[[ -d "$VLLM_ROOT" ]] || {
+    printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
+    exit 2
+}
+[[ -f "$MANIFEST" ]] || {
+    printf 'prebuilt CoreX manifest is missing: %s\n' "$MANIFEST" >&2
+    exit 2
+}
+
+mapfile -t artifacts < <(awk '{print $2}' "$MANIFEST")
+[[ "${#artifacts[@]}" -eq 12 ]] || {
+    printf 'expected 12 prebuilt CoreX artifacts, found %s\n' \
+        "${#artifacts[@]}" >&2
+    exit 2
+}
+
+for artifact in "${artifacts[@]}"; do
+    [[ "$artifact" == corex_*.so && "$artifact" != */* ]] || {
+        printf 'invalid prebuilt artifact name: %s\n' "$artifact" >&2
+        exit 2
+    }
+done
+
+(
+    cd "$BUNDLE_DIR"
+    sha256sum --strict --check SHA256SUMS
+)
+
+for artifact in "${artifacts[@]}"; do
+    install -m 0755 "$BUNDLE_DIR/$artifact" "$VLLM_ROOT/$artifact"
+done
+
+python3 - "$VLLM_ROOT" "${artifacts[@]}" <<'PY'
+import pathlib
+import struct
+import sys
+
+root = pathlib.Path(sys.argv[1])
+for name in sys.argv[2:]:
+    path = root / name
+    if not path.is_file() or path.stat().st_size == 0:
+        raise SystemExit(f"installed CoreX extension is empty: {path}")
+    header = path.read_bytes()[:20]
+    if len(header) < 20 or header[:4] != b"\x7fELF":
+        raise SystemExit(f"installed CoreX extension is not ELF: {path}")
+    if header[4:6] != b"\x02\x01":
+        raise SystemExit(
+            f"installed CoreX extension is not 64-bit little-endian ELF: {path}")
+    machine = struct.unpack_from("<H", header, 18)[0]
+    if machine != 62:
+        raise SystemExit(
+            f"installed CoreX extension is not x86-64 ELF: {path} machine={machine}")
+    print(f"[ok] installed prebuilt CoreX extension {path}")
+PY

qwen3_6_scripts/logits_processor.py

@@ -1,158 +0,0 @@
-"""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

@@ -70,15 +70,10 @@ class MambaCacheManager:
         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)
+            cache_t[:, [to_index,from_index]] = \
+             cache_t[:, [from_index,to_index]]
 
     def _copy_mamba_cache(self, from_index: int, to_index: int):
         assert len(self.mamba_cache) > 0

qwen3_6_scripts/paged_attention_v2_pytorch.py

@@ -1,343 +0,0 @@
-"""
-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/patch_block_major_cache_engine.py

@@ -0,0 +1,91 @@
+from patch_utils import package_root, replace_once
+
+
+CACHE_ENGINE = package_root("vllm") / "worker" / "cache_engine.py"
+
+IMPORT_ANCHOR = """\
+from vllm.logger import init_logger
+"""
+
+IMPORT_REPLACEMENT = """\
+from vllm.block_major_kv_cache import (
+    BlockMajorCpuKVCache,
+    block_major_cpu_kv_enabled,
+)
+from vllm.logger import init_logger
+"""
+
+ALLOCATION_ANCHOR = """\
+        self.gpu_cache = self._allocate_kv_cache(
+            self.num_gpu_blocks, self.device_config.device_type)
+        self.cpu_cache = self._allocate_kv_cache(self.num_cpu_blocks, "cpu")
+"""
+
+ALLOCATION_REPLACEMENT = """\
+        self.gpu_cache = self._allocate_kv_cache(
+            self.num_gpu_blocks, self.device_config.device_type)
+        self._bi100_block_major_cpu_kv = None
+        if block_major_cpu_kv_enabled():
+            self._bi100_block_major_cpu_kv = BlockMajorCpuKVCache(
+                self.gpu_cache,
+                self.num_cpu_blocks,
+                pin_memory=is_pin_memory_available(),
+            )
+            self.cpu_cache = self._bi100_block_major_cpu_kv.layer_views
+        else:
+            self.cpu_cache = self._allocate_kv_cache(
+                self.num_cpu_blocks, "cpu")
+"""
+
+SWAP_ANCHOR = """\
+    def swap_in(self, src_to_dst: torch.Tensor) -> None:
+        for i in range(self.num_attention_layers):
+            self.attn_backend.swap_blocks(self.cpu_cache[i], self.gpu_cache[i],
+                                          src_to_dst)
+
+    def swap_out(self, src_to_dst: torch.Tensor) -> None:
+        for i in range(self.num_attention_layers):
+            self.attn_backend.swap_blocks(self.gpu_cache[i], self.cpu_cache[i],
+                                          src_to_dst)
+"""
+
+SWAP_REPLACEMENT = """\
+    def swap_in(self, src_to_dst: torch.Tensor) -> None:
+        if self._bi100_block_major_cpu_kv is not None:
+            self._bi100_block_major_cpu_kv.swap_in(src_to_dst)
+            return
+        for i in range(self.num_attention_layers):
+            self.attn_backend.swap_blocks(self.cpu_cache[i], self.gpu_cache[i],
+                                          src_to_dst)
+
+    def swap_out(self, src_to_dst: torch.Tensor) -> None:
+        if self._bi100_block_major_cpu_kv is not None:
+            self._bi100_block_major_cpu_kv.swap_out(src_to_dst)
+            return
+        for i in range(self.num_attention_layers):
+            self.attn_backend.swap_blocks(self.gpu_cache[i], self.cpu_cache[i],
+                                          src_to_dst)
+"""
+
+
+replace_once(
+    CACHE_ENGINE,
+    IMPORT_ANCHOR,
+    IMPORT_REPLACEMENT,
+    required=True,
+    already_contains="from vllm.block_major_kv_cache import",
+)
+replace_once(
+    CACHE_ENGINE,
+    ALLOCATION_ANCHOR,
+    ALLOCATION_REPLACEMENT,
+    required=True,
+    already_contains="self._bi100_block_major_cpu_kv = None",
+)
+replace_once(
+    CACHE_ENGINE,
+    SWAP_ANCHOR,
+    SWAP_REPLACEMENT,
+    required=True,
+    already_contains="self._bi100_block_major_cpu_kv.swap_in",
+)

qwen3_6_scripts/patch_block_major_worker_capacity.py

@@ -0,0 +1,46 @@
+from patch_utils import package_root, replace_once
+
+
+WORKER = package_root("vllm") / "worker" / "worker.py"
+
+IMPORT_ANCHOR = """\
+from vllm.logger import init_logger
+"""
+
+IMPORT_REPLACEMENT = """\
+from vllm.block_major_kv_cache import reserve_block_major_gpu_blocks
+from vllm.logger import init_logger
+"""
+
+CAPACITY_ANCHOR = """\
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+"""
+
+CAPACITY_REPLACEMENT = """\
+        num_gpu_blocks = reserve_block_major_gpu_blocks(
+            num_gpu_blocks, cache_block_size)
+        num_gpu_blocks = max(num_gpu_blocks, 0)
+        num_cpu_blocks = max(num_cpu_blocks, 0)
+"""
+
+
+replace_once(
+    WORKER,
+    IMPORT_ANCHOR,
+    IMPORT_REPLACEMENT,
+    required=True,
+    already_contains=(
+        "from vllm.block_major_kv_cache import "
+        "reserve_block_major_gpu_blocks"
+    ),
+)
+replace_once(
+    WORKER,
+    CAPACITY_ANCHOR,
+    CAPACITY_REPLACEMENT,
+    required=True,
+    already_contains=(
+        "num_gpu_blocks = reserve_block_major_gpu_blocks("
+    ),
+)

qwen3_6_scripts/patch_block_manager_cache_trace.py

@@ -0,0 +1,210 @@
+"""Install the optional BI100 prefix-cache diagnostic trace."""
+from patch_utils import package_root, replace_once, replace_one_of
+
+VLLM_ROOT = package_root("vllm")
+TARGET = VLLM_ROOT / "core" / "block_manager_v2.py"
+OUTPUTS_TARGET = VLLM_ROOT / "outputs.py"
+
+HELPER = '''
+    def _bi100_capture_cache_trace(self, seq_group, seq, block_table) -> None:
+        if os.getenv("BI100_CACHE_TRACE", "0") != "1":
+            return
+
+        session = getattr(self, "_bi100_trace_session", None)
+        if session is None:
+            session = hashlib.sha256(os.urandom(16)).hexdigest()[:16]
+            self._bi100_trace_session = session
+
+        self._bi100_trace_ordinal = getattr(self, "_bi100_trace_ordinal", 0) + 1
+        request_id_sha256 = hashlib.sha256(
+            str(seq_group.request_id).encode("utf-8")).hexdigest()[:16]
+
+        prompt_tokens = len(seq.get_token_ids())
+        requests = getattr(self, "_bi100_trace_requests", None)
+        if requests is None:
+            requests = {}
+            self._bi100_trace_requests = requests
+
+        requests[seq.seq_id] = {
+            "version": 4,
+            "trace_session_sha256": session,
+            "ordinal": self._bi100_trace_ordinal,
+            "request_id_sha256": request_id_sha256,
+            "prompt_tokens": prompt_tokens,
+            "prompt_allocated_blocks": (
+                (prompt_tokens + self.block_size - 1) // self.block_size
+            ),
+            "block_size": self.block_size,
+            "capacity_blocks": self.num_total_gpu_blocks,
+        }
+        setattr(seq_group, "_bi100_cache_trace_seq_id", seq.seq_id)
+        setattr(seq_group, "_bi100_cache_trace_emit",
+                self._bi100_emit_cache_trace)
+
+    def _bi100_update_cache_trace(
+            self, seq, raw_kv_hit_blocks, restore_key, capture_actions,
+            evict_keys, policy) -> None:
+        if os.getenv("BI100_CACHE_TRACE", "0") != "1":
+            return
+        requests = getattr(self, "_bi100_trace_requests", None)
+        if not requests or seq.seq_id not in requests:
+            return
+        record = requests[seq.seq_id]
+        record["gdn_policy"] = policy
+        if "initial_raw_kv_contiguous_hit_blocks" not in record:
+            record["initial_raw_kv_contiguous_hit_blocks"] = max(
+                0, int(raw_kv_hit_blocks))
+            record["gdn_restore_digest_base64"] = (
+                base64.b64encode(restore_key[1]).decode("ascii")
+                if restore_key is not None else None)
+        record["raw_kv_contiguous_hit_blocks"] = max(
+            int(raw_kv_hit_blocks),
+            int(record.get("raw_kv_contiguous_hit_blocks", 0)))
+        effective_blocks = int(restore_key[0]) if restore_key is not None else 0
+        record["effective_gdn_hit_blocks"] = max(
+            effective_blocks, int(record.get("effective_gdn_hit_blocks", 0)))
+
+        admissions = record.setdefault("gdn_admissions", [])
+        for key, reason in capture_actions:
+            admissions.append({
+                "block_count": int(key[0]),
+                "digest_base64": base64.b64encode(key[1]).decode("ascii"),
+                "reason": str(reason),
+            })
+        evictions = record.setdefault("gdn_evictions", [])
+        for key in evict_keys:
+            evictions.append({
+                "block_count": int(key[0]),
+                "digest_base64": base64.b64encode(key[1]).decode("ascii"),
+                "reason": "capacity_lru",
+            })
+
+    def _bi100_finalize_cache_trace(self, seq, block_table) -> None:
+        if os.getenv("BI100_CACHE_TRACE", "0") != "1":
+            return
+
+        requests = getattr(self, "_bi100_trace_requests", None)
+        if not requests:
+            return
+
+        record = requests.get(seq.seq_id)
+        if record is None:
+            return
+
+        total_tokens = len(seq.get_token_ids())
+        block_hashes = block_table.get_content_hashes()
+        for block_hash in block_hashes:
+            if not isinstance(block_hash, bytes) or len(block_hash) != 32:
+                raise RuntimeError(
+                    "BI100 cache trace requires 32-byte content hashes")
+        full_blocks = len(block_hashes)
+        record.update({
+            "total_tokens": total_tokens,
+            "allocated_blocks": (
+                (total_tokens + self.block_size - 1) // self.block_size
+            ),
+            "full_blocks": full_blocks,
+            "hash_encoding": "sha256_base64",
+            "block_hashes": base64.b64encode(b"".join(block_hashes)).decode("ascii"),
+            "_finalized": True,
+        })
+        generated_tokens = max(0, total_tokens - record["prompt_tokens"])
+        record["generated_tokens"] = generated_tokens
+
+    def _bi100_emit_cache_trace(self, seq_group) -> None:
+        if os.getenv("BI100_CACHE_TRACE", "0") != "1":
+            return
+        seq_id = getattr(seq_group, "_bi100_cache_trace_seq_id", None)
+        requests = getattr(self, "_bi100_trace_requests", None)
+        if seq_id is None or not requests:
+            return
+        record = requests.pop(seq_id, None)
+        if record is None:
+            return
+        if record.pop("_finalized", False) is not True:
+            raise RuntimeError(
+                "BI100 cache trace emitted before block finalization")
+
+        metrics = getattr(seq_group, "metrics", None)
+        arrival = getattr(metrics, "arrival_time", None)
+        first_token = getattr(metrics, "first_token_time", None)
+        finished = getattr(metrics, "finished_time", None)
+        queue = getattr(metrics, "time_in_queue", None)
+        cached = getattr(metrics, "num_cached_tokens", None)
+        if any(value is None for value in (
+                arrival, first_token, finished, queue)):
+            raise RuntimeError(
+                "BI100 cache trace requires finalized request metrics")
+        record["ttft_s"] = max(0.0, float(first_token - arrival))
+        record["request_latency_s"] = max(
+            0.0, float(finished - arrival))
+        record["time_in_queue_s"] = max(0.0, float(queue))
+        record["observed_effective_cached_tokens"] = max(
+            0, int(cached or 0))
+        ttft_s = record["ttft_s"]
+        if ttft_s > 0:
+            record["observed_input_tps"] = record["prompt_tokens"] / ttft_s
+        generated_tokens = record["generated_tokens"]
+        if generated_tokens > 1:
+            decode_s = finished - first_token
+            if decode_s > 0:
+                record["observed_output_tps"] = (
+                    (generated_tokens - 1) / decode_s)
+        print("[BI100_CACHE_TRACE] " + json.dumps(record, separators=(",", ":"),
+                                             sort_keys=True), flush=True)
+'''
+
+
+def main():
+    replace_once(TARGET, "from collections.abc import Mapping\n",
+                 "from collections.abc import Mapping\nimport base64\nimport json\nimport os\n",
+                 required=True, already_contains="import base64\n")
+    replace_once(TARGET, "class BlockSpaceManagerV2(BlockSpaceManager):\n",
+                 "class BlockSpaceManagerV2(BlockSpaceManager):\n" + HELPER,
+                 required=True, already_contains="def _bi100_capture_cache_trace(")
+    replace_once(TARGET,
+                 "        self.block_tables[seq.seq_id] = block_table\n\n        # Track seq",
+                 "        self.block_tables[seq.seq_id] = block_table\n        self._bi100_capture_cache_trace(\n            seq_group, seq, block_table)\n\n        # Track seq",
+                 required=True,
+                 already_contains="self.block_tables[seq.seq_id] = block_table\n"
+                                "        self._bi100_capture_cache_trace(")
+    replacements = []
+    for table_key in ("seq_id", "seq.seq_id"):
+        prefix = (
+            "        self._last_access_blocks_tracker."
+            "update_seq_blocks_last_access(\n"
+            f"            seq_id, self.block_tables[{table_key}]."
+            "physical_block_ids)\n")
+        replacements.append((
+            prefix + "\n        # Untrack seq",
+            prefix + "        self._bi100_finalize_cache_trace(\n"
+            f"            seq, self.block_tables[{table_key}])\n\n"
+            "        # Untrack seq",
+        ))
+    replace_one_of(
+        TARGET,
+        replacements,
+        required=True,
+        already_contains="        self._bi100_finalize_cache_trace(\n"
+                         "            seq, self.block_tables[")
+    replace_once(
+        OUTPUTS_TARGET,
+        "        seq_group.set_finished_time(finished_time)\n\n"
+        "        init_args = (seq_group.request_id, prompt, prompt_token_ids,\n",
+        "        seq_group.set_finished_time(finished_time)\n"
+        "        if finished_time is not None:\n"
+        "            cache_trace_emit = getattr(\n"
+        "                seq_group, \"_bi100_cache_trace_emit\", None)\n"
+        "            if callable(cache_trace_emit):\n"
+        "                cache_trace_emit(seq_group)\n"
+        "                delattr(seq_group, \"_bi100_cache_trace_emit\")\n"
+        "                delattr(seq_group, \"_bi100_cache_trace_seq_id\")\n\n"
+        "        init_args = (seq_group.request_id, prompt, prompt_token_ids,\n",
+        required=True,
+        already_contains="if finished_time is not None:\n"
+                         "            cache_trace_emit = getattr(\n",
+    )
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_corex_swap_blocks.py

@@ -0,0 +1,65 @@
+from patch_utils import package_root, replace_once
+
+
+CUSTOM_OPS = package_root("vllm") / "_custom_ops.py"
+
+CLEAN_BLOCK = """\
+def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
+                block_mapping: torch.Tensor) -> None:
+    ixf_F.swap_blocks(src, dst, block_mapping)
+"""
+
+COMPATIBLE_BLOCK = """\
+def swap_blocks(src: torch.Tensor, dst: torch.Tensor,
+                block_mapping: torch.Tensor) -> None:
+    # BI100 CoreX 3.2.3 exposes vllm_swap_blocks, while this vLLM build calls
+    # the newer swap_blocks name. Normalize the worker's CPU int64 [N, 2]
+    # tensor only for the legacy public API and fail fast on malformed maps.
+    native_swap_blocks = getattr(ixf_F, "swap_blocks", None)
+    if native_swap_blocks is not None:
+        native_swap_blocks(src, dst, block_mapping)
+        return
+
+    vendor_swap_blocks = getattr(ixf_F, "vllm_swap_blocks", None)
+    if vendor_swap_blocks is None:
+        raise RuntimeError(
+            "ixformer exposes neither swap_blocks nor vllm_swap_blocks")
+
+    if isinstance(block_mapping, torch.Tensor):
+        if block_mapping.device.type != "cpu":
+            raise ValueError("swap block mapping must be a CPU tensor")
+        if block_mapping.dtype != torch.int64:
+            raise ValueError("swap block mapping must use torch.int64")
+        if block_mapping.dim() != 2 or block_mapping.shape[1] != 2:
+            raise ValueError("swap block mapping must have shape [N, 2]")
+        pairs = block_mapping.tolist()
+    elif isinstance(block_mapping, dict):
+        pairs = list(block_mapping.items())
+    else:
+        raise TypeError("swap block mapping must be a tensor or dict")
+
+    normalized_mapping = {}
+    destinations = set()
+    for source, destination in pairs:
+        source = int(source)
+        destination = int(destination)
+        if source < 0 or destination < 0:
+            raise ValueError("swap block indices must be non-negative")
+        if source in normalized_mapping:
+            raise ValueError(f"duplicate swap source block: {source}")
+        if destination in destinations:
+            raise ValueError(
+                f"duplicate swap destination block: {destination}")
+        normalized_mapping[source] = destination
+        destinations.add(destination)
+    vendor_swap_blocks(src, dst, normalized_mapping)
+"""
+
+
+replace_once(
+    CUSTOM_OPS,
+    CLEAN_BLOCK,
+    COMPATIBLE_BLOCK,
+    required=True,
+    already_contains="BI100 CoreX 3.2.3 exposes vllm_swap_blocks",
+)

qwen3_6_scripts/patch_executor_startup_debug.py

@@ -0,0 +1,61 @@
+from patch_utils import package_root, replace_once
+
+VLLM_ROOT = package_root("vllm")
+
+MULTIPROC_GPU_EXECUTOR = VLLM_ROOT / "executor" / "multiproc_gpu_executor.py"
+MULTIPROC_WORKER_UTILS = VLLM_ROOT / "executor" / "multiproc_worker_utils.py"
+
+
+def ensure_import_os(path):
+    text = path.read_text()
+    if "import os\n" in text:
+        print(f"[skip] import os already present: {path}")
+        return
+    for anchor in ("import time\n", "import signal\n", "import sys\n"):
+        if anchor in text:
+            replace_once(
+                path,
+                anchor,
+                anchor + "import os\n",
+                required=True,
+                already_contains="import os\n",
+            )
+            return
+    raise RuntimeError(f"no import anchor found for os in {path}")
+
+
+ensure_import_os(MULTIPROC_GPU_EXECUTOR)
+ensure_import_os(MULTIPROC_WORKER_UTILS)
+
+
+replace_once(
+    MULTIPROC_GPU_EXECUTOR,
+    """logger = init_logger(__name__)\n""",
+    """logger = init_logger(__name__)\n\n\ndef _bi100_startup_debug(message: str, *args) -> None:\n    if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n        logger.info(\"[BI100 startup] \" + message, *args)\n""",
+    required=True,
+    already_contains="def _bi100_startup_debug(",
+)
+
+replace_once(
+    MULTIPROC_GPU_EXECUTOR,
+    """        self.driver_worker = self._create_worker(\n            distributed_init_method=distributed_init_method)\n        self._run_workers(\"init_device\")\n        self._run_workers(\"load_model\",\n                          max_concurrent_workers=self.parallel_config.\n                          max_parallel_loading_workers)\n""",
+    """        _bi100_startup_debug(\"creating driver worker\")\n        self.driver_worker = self._create_worker(\n            distributed_init_method=distributed_init_method)\n        _bi100_startup_debug(\"created driver worker\")\n        _bi100_startup_debug(\"starting init_device\")\n        self._run_workers(\"init_device\")\n        _bi100_startup_debug(\"finished init_device\")\n        _bi100_startup_debug(\"starting load_model\")\n        self._run_workers(\"load_model\",\n                          max_concurrent_workers=self.parallel_config.\n                          max_parallel_loading_workers)\n        _bi100_startup_debug(\"finished load_model\")\n""",
+    required=True,
+    already_contains='_bi100_startup_debug("starting init_device")',
+)
+
+replace_once(
+    MULTIPROC_GPU_EXECUTOR,
+    """        # Start all remote workers first.\n        worker_outputs = [\n            worker.execute_method(method, *args, **kwargs)\n            for worker in self.workers\n        ]\n\n        driver_worker_method = getattr(self.driver_worker, method)\n        driver_worker_output = driver_worker_method(*args, **kwargs)\n\n        # Get the results of the workers.\n        return [driver_worker_output\n                ] + [output.get() for output in worker_outputs]\n""",
+    """        _bi100_startup_debug(\"enqueue remote method=%s workers=%d\", method,\n                              len(self.workers))\n        # Start all remote workers first.\n        worker_outputs = [\n            worker.execute_method(method, *args, **kwargs)\n            for worker in self.workers\n        ]\n        _bi100_startup_debug(\"remote enqueued method=%s\", method)\n\n        driver_worker_method = getattr(self.driver_worker, method)\n        _bi100_startup_debug(\"driver start method=%s\", method)\n        driver_worker_output = driver_worker_method(*args, **kwargs)\n        _bi100_startup_debug(\"driver done method=%s\", method)\n\n        # Get the results of the workers.\n        _bi100_startup_debug(\"waiting remote results method=%s\", method)\n        remote_outputs = [output.get() for output in worker_outputs]\n        _bi100_startup_debug(\"remote done method=%s\", method)\n        return [driver_worker_output] + remote_outputs\n""",
+    required=True,
+    already_contains='_bi100_startup_debug("enqueue remote method=%s workers=%d"',
+)
+
+replace_once(
+    MULTIPROC_WORKER_UTILS,
+    """            task_id, method, args, kwargs = items\n            try:\n                executor = getattr(worker, method)\n                output = executor(*args, **kwargs)\n            except SystemExit:\n""",
+    """            task_id, method, args, kwargs = items\n            if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n                logger.info(\"[BI100 worker] start method=%s\", method)\n            try:\n                executor = getattr(worker, method)\n                output = executor(*args, **kwargs)\n                if os.getenv(\"BI100_EXECUTOR_STARTUP_DEBUG\") == \"1\":\n                    logger.info(\"[BI100 worker] done method=%s\", method)\n            except SystemExit:\n""",
+    required=True,
+    already_contains='logger.info("[BI100 worker] start method=%s", method)',
+)

qwen3_6_scripts/patch_model_runner.py

@@ -0,0 +1,408 @@
+"""Patch vLLM 0.6.3 prefix-cache and MRoPE chunk alignment bugs."""
+
+from __future__ import annotations
+
+import pathlib
+
+from patch_utils import package_root, replace_once
+
+
+HELPER_ANCHOR = """\
+logger = init_logger(__name__)
+
+LORA_WARMUP_RANK = 8"""
+
+HELPER_REPLACEMENT = """\
+logger = init_logger(__name__)
+
+
+def _slice_mrope_positions(positions, start, stop, expected_len):
+    if positions is None or len(positions) != 3:
+        raise RuntimeError("MRoPE positions must contain three axes")
+    sliced = [axis[start:stop] for axis in positions]
+    lengths = [len(axis) for axis in sliced]
+    if lengths != [expected_len] * 3:
+        raise RuntimeError(
+            "MRoPE/input token length mismatch after chunk alignment: "
+            f"positions={lengths}, input_tokens={expected_len}, "
+            f"slice=({start}, {stop})")
+    return sliced
+
+
+LORA_WARMUP_RANK = 8"""
+
+PREFIX_PAST_ANCHOR = """\
+        if prefix_cache_len <= context_len:
+            # We already passed the cache hit region,
+            # so do normal computation.
+            pass"""
+
+PREFIX_PAST_REPLACEMENT = """\
+        if prefix_cache_len <= context_len:
+            # We already passed the cache hit region,
+            # so do normal computation.
+            # Must clear prefix_cache_hit so _add_seq_group uses the full
+            # block_tables (prefix + previous-chunk blocks) instead of only
+            # computed_block_nums (prefix only).  Without this, block_tables
+            # passed to _forward_prefix_pytorch is too narrow for context_len,
+            # causing an empty blk_ids slice and a zero-dim amax() crash.
+            inter_data.prefix_cache_hit = False"""
+
+PARTIAL_HIT_ANCHOR = """\
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
+                seq_idx][uncomputed_start:]
+            context_len = prefix_cache_len
+
+            inter_data.context_lens[seq_idx] = context_len
+            inter_data.query_lens[
+                seq_idx] = inter_data.seq_lens[seq_idx] - context_len"""
+
+PARTIAL_HIT_REPLACEMENT = """\
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
+                seq_idx][uncomputed_start:]
+            context_len = prefix_cache_len
+
+            inter_data.context_lens[seq_idx] = context_len
+            inter_data.query_lens[
+                seq_idx] = inter_data.seq_lens[seq_idx] - context_len
+            if inter_data.mrope_input_positions is not None:
+                positions = inter_data.mrope_input_positions[seq_idx]
+                if positions is not None:
+                    inter_data.mrope_input_positions[seq_idx] = \\
+                        _slice_mrope_positions(
+                            positions, uncomputed_start, None,
+                            inter_data.query_lens[seq_idx])"""
+
+FULL_HIT_ANCHOR = """\
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
+                seq_idx][-1:]
+            inter_data.query_lens[seq_idx] = 1
+            inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1"""
+
+FULL_HIT_REPLACEMENT = """\
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
+                seq_idx][-1:]
+            inter_data.query_lens[seq_idx] = 1
+            inter_data.context_lens[seq_idx] = inter_data.seq_lens[seq_idx] - 1
+            if inter_data.mrope_input_positions is not None:
+                positions = inter_data.mrope_input_positions[seq_idx]
+                if positions is not None:
+                    inter_data.mrope_input_positions[seq_idx] = \\
+                        _slice_mrope_positions(positions, -1, None, 1)"""
+
+MULTIMODAL_MROPE_ANCHOR = """\
+                mrope_input_positions, mrope_position_delta = \\
+                    MRotaryEmbedding.get_input_positions(
+                        token_ids,
+                        image_grid_thw=image_grid_thw,
+                        video_grid_thw=video_grid_thw,
+                        image_token_id=hf_config.image_token_id,
+                        video_token_id=hf_config.video_token_id,
+                        vision_start_token_id=hf_config.vision_start_token_id,
+                        vision_end_token_id=hf_config.vision_end_token_id,
+                        spatial_merge_size=hf_config.vision_config.
+                        spatial_merge_size,
+                        context_len=inter_data.context_lens[seq_idx],
+                    )
+
+                seq_data.mrope_position_delta = mrope_position_delta
+                inter_data.mrope_input_positions[
+                    seq_idx] = mrope_input_positions"""
+
+MULTIMODAL_MROPE_REPLACEMENT = """\
+                # vLLM 0.6.3 returns positions through the end of token_ids,
+                # while chunked prefill sends only [context_len:seq_len].
+                # Compute the full MRoPE map once so the delta remains tied to
+                # the complete request, then select exactly the physical query.
+                mrope_input_positions, mrope_position_delta = \\
+                    MRotaryEmbedding.get_input_positions(
+                        token_ids,
+                        image_grid_thw=image_grid_thw,
+                        video_grid_thw=video_grid_thw,
+                        image_token_id=hf_config.image_token_id,
+                        video_token_id=hf_config.video_token_id,
+                        vision_start_token_id=hf_config.vision_start_token_id,
+                        vision_end_token_id=hf_config.vision_end_token_id,
+                        spatial_merge_size=hf_config.vision_config.
+                        spatial_merge_size,
+                        context_len=0,
+                    )
+                mrope_input_positions = _slice_mrope_positions(
+                    mrope_input_positions,
+                    inter_data.context_lens[seq_idx],
+                    inter_data.seq_lens[seq_idx],
+                    len(inter_data.input_tokens[seq_idx]))
+
+                seq_data.mrope_position_delta = mrope_position_delta
+                inter_data.mrope_input_positions[
+                    seq_idx] = mrope_input_positions"""
+
+MODEL_INPUT_FIELDS_ANCHOR = """\
+    multi_modal_kwargs: Optional[BatchedTensorInputs] = None
+    request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
+
+MODEL_INPUT_FIELDS_REPLACEMENT = """\
+    multi_modal_kwargs: Optional[BatchedTensorInputs] = None
+    # BI100 scheduler-owned GDN prefix-cache actions. These plain Python
+    # objects are included in the multiprocess model-input broadcast.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
+    request_ids_to_seq_ids: Optional[Dict[str, List[int]]] = None"""
+
+BASE_BROADCAST_ANCHOR = """\
+            \"multi_modal_kwargs\": self.multi_modal_kwargs,
+            \"prompt_adapter_mapping\": self.prompt_adapter_mapping,
+            \"prompt_adapter_requests\": self.prompt_adapter_requests,
+            \"virtual_engine\": self.virtual_engine,
+            \"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
+            \"finished_requests_ids\": self.finished_requests_ids,
+        }
+        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
+        return tensor_dict
+
+    @classmethod"""
+
+BASE_BROADCAST_REPLACEMENT = """\
+            \"multi_modal_kwargs\": self.multi_modal_kwargs,
+            \"gdn_restore_key\": self.gdn_restore_key,
+            \"gdn_capture_points\": self.gdn_capture_points,
+            \"gdn_evict_keys\": self.gdn_evict_keys,
+            \"gdn_segment_offsets\": self.gdn_segment_offsets,
+            \"prompt_adapter_mapping\": self.prompt_adapter_mapping,
+            \"prompt_adapter_requests\": self.prompt_adapter_requests,
+            \"virtual_engine\": self.virtual_engine,
+            \"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
+            \"finished_requests_ids\": self.finished_requests_ids,
+        }
+        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
+        return tensor_dict
+
+    @classmethod"""
+
+SAMPLING_BROADCAST_ANCHOR = """\
+            \"multi_modal_kwargs\": self.multi_modal_kwargs,
+            \"prompt_adapter_mapping\": self.prompt_adapter_mapping,
+            \"prompt_adapter_requests\": self.prompt_adapter_requests,
+            \"virtual_engine\": self.virtual_engine,
+            \"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
+            \"finished_requests_ids\": self.finished_requests_ids,
+        }
+        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
+        _add_sampling_metadata_broadcastable_dict(tensor_dict,
+                                                  self.sampling_metadata)"""
+
+SAMPLING_BROADCAST_REPLACEMENT = """\
+            \"multi_modal_kwargs\": self.multi_modal_kwargs,
+            \"gdn_restore_key\": self.gdn_restore_key,
+            \"gdn_capture_points\": self.gdn_capture_points,
+            \"gdn_evict_keys\": self.gdn_evict_keys,
+            \"gdn_segment_offsets\": self.gdn_segment_offsets,
+            \"prompt_adapter_mapping\": self.prompt_adapter_mapping,
+            \"prompt_adapter_requests\": self.prompt_adapter_requests,
+            \"virtual_engine\": self.virtual_engine,
+            \"request_ids_to_seq_ids\": self.request_ids_to_seq_ids,
+            \"finished_requests_ids\": self.finished_requests_ids,
+        }
+        _add_attn_metadata_broadcastable_dict(tensor_dict, self.attn_metadata)
+        _add_sampling_metadata_broadcastable_dict(tensor_dict,
+                                                  self.sampling_metadata)"""
+
+BUILDER_INIT_ANCHOR = """\
+        self.finished_requests_ids = finished_requests_ids
+        self.decode_only = True
+
+        # Intermediate data"""
+
+BUILDER_INIT_REPLACEMENT = """\
+        self.finished_requests_ids = finished_requests_ids
+        self.decode_only = True
+        self.gdn_restore_key = None
+        self.gdn_capture_points = None
+        self.gdn_evict_keys = None
+        self.gdn_segment_offsets = None
+
+        # Intermediate data"""
+
+ADD_SEQ_GROUP_ANCHOR = """\
+    def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
+        \"\"\"Add a sequence group to the builder.\"\"\"
+        seq_ids = seq_group_metadata.seq_data.keys()"""
+
+ADD_SEQ_GROUP_REPLACEMENT = """\
+    def add_seq_group(self, seq_group_metadata: SequenceGroupMetadata):
+        \"\"\"Add a sequence group to the builder.\"\"\"
+        gdn_actions = (
+            seq_group_metadata.gdn_restore_key,
+            seq_group_metadata.gdn_capture_points,
+            seq_group_metadata.gdn_evict_keys,
+            seq_group_metadata.gdn_segment_offsets,
+        )
+        if any(value is not None for value in gdn_actions):
+            if not seq_group_metadata.is_prompt:
+                raise RuntimeError(\"GDN prefix-cache actions require prefill\")
+            if any(value is not None for value in (
+                    self.gdn_restore_key, self.gdn_capture_points,
+                    self.gdn_evict_keys, self.gdn_segment_offsets)):
+                raise RuntimeError(
+                    \"only one GDN prefix-cache action group is supported\")
+            (self.gdn_restore_key, self.gdn_capture_points,
+             self.gdn_evict_keys, self.gdn_segment_offsets) = gdn_actions
+        seq_ids = seq_group_metadata.seq_data.keys()"""
+
+BUILD_RESULT_ANCHOR = """\
+            lora_mapping=lora_mapping,
+            lora_requests=lora_requests,
+            multi_modal_kwargs=multi_modal_kwargs,
+            request_ids_to_seq_ids=request_ids_to_seq_ids,"""
+
+BUILD_RESULT_REPLACEMENT = """\
+            lora_mapping=lora_mapping,
+            lora_requests=lora_requests,
+            multi_modal_kwargs=multi_modal_kwargs,
+            gdn_restore_key=self.gdn_restore_key,
+            gdn_capture_points=self.gdn_capture_points,
+            gdn_evict_keys=self.gdn_evict_keys,
+            gdn_segment_offsets=self.gdn_segment_offsets,
+            request_ids_to_seq_ids=request_ids_to_seq_ids,"""
+
+EXECUTE_KWARGS_ANCHOR = """\
+        seqlen_agnostic_kwargs = {
+            \"finished_requests_ids\": model_input.finished_requests_ids,
+            \"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
+        } if self.has_inner_state else {}
+        if (self.observability_config is not None"""
+
+EXECUTE_KWARGS_REPLACEMENT = """\
+        seqlen_agnostic_kwargs = {
+            \"finished_requests_ids\": model_input.finished_requests_ids,
+            \"request_ids_to_seq_ids\": model_input.request_ids_to_seq_ids,
+        } if self.has_inner_state else {}
+        gdn_prefix_kwargs = {}
+        if model_input.gdn_restore_key is not None:
+            gdn_prefix_kwargs[\"gdn_restore_key\"] = model_input.gdn_restore_key
+        if model_input.gdn_capture_points is not None:
+            gdn_prefix_kwargs[\"gdn_capture_points\"] = (
+                model_input.gdn_capture_points)
+        if model_input.gdn_evict_keys is not None:
+            gdn_prefix_kwargs[\"gdn_evict_keys\"] = model_input.gdn_evict_keys
+        if model_input.gdn_segment_offsets is not None:
+            gdn_prefix_kwargs[\"gdn_segment_offsets\"] = (
+                model_input.gdn_segment_offsets)
+        if (self.observability_config is not None"""
+
+MODEL_CALL_ANCHOR = """\
+                **MultiModalInputs.as_kwargs(multi_modal_kwargs,
+                                             device=self.device),
+                **seqlen_agnostic_kwargs)"""
+
+MODEL_CALL_REPLACEMENT = """\
+                **MultiModalInputs.as_kwargs(multi_modal_kwargs,
+                                             device=self.device),
+                **seqlen_agnostic_kwargs,
+                **gdn_prefix_kwargs)"""
+
+PROFILE_KV_LAYERS_ANCHOR = """\
+        num_layers = self.model_config.get_num_layers(self.parallel_config)"""
+
+PROFILE_KV_LAYERS_REPLACEMENT = """\
+        num_layers = self.model_config.get_num_attention_layers(
+            self.parallel_config)"""
+
+
+def patch_model_runner(model_runner: pathlib.Path) -> None:
+    replace_once(
+        model_runner,
+        HELPER_ANCHOR,
+        HELPER_REPLACEMENT,
+        required=True,
+        already_contains="def _slice_mrope_positions(",
+    )
+    replace_once(
+        model_runner,
+        PREFIX_PAST_ANCHOR,
+        PREFIX_PAST_REPLACEMENT,
+        required=True,
+        already_contains="Must clear prefix_cache_hit so _add_seq_group",
+    )
+    replace_once(
+        model_runner,
+        PARTIAL_HIT_ANCHOR,
+        PARTIAL_HIT_REPLACEMENT,
+        required=True,
+        already_contains="positions, uncomputed_start, None,",
+    )
+    replace_once(
+        model_runner,
+        FULL_HIT_ANCHOR,
+        FULL_HIT_REPLACEMENT,
+        required=True,
+        already_contains="_slice_mrope_positions(positions, -1, None, 1)",
+    )
+    replace_once(
+        model_runner,
+        MULTIMODAL_MROPE_ANCHOR,
+        MULTIMODAL_MROPE_REPLACEMENT,
+        required=True,
+        already_contains="Compute the full MRoPE map once",
+    )
+    replace_once(
+        model_runner,
+        MODEL_INPUT_FIELDS_ANCHOR,
+        MODEL_INPUT_FIELDS_REPLACEMENT,
+        already_contains="gdn_restore_key: Optional[Tuple[int, bytes]]",
+    )
+    replace_once(
+        model_runner,
+        BASE_BROADCAST_ANCHOR,
+        BASE_BROADCAST_REPLACEMENT,
+        already_contains=BASE_BROADCAST_REPLACEMENT,
+    )
+    replace_once(
+        model_runner,
+        SAMPLING_BROADCAST_ANCHOR,
+        SAMPLING_BROADCAST_REPLACEMENT,
+        already_contains=SAMPLING_BROADCAST_REPLACEMENT,
+    )
+    replace_once(
+        model_runner,
+        BUILDER_INIT_ANCHOR,
+        BUILDER_INIT_REPLACEMENT,
+        already_contains="self.gdn_restore_key = None",
+    )
+    replace_once(
+        model_runner,
+        ADD_SEQ_GROUP_ANCHOR,
+        ADD_SEQ_GROUP_REPLACEMENT,
+        already_contains="gdn_actions = (",
+    )
+    replace_once(
+        model_runner,
+        BUILD_RESULT_ANCHOR,
+        BUILD_RESULT_REPLACEMENT,
+        already_contains="gdn_restore_key=self.gdn_restore_key",
+    )
+    replace_once(
+        model_runner,
+        EXECUTE_KWARGS_ANCHOR,
+        EXECUTE_KWARGS_REPLACEMENT,
+        already_contains="gdn_prefix_kwargs = {}",
+    )
+    replace_once(
+        model_runner,
+        MODEL_CALL_ANCHOR,
+        MODEL_CALL_REPLACEMENT,
+        already_contains="**gdn_prefix_kwargs)",
+    )
+    replace_once(
+        model_runner,
+        PROFILE_KV_LAYERS_ANCHOR,
+        PROFILE_KV_LAYERS_REPLACEMENT,
+        required=True,
+        already_contains=PROFILE_KV_LAYERS_REPLACEMENT,
+    )
+
+
+if __name__ == "__main__":
+    patch_model_runner(package_root("vllm") / "worker" / "model_runner.py")

qwen3_6_scripts/patch_ops.sh

@@ -1,141 +1,251 @@
-#!/bin/bash
-set -eo pipefail
-# BI-V100 engine patches for Qwen3.6-35B-A3B (Qwen3_5 architecture)
+#!/usr/bin/env bash
+# BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE 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/
+# Triton situation on BI-V100:
+#   - Standard Triton 2.3.1 is already present in the image.
+#   - HAS_TRITON = False (hardcoded in vendor vllm), but Triton is still used
+#     for TP-mode cache management (custom_cache_manager / libentry).
+#   - The vendor's triton_utils/__init__.py, custom_cache_manager.py, libentry.py
+#     are already correct for standard Triton 2.3.1 — do NOT overwrite them.
+#   - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
+#     GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.
 
-# CRITICAL: cd into this script's directory so all ./relative paths work
-# regardless of WORKDIR in Dockerfile or caller's cwd.
-cd "$(dirname "$0")"
+# Recommended server start command for TP=4 support 256K, needs chunked prefill
+# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
+#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
+#     --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
+#     --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
+#     --tool-call-parser qwen3_coder --reasoning-parser qwen3
+#
+# With prefix caching (GDN align-mode, requires chunked prefill):
+# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
+#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
+#     --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
+#     --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
+#     --tool-call-parser qwen3_coder --reasoning-parser qwen3
+
+set -eo pipefail
+
+# cd into this script's directory so ./relative paths work
+cd "$(dirname "${BASH_SOURCE[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
+build_stage() { printf '[BI100 BUILD] %s\n' "$1" >&2; }
+require_file() {
+    local path=$1
+    [[ -f "$path" ]] || {
+        printf 'required patch source is missing: %s\n' "$path" >&2
+        exit 2
+    }
+}
+install_patch_file() {
+    local source=$1
+    local target=$2
 
-# 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
+    require_file "$source"
+    mkdir -p "$(dirname "$target")"
+    install -m 0644 "$source" "$target"
 }
 
-# --- _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/"
+build_stage "patch script entered"
 
+build_stage "checking offline transformers dependency"
 # --- 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/
+TRANSFORMERS_REQUIRED_VERSION="4.55.3"
+if ! python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
+import importlib.metadata
+import sys
+
+required = sys.argv[1]
+try:
+    installed = importlib.metadata.version("transformers")
+except importlib.metadata.PackageNotFoundError:
+    raise SystemExit(1)
+raise SystemExit(0 if installed == required else 1)
+PY
+then
+  WHEEL_DIR="./wheels"
+  if ! ls "${WHEEL_DIR}/transformers-${TRANSFORMERS_REQUIRED_VERSION}"*.whl >/dev/null 2>&1; then
+    echo "transformers ${TRANSFORMERS_REQUIRED_VERSION} is required, but no offline wheel was found in ${WHEEL_DIR}" >&2
+    exit 2
+  fi
+  python3 -m pip install --no-index --no-deps --find-links="${WHEEL_DIR}" \
+    "transformers==${TRANSFORMERS_REQUIRED_VERSION}"
+fi
+
+python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
+import importlib.metadata
+import sys
+
+required = sys.argv[1]
+installed = importlib.metadata.version("transformers")
+if installed != required:
+    raise SystemExit(
+        f"transformers version mismatch: expected {required}, got {installed}")
+print(f"[ok] transformers {installed}")
+PY
+
+build_stage "discovering Python package roots"
+python3 - <<'PY' > /tmp/qwen36_patch_paths.env
+from patch_utils import package_root, shell_env_line
+
+print(shell_env_line("VLLM_ROOT", package_root("vllm")))
+print(shell_env_line("TRANSFORMERS_ROOT", package_root("transformers")))
+PY
+source /tmp/qwen36_patch_paths.env
+
+echo "VLLM_ROOT=${VLLM_ROOT}"
+echo "TRANSFORMERS_ROOT=${TRANSFORMERS_ROOT}"
+[[ -d "$VLLM_ROOT" ]] || {
+    printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
+    exit 2
+}
+
+VLLM_OVERRIDE_ROOT="./vendor_overrides/vllm"
+[[ -d "$VLLM_OVERRIDE_ROOT" ]] || {
+    printf 'vLLM override directory missing: %s\n' "$VLLM_OVERRIDE_ROOT" >&2
+    exit 2
+}
+
+build_stage "installing authoritative vLLM core block overrides"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/evictor_v2.py" \
+    "${VLLM_ROOT}/core/evictor_v2.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_kv_content_cache.py" \
+    "${VLLM_ROOT}/core/block/cpu_kv_content_cache.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_gpu_block_allocator.py" \
+    "${VLLM_ROOT}/core/block/cpu_gpu_block_allocator.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/prefix_caching_block.py" \
+    "${VLLM_ROOT}/core/block/prefix_caching_block.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/block_table.py" \
+    "${VLLM_ROOT}/core/block/block_table.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block_manager_v2.py" \
+    "${VLLM_ROOT}/core/block_manager_v2.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/sampling_params.py" \
+    "${VLLM_ROOT}/sampling_params.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/model_executor/sampling_metadata.py" \
+    "${VLLM_ROOT}/model_executor/sampling_metadata.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/model_executor/layers/sampler.py" \
+    "${VLLM_ROOT}/model_executor/layers/sampler.py"
+
+build_stage "installing hash-pinned CoreX 3.2.3 extensions"
+bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
+
+build_stage "installing BI100 runtime modules"
+cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
+cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
+cp ./block_major_kv_cache.py "${VLLM_ROOT}/block_major_kv_cache.py"
+cp ./gdn_prefix.py "${VLLM_ROOT}/gdn_prefix.py"
+
+build_stage "installing CoreX paged-KV swap compatibility"
+python3 ./patch_corex_swap_blocks.py
+python3 ./patch_block_major_cache_engine.py
+python3 ./patch_worker_cache_transfer_order.py
+
+# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
+# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
+# (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
+# Our paged_attn.py bypasses it entirely via _forward_prefix_pytorch, which
+# utilizes K-tiling techniques, and also have _forward_decode_pytorch to bypass kernel
+# when context length is high
+cp ./paged_attn.py "${VLLM_ROOT}/attention/ops/paged_attn.py"
+
+# --- model_runner.py: fix prefix_cache_hit stays True in chunked-prefill chunk 2+ ---
+# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
+# leaves prefix_cache_hit=True. Then _add_seq_group uses block_table=computed_block_nums
+# (only the original prefix blocks), ignoring chunk-1 KV cache blocks.
+# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
+# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
+# Fix: set prefix_cache_hit=False for Case 1 so the full block_tables is used.
+python3 ./patch_model_runner.py
+
+build_stage "installing executor startup diagnostics"
+python3 ./patch_executor_startup_debug.py
+python3 ./patch_worker_startup_profile_guard.py
+python3 ./patch_block_major_worker_capacity.py
+
+build_stage "installing transformers Qwen3.5 model support"
+cp -r ./qwen3_5 "${TRANSFORMERS_ROOT}/models/"
+cp -r ./qwen3_5_moe "${TRANSFORMERS_ROOT}/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"
+build_stage "installing vLLM Qwen3.6 model implementation"
+# --- vllm model: Qwen3.6-35B-A3B (Qwen3_5 MoE arch) -------------------------
+cp ./mamba_cache.py "${VLLM_ROOT}/model_executor/models/"
+cp ./qwen3_5.py "${VLLM_ROOT}/model_executor/models/qwen3_5.py"
+python3 ./patch_vllm_qwen3_5.py
 
-# --- 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 under chunked prefill ------
+# Bug: get_output_token_ids_to_return(delta=True) with num_new_tokens=0
+# returns _cached_all_token_ids[-0:] == [0:] (the ENTIRE prompt+output list).
+# Each prefill chunk step adds prompt_len to previous_num_tokens, so a 10K
+# prompt processed in 3 chunks inflates completion_tokens by ~30K.
+# Also adds num_cached_tokens field to RequestMetrics for prefix-cache stats.
+cp ./sequence.py "${VLLM_ROOT}/sequence.py"
 
-# --- sequence.py: fix completion_tokens inflation ----------------------------
-deploy ./sequence.py "sequence.py"
-echo "[patch_ops] sequence.py → /"
+# --- scheduler.py: record num_cached_tokens in RequestMetrics ----------------
+# Reports only the longest prefix backed by both live KV blocks and an exact
+# GDN restore state. Raw KV-only hits must not inflate cached_tokens.
+# serving_chat.py exposes the value in the OpenAI-compatible usage details.
+cp ./scheduler.py "${VLLM_ROOT}/core/scheduler.py"
 
-# --- scheduler.py: record num_cached_tokens ---------------------------------
-deploy ./scheduler.py "core/scheduler.py"
-echo "[patch_ops] scheduler.py → core/"
+build_stage "installing diagnostic initial allocation trace"
+python3 ./patch_block_manager_cache_trace.py
 
-# --- 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"
+build_stage "installing scheduler and attention patches"
+# --- xformers: bypass cudnnFlashAttnForward (head_dim=256 > 128 limit) ------
+# Injects _run_sdpa_fallback (pure matmul+softmax) into xformers.py.
+# Required because head_dim=256 > 128 and ixformer flash attention either
+# crashes (is_causal=True) or produces wrong output (attn_mask path).
+# The fallback uses query_start_loc to derive actual query lengths, so it
+# works correctly during profiling runs with chunked-prefill-style batches.
+# also bypasses auto chunked prefill on
+python3 ./patch_xformers_sdpa_seq.py
+python3 ./patch_xformers_profile.py
 
-# --- 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"
+build_stage "installing API parsers and serving modules"
+# --- tool parser: Qwen3 XML tool call format ---------------------------------
+# Registers "qwen3_coder" parser for Qwen3.6 XML-style tool calls:
+#   <tool_call><function=name><parameter=key>\nvalue\n</parameter></function></tool_call>
+# Use at server start: --tool-call-parser qwen3_coder --enable-auto-tool-choice
+cp ./qwen3coder_tool_parser.py "${VLLM_ROOT}/entrypoints/openai/tool_parsers/"
+python3 ./patch_vllm_tool_parser.py
 
-echo "[patch_ops] DONE — all patches applied via full file replacement"
+# --- reasoning parser: Qwen3 <think>...</think> split ------------------------
+# Adds --reasoning-parser qwen3 support.
+# Routes thinking tokens to reasoning_content, rest to content in the delta.
+# Works together with --tool-call-parser qwen3_coder (think → tool call flow).
+cp -r ./reasoning "${VLLM_ROOT}/"
+cp ./protocol.py "${VLLM_ROOT}/entrypoints/openai/protocol.py"
+cp ./cli_args.py "${VLLM_ROOT}/entrypoints/openai/cli_args.py"
+cp ./serving_chat.py "${VLLM_ROOT}/entrypoints/openai/serving_chat.py"
+cp ./serving_tokenization.py \
+    "${VLLM_ROOT}/entrypoints/openai/serving_tokenization.py"
+cp ./api_server.py "${VLLM_ROOT}/entrypoints/openai/api_server.py"
+cp ./chat_utils.py "${VLLM_ROOT}/entrypoints/chat_utils.py"
+python3 - ./api_server.py \
+        "${VLLM_ROOT}/entrypoints/openai/api_server.py" <<'PY'
+from pathlib import Path
+import sys
+
+source = Path(sys.argv[1]).read_bytes()
+installed = Path(sys.argv[2]).read_bytes()
+if source != installed:
+    raise SystemExit("runtime api_server overlay identity mismatch")
+PY
+
+build_stage "compiling submission Python sources"
+find . -path './wheels' -prune -o -name '*.py' -print0 | xargs -0 python3 -m py_compile
+build_stage "patch script completed"

qwen3_6_scripts/patch_transformers_qwen3_5.py

@@ -2,45 +2,23 @@
 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
+  1. patch_ops.sh locates transformers with importlib.util.find_spec.
+  2. cp -r modified_scripts/qwen3_5* into the detected transformers/models.
   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"
+from patch_utils import package_root, replace_once, replace_one_of
 
-
-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)
+TRANSFORMERS_ROOT = package_root("transformers")
+AUTO_CONFIG = TRANSFORMERS_ROOT / "models" / "auto" / "configuration_auto.py"
+MODELS_INIT = TRANSFORMERS_ROOT / "models" / "__init__.py"
 
 
 def main():
     print(f"=== Patching {AUTO_CONFIG} ===")
-    patch_file(AUTO_CONFIG, [
+    replace_one_of(AUTO_CONFIG, [
         # CONFIG_MAPPING_NAMES: insert qwen3_5 + qwen3_5_moe right after qwen3
         (
             '("qwen3", "Qwen3Config"),',
@@ -50,6 +28,8 @@ def main():
             '("qwen3", "Qwen3Config")\n',
             '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n        ("qwen3_5_moe", "Qwen3_5MoeConfig"),\n',
         ),
+    ], required=True, already_contains='("qwen3_5_moe", "Qwen3_5MoeConfig")')
+    replace_one_of(AUTO_CONFIG, [
         # MODEL_NAMES_MAPPING (model_type -> human readable name)
         (
             '("qwen3", "Qwen3"),',
@@ -59,15 +39,15 @@ def main():
             '("qwen3", "Qwen3")\n',
             '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n        ("qwen3_5_moe", "Qwen3_5_MoE"),\n',
         ),
-    ])
+    ], required=True, already_contains='("qwen3_5_moe", "Qwen3_5_MoE")')
 
     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",
-        ),
-    ])
+    replace_once(
+        MODELS_INIT,
+        "from .qwen3 import *\n",
+        "from .qwen3 import *\n    from .qwen3_5 import *\n    from .qwen3_5_moe import *\n",
+        required=True,
+        already_contains="from .qwen3_5_moe import *")
 
     # Verification
     print("\n=== Verification ===")
@@ -79,28 +59,31 @@ def main():
             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]
+            pkg.__path__ = [str(TRANSFORMERS_ROOT)]
             cu = sys.modules.setdefault(
                 "transformers.configuration_utils", types.ModuleType("transformers.configuration_utils"))
             class _PC:
-                def __init__(self, **kwargs): pass
+                def __init__(self, **kwargs):
+                    return None
             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]
+                m.__path__ = [str(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",
+            str(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",
+            str(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})")
@@ -108,7 +91,7 @@ def main():
         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(f"  [optional] smoke-test failed (may be fine at runtime): {e}")
 
     print("\nDone.")
 

qwen3_6_scripts/patch_utils.py

@@ -0,0 +1,81 @@
+from __future__ import annotations
+
+import importlib.util
+import pathlib
+import shlex
+from typing import Iterable, Optional, Sequence, Tuple
+
+
+def package_root(pkg: str) -> pathlib.Path:
+    spec = importlib.util.find_spec(pkg)
+    if spec is None:
+        raise RuntimeError(f"package not found: {pkg}")
+    if not spec.submodule_search_locations:
+        raise RuntimeError(f"package has no package root: {pkg}")
+    return pathlib.Path(next(iter(spec.submodule_search_locations))).resolve()
+
+
+def ensure_file(path: pathlib.Path) -> pathlib.Path:
+    if not path.is_file():
+        raise FileNotFoundError(str(path))
+    return path
+
+
+def ensure_dir(path: pathlib.Path) -> pathlib.Path:
+    if not path.is_dir():
+        raise FileNotFoundError(str(path))
+    return path
+
+
+def replace_once(path: pathlib.Path,
+                 old: str,
+                 new: str,
+                 *,
+                 required: bool = True,
+                 already_contains: Optional[str] = None) -> bool:
+    path = ensure_file(path)
+    text = path.read_text()
+    marker = already_contains if already_contains is not None else new
+    if marker in text:
+        print(f"[skip] already patched: {path}")
+        return False
+    if old not in text:
+        msg = f"anchor not found in {path}: {old[:120]!r}"
+        if required:
+            raise RuntimeError(msg)
+        print(f"[warn] {msg}")
+        return False
+    path.write_text(text.replace(old, new, 1))
+    print(f"[ok] patched: {path}")
+    return True
+
+
+def replace_one_of(path: pathlib.Path,
+                   replacements: Sequence[Tuple[str, str]],
+                   *,
+                   required: bool = True,
+                   already_contains: Optional[str] = None) -> bool:
+    path = ensure_file(path)
+    text = path.read_text()
+    if already_contains is not None and already_contains in text:
+        print(f"[skip] already patched: {path}")
+        return False
+    for _, new in replacements:
+        if new in text:
+            print(f"[skip] already patched: {path}")
+            return False
+    for old, new in replacements:
+        if old in text:
+            path.write_text(text.replace(old, new, 1))
+            print(f"[ok] patched: {path}")
+            return True
+    anchors = ", ".join(repr(old[:80]) for old, _ in replacements)
+    msg = f"anchor not found in {path}; tried: {anchors}"
+    if required:
+        raise RuntimeError(msg)
+    print(f"[warn] {msg}")
+    return False
+
+
+def shell_env_line(name: str, value: pathlib.Path) -> str:
+    return f"{name}={shlex.quote(str(value))}"

qwen3_6_scripts/patch_vllm_qwen3_5.py

@@ -0,0 +1,73 @@
+"""
+Patches the vLLM model registry and deploys the Qwen3_5 model file.
+
+Deploy steps on the remote machine:
+  1. patch_ops.sh locates vLLM with importlib.util.find_spec.
+  2. cp modified_scripts/qwen3_5.py into the detected vllm model directory.
+  2. python3 modified_scripts/patch_vllm_qwen3_5.py
+
+The registry patch installs Qwen3.6 aliases so /model/config.json does not
+need to be edited by hand.
+"""
+
+import ast
+
+from patch_utils import package_root, replace_once
+
+VLLM_ROOT = package_root("vllm")
+REGISTRY = VLLM_ROOT / "model_executor" / "models" / "registry.py"
+MODEL = VLLM_ROOT / "model_executor" / "models" / "qwen3_5.py"
+
+EXPECTED_REGISTRY_ENTRIES = (
+    '"Qwen3ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
+    '"Qwen3MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
+    '"Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
+    '"Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
+    '"Qwen3_6ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM")',
+    '"Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM")',
+)
+
+
+def main():
+    print(f"=== Patching {REGISTRY} ===")
+    replace_once(
+        REGISTRY,
+        '    "Qwen3ForCausalLM": ("qwen3", "Qwen3ForCausalLM"),\n'
+        '    "Qwen3MoeForCausalLM": ("qwen3_moe", "Qwen3MoeForCausalLM"),',
+        '    "Qwen3ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
+        '    "Qwen3MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),\n'
+        '    "Qwen3_5ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
+        '    "Qwen3_5MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),\n'
+        '    "Qwen3_6ForCausalLM": ("qwen3_5", "Qwen3_5ForCausalLM"),\n'
+        '    "Qwen3_6MoeForCausalLM": ("qwen3_5", "Qwen3_5MoeForCausalLM"),',
+        required=True,
+        already_contains='"Qwen3_6MoeForCausalLM"')
+
+    print("\n=== Static verification ===")
+    model_source = MODEL.read_text(encoding="utf-8")
+    tree = ast.parse(model_source, filename=str(MODEL))
+    class_names = {
+        node.name for node in tree.body if isinstance(node, ast.ClassDef)
+    }
+    required_classes = {"Qwen3_5ForCausalLM", "Qwen3_5MoeForCausalLM"}
+    missing_classes = required_classes - class_names
+    if missing_classes:
+        raise RuntimeError(
+            f"Qwen3.5 model classes missing: {sorted(missing_classes)}")
+
+    registry_source = REGISTRY.read_text(encoding="utf-8")
+    missing_entries = [
+        entry for entry in EXPECTED_REGISTRY_ENTRIES
+        if entry not in registry_source
+    ]
+    if missing_entries:
+        raise RuntimeError(
+            f"Qwen3.5 registry entries missing: {missing_entries}")
+    print("  model syntax and class declarations verified without import")
+    print(f"  registry aliases verified: {len(EXPECTED_REGISTRY_ENTRIES)}")
+
+    print("\nDone. Registry aliases installed; do not edit /model/config.json.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_vllm_tool_parser.py

@@ -0,0 +1,57 @@
+"""
+Patches vLLM 0.6.3 to register Qwen3CoderToolParser under the name "qwen3_coder".
+
+Deploy steps on the remote machine (already called by patch_ops.sh):
+  1. patch_ops.sh locates vLLM with importlib.util.find_spec.
+  2. cp qwen3coder_tool_parser.py into the detected vllm tool_parsers.
+  2. python3 patch_vllm_tool_parser.py
+
+Usage after patching:
+  --tool-call-parser qwen3_coder --enable-auto-tool-choice
+"""
+
+from patch_utils import ensure_dir, package_root, replace_once
+
+VLLM_ROOT = package_root("vllm")
+TOOL_PARSERS_DIR = VLLM_ROOT / "entrypoints" / "openai" / "tool_parsers"
+INIT_FILE = TOOL_PARSERS_DIR / "__init__.py"
+
+
+def main():
+    ensure_dir(TOOL_PARSERS_DIR)
+
+    print(f"=== Patching {INIT_FILE} ===")
+    replace_once(
+        INIT_FILE,
+        "from .mistral_tool_parser import MistralToolParser",
+        "from .mistral_tool_parser import MistralToolParser\n"
+        "from .qwen3coder_tool_parser import Qwen3CoderToolParser",
+        required=True,
+        already_contains="from .qwen3coder_tool_parser import Qwen3CoderToolParser")
+    replace_once(
+        INIT_FILE,
+        '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser"\n]',
+        '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",\n'
+        '    "Qwen3CoderToolParser"\n]',
+        required=True,
+        already_contains='"Qwen3CoderToolParser"')
+
+    print("\n=== Verification ===")
+    try:
+        import importlib.util
+        spec = importlib.util.spec_from_file_location(
+            "qwen3coder_tool_parser",
+            str(TOOL_PARSERS_DIR / "qwen3coder_tool_parser.py"),
+        )
+        mod = importlib.util.module_from_spec(spec)
+        print(f"  Module spec loaded: {spec.name}")
+        print("  (full import requires torch/vllm runtime — skipping exec)")
+    except Exception as e:
+        print(f"  [optional] spec check failed: {e}")
+
+    print("\nDone. Start vLLM server with:")
+    print("  --tool-call-parser qwen3_coder --enable-auto-tool-choice")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_worker_cache_transfer_order.py

@@ -0,0 +1,37 @@
+from patch_utils import package_root, replace_once
+
+
+WORKER = package_root("vllm") / "worker" / "worker.py"
+
+CLEAN_BLOCK = """\
+        if (worker_input.blocks_to_swap_in is not None
+                and worker_input.blocks_to_swap_in.numel() > 0):
+            self.cache_engine[virtual_engine].swap_in(
+                worker_input.blocks_to_swap_in)
+        if (worker_input.blocks_to_swap_out is not None
+                and worker_input.blocks_to_swap_out.numel() > 0):
+            self.cache_engine[virtual_engine].swap_out(
+                worker_input.blocks_to_swap_out)
+"""
+
+ORDERED_BLOCK = """\
+        # BI100 content-addressed CPU KV tier may preserve a victim and reuse
+        # that same GPU slot in one step. Complete every D2H before any H2D.
+        if (worker_input.blocks_to_swap_out is not None
+                and worker_input.blocks_to_swap_out.numel() > 0):
+            self.cache_engine[virtual_engine].swap_out(
+                worker_input.blocks_to_swap_out)
+        if (worker_input.blocks_to_swap_in is not None
+                and worker_input.blocks_to_swap_in.numel() > 0):
+            self.cache_engine[virtual_engine].swap_in(
+                worker_input.blocks_to_swap_in)
+"""
+
+
+replace_once(
+    WORKER,
+    CLEAN_BLOCK,
+    ORDERED_BLOCK,
+    required=True,
+    already_contains="Complete every D2H before any H2D",
+)

qwen3_6_scripts/patch_worker_profile_override.py

@@ -0,0 +1,81 @@
+from patch_utils import package_root, replace_one_of
+
+WORKER = package_root("vllm") / "worker" / "worker.py"
+
+CLEAN_BLOCK = """\
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.cuda.empty_cache()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        self.model_runner.profile_run()
+"""
+
+GUARDED_BLOCK = """\
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.cuda.empty_cache()
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model. Mark this synthetic pass so BI100_PROFILE can skip
+        # timing it by default; profiling real requests is the useful signal.
+        _bi100_prev_startup_profile = os.environ.get("BI100_IN_STARTUP_PROFILE")
+        os.environ["BI100_IN_STARTUP_PROFILE"] = "1"
+        try:
+            self.model_runner.profile_run()
+        finally:
+            if _bi100_prev_startup_profile is None:
+                os.environ.pop("BI100_IN_STARTUP_PROFILE", None)
+            else:
+                os.environ["BI100_IN_STARTUP_PROFILE"] = _bi100_prev_startup_profile
+"""
+
+NEW_BLOCK = """\
+        # Profile the memory usage of the model and get the maximum number of
+        # cache blocks that can be allocated with the remaining free memory.
+        torch.cuda.empty_cache()
+
+        # BI100: Qwen3.6 batched dummy profile_run can trip GDN non-finite
+        # checks before the server starts. If the operator explicitly provides
+        # --num-gpu-blocks-override, trust that conservative capacity value and
+        # skip only the synthetic profile pass. Real inference still uses the
+        # normal GDN fail-fast path.
+        if self.cache_config.num_gpu_blocks_override is not None:
+            cache_block_size = self.get_cache_block_size_bytes()
+            if cache_block_size == 0:
+                num_cpu_blocks = 0
+            else:
+                num_cpu_blocks = int(self.cache_config.swap_space_bytes //
+                                     cache_block_size)
+            logger.warning(
+                "[BI100] skipping worker.profile_run because "
+                "num_gpu_blocks_override=%d was explicitly set",
+                self.cache_config.num_gpu_blocks_override)
+            gc.collect()
+            torch.cuda.empty_cache()
+            return self.cache_config.num_gpu_blocks_override, max(num_cpu_blocks, 0)
+
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model. Mark this synthetic pass so BI100_PROFILE can skip
+        # timing it by default; profiling real requests is the useful signal.
+        _bi100_prev_startup_profile = os.environ.get("BI100_IN_STARTUP_PROFILE")
+        os.environ["BI100_IN_STARTUP_PROFILE"] = "1"
+        try:
+            self.model_runner.profile_run()
+        finally:
+            if _bi100_prev_startup_profile is None:
+                os.environ.pop("BI100_IN_STARTUP_PROFILE", None)
+            else:
+                os.environ["BI100_IN_STARTUP_PROFILE"] = _bi100_prev_startup_profile
+"""
+
+replace_one_of(
+    WORKER,
+    [
+        (GUARDED_BLOCK, NEW_BLOCK),
+        (CLEAN_BLOCK, NEW_BLOCK),
+    ],
+    required=True,
+    already_contains="[BI100] skipping worker.profile_run",
+)

qwen3_6_scripts/patch_worker_startup_profile_guard.py

@@ -0,0 +1,34 @@
+from patch_utils import package_root, replace_one_of
+
+
+WORKER = package_root("vllm") / "worker" / "worker.py"
+
+CLEAN_BLOCK = """\
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model.
+        self.model_runner.profile_run()
+"""
+
+GUARDED_BLOCK = """\
+        # Execute a forward pass with dummy inputs to profile the memory usage
+        # of the model. Mark this synthetic pass so BI100_PROFILE can exclude
+        # it without changing vLLM's normal capacity calculation.
+        _bi100_prev_startup_profile = os.environ.get("BI100_IN_STARTUP_PROFILE")
+        os.environ["BI100_IN_STARTUP_PROFILE"] = "1"
+        try:
+            self.model_runner.profile_run()
+        finally:
+            if _bi100_prev_startup_profile is None:
+                os.environ.pop("BI100_IN_STARTUP_PROFILE", None)
+            else:
+                os.environ["BI100_IN_STARTUP_PROFILE"] = _bi100_prev_startup_profile
+"""
+
+
+replace_one_of(
+    WORKER,
+    [(CLEAN_BLOCK, GUARDED_BLOCK)],
+    required=True,
+    already_contains=(
+        "Mark this synthetic pass so BI100_PROFILE can exclude"),
+)

qwen3_6_scripts/patch_xformers_profile.py

@@ -0,0 +1,121 @@
+"""Install disabled-by-default M1-48 XFormers timing boundaries."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+try:
+    from patch_utils import package_root, replace_once
+except ModuleNotFoundError:
+    from .patch_utils import package_root, replace_once
+
+
+IMPORT_OLD = "from vllm.logger import init_logger"
+IMPORT_NEW = """\
+from vllm.bi100_profile import bi100_timer
+from vllm.logger import init_logger"""
+
+KV_WRITE_OLD = """\
+                PagedAttention.write_to_paged_cache(key, value, key_cache,
+                                                    value_cache,
+                                                    updated_slot_mapping,
+                                                    self.kv_cache_dtype,
+                                                    k_scale, v_scale)"""
+KV_WRITE_NEW = """\
+                with bi100_timer("xformers.kv_write"):
+                    PagedAttention.write_to_paged_cache(
+                        key, value, key_cache, value_cache,
+                        updated_slot_mapping, self.kv_cache_dtype,
+                        k_scale, v_scale)"""
+
+DENSE_OLD = """\
+                out = self._run_memory_efficient_xformers_forward(
+                    query, key, value, prefill_meta, attn_type=attn_type)"""
+DENSE_NEW = """\
+                with bi100_timer("xformers.dense_prefill"):
+                    out = self._run_memory_efficient_xformers_forward(
+                        query, key, value, prefill_meta, attn_type=attn_type)"""
+
+PAGED_OLD = """\
+                out = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    self.kv_cache_dtype,
+                    key_cache,
+                    value_cache,
+                    prefill_meta.block_tables,
+                    prefill_meta.query_start_loc,
+                    prefill_meta.seq_lens_tensor,
+                    prefill_meta.context_lens_tensor,
+                    prefill_meta.max_query_len,
+                    self.alibi_slopes,
+                    self.sliding_window,
+                    k_scale,
+                    v_scale,
+                    is_causal_decoder=(attn_type == AttentionType.DECODER),
+                )"""
+PAGED_NEW = """\
+                with bi100_timer("xformers.paged_prefill"):
+                    out = PagedAttention.forward_prefix(
+                        query,
+                        key,
+                        value,
+                        self.kv_cache_dtype,
+                        key_cache,
+                        value_cache,
+                        prefill_meta.block_tables,
+                        prefill_meta.query_start_loc,
+                        prefill_meta.seq_lens_tensor,
+                        prefill_meta.context_lens_tensor,
+                        prefill_meta.max_query_len,
+                        self.alibi_slopes,
+                        self.sliding_window,
+                        k_scale,
+                        v_scale,
+                        is_causal_decoder=(attn_type == AttentionType.DECODER),
+                    )"""
+
+
+def patch_file(path: Path) -> None:
+    replace_once(
+        path,
+        IMPORT_OLD,
+        IMPORT_NEW,
+        already_contains="from vllm.bi100_profile import bi100_timer",
+    )
+    replace_once(
+        path,
+        KV_WRITE_OLD,
+        KV_WRITE_NEW,
+        already_contains='bi100_timer("xformers.kv_write")',
+    )
+    replace_once(
+        path,
+        DENSE_OLD,
+        DENSE_NEW,
+        already_contains='bi100_timer("xformers.dense_prefill")',
+    )
+    replace_once(
+        path,
+        PAGED_OLD,
+        PAGED_NEW,
+        already_contains='bi100_timer("xformers.paged_prefill")',
+    )
+    text = path.read_text(encoding="utf-8")
+    canonical = "\n".join(line.rstrip(" \t") for line in text.split("\n"))
+    if not canonical.endswith("\n"):
+        canonical += "\n"
+    if canonical != text:
+        path.write_text(canonical, encoding="utf-8")
+
+
+def main() -> None:
+    path = package_root("vllm") / "attention" / "backends" / "xformers.py"
+    print("=== patch_xformers_profile (M1-48 diagnostic timers) ===")
+    print(f"Target: {path}")
+    patch_file(path)
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_xformers_sdpa_batch.py

@@ -0,0 +1,177 @@
+"""
+策略：批量（block-diagonal）fallback — 纯 PyTorch 数学实现
+=============================================================
+构建块对角 causal mask，对整批序列一次 matmul + softmax，
+完全绕开所有硬件 flash attention kernel。
+
+背景：
+  ixformer flshattF:         head_dim > 128 报错拒绝
+  cudnnFlashAttnForward:     接受 head_dim=256，但数值结果错误（输出全"!"）
+  两者大概率是同一硬件单元，ixformer 提前拦截了硬件不支持的配置。
+  纯 matmul 路径完全绕开硬件 flash attention，数值正确。
+
+优点：
+  数值正确。
+  并发请求 prefill attention 在 GPU 上真正并行（一次大 matmul）。
+
+缺点：
+  峰值显存 = total_tokens² × H × dtype_size
+  total_tokens 受 --max-num-batched-tokens 控制，max-model-len 控制不住。
+
+内存参考（fp16，H_local=6，--max-num-batched-tokens=T）：
+  T=2048  → 峰值 ~50  MB
+  T=4096  → 峰值 ~200 MB
+  T=8192  → 峰值 ~800 MB
+  T=16384 → 峰值 ~3.2 GB
+
+Deploy:
+  python3 modified_scripts/patch_xformers_sdpa_batch.py
+"""
+
+from patch_utils import package_root, replace_once
+
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
+
+FALLBACK_METHOD = '''
+    def _run_sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: "XFormersMetadata",
+    ) -> torch.Tensor:
+        """批量纯数学 attention fallback。
+
+        构建块对角 causal mask（等价于 ixformer BlockDiagonalCausalMask），
+        对整批序列一次 matmul + softmax，GPU 并行处理所有序列。
+
+        块对角 mask 结构（seq1 len=3，seq2 len=2）：
+                 s1,0  s1,1  s1,2  s2,0  s2,1
+          s1,0 [  0   -inf  -inf  -inf  -inf ]
+          s1,1 [  0    0   -inf  -inf  -inf ]
+          s1,2 [  0    0    0   -inf  -inf ]
+          s2,0 [-inf  -inf  -inf   0   -inf ]
+          s2,1 [-inf  -inf  -inf   0    0  ]
+
+        softmax 在 float32 下计算防止 float16 溢出，结果转回原始 dtype。
+
+        Args:
+            query : [1, total_prefill_tokens, num_heads,    head_dim]
+            key   : [1, total_prefill_tokens, num_kv_heads, head_dim]
+            value : [1, total_prefill_tokens, num_kv_heads, head_dim]
+        Returns:
+            [1, total_prefill_tokens, num_heads, head_dim]
+        """
+        assert attn_metadata.seq_lens is not None
+        orig_dtype = query.dtype
+        total_tokens = query.shape[1]
+
+        # ── 构建块对角 causal mask [T, T] ────────────────────────────────
+        # 全部初始化为 -inf，再对每条序列的对角块填入下三角 0
+        mask = torch.full(
+            (total_tokens, total_tokens),
+            float("-inf"),
+            dtype=torch.float32,
+            device=query.device,
+        )
+        start = 0
+        for seq_len in attn_metadata.seq_lens:
+            end = start + seq_len
+            mask[start:end, start:end] = torch.tril(
+                torch.zeros(seq_len, seq_len,
+                            dtype=torch.float32, device=query.device)
+            )
+            start = end
+
+        # ── [1, H, T, D]，.contiguous() ──────────────────────────────────
+        q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+        k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+        v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+
+        # ── GQA：展开 KV heads ────────────────────────────────────────────
+        if k_all.shape[1] != q_all.shape[1]:
+            n = q_all.shape[1] // k_all.shape[1]
+            k_all = k_all.repeat_interleave(n, dim=1).contiguous()
+            v_all = v_all.repeat_interleave(n, dim=1).contiguous()
+
+        # ── 纯数学 attention（float32 防溢出）────────────────────────────
+        # [1, H, T, T]
+        attn_w = torch.matmul(q_all.float(), k_all.float().transpose(-2, -1))
+        attn_w = attn_w * self.scale
+        attn_w = attn_w + mask  # 加法广播：mask [T,T] → [1, H, T, T]
+        attn_w = torch.softmax(attn_w, dim=-1)
+
+        out = torch.matmul(attn_w, v_all.float()).to(orig_dtype)
+        # [1, H, T, D] → [1, T, H, D]
+        return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+
+'''
+
+OLD_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            out = xops.memory_efficient_attention_forward(
+                query,
+                key,
+                value,
+                attn_bias=attn_bias[0],
+                p=0.0,
+                scale=self.scale,
+                op = self.attn_op
+                )
+            return out.view_as(original_query)\
+"""
+
+NEW_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            if self.head_size > 128:
+                out = self._run_sdpa_fallback(query, key, value, attn_metadata)
+            else:
+                out = xops.memory_efficient_attention_forward(
+                    query,
+                    key,
+                    value,
+                    attn_bias=attn_bias[0],
+                    p=0.0,
+                    scale=self.scale,
+                    op=self.attn_op,
+                )
+            return out.view_as(original_query)\
+"""
+
+INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
+
+
+def patch_file(path):
+    replace_once(
+        path,
+        INJECT_ANCHOR,
+        FALLBACK_METHOD + INJECT_ANCHOR,
+        required=True,
+        already_contains="def _run_sdpa_fallback(")
+    replace_once(
+        path,
+        OLD_XFORMER_BLOCK,
+        NEW_XFORMER_BLOCK,
+        required=True,
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
+
+
+def main():
+    print("=== patch_xformers_sdpa_batch (batch, pure-math) ===")
+    print(f"Target: {XFORMERS_PATH}")
+    patch_file(XFORMERS_PATH)
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py

@@ -0,0 +1,176 @@
+"""
+策略：批量（block-diagonal）— F.scaled_dot_product_attention，可走硬件 kernel
+=============================================================================
+构建块对角 causal mask，对整批序列一次 F.scaled_dot_product_attention。
+与 patch_xformers_sdpa_batch.py（纯 matmul）的区别：
+  SDPA 会根据 PyTorch/驱动能力分发到最优 kernel（Flash Attention /
+  mem-efficient attention / math fallback），而不是固定走 cublas matmul。
+
+历史说明：
+  该方案最早因输出全"!"而被弃用，后续排查确认"!"由 mamba_cache.py bug
+  引起，与 attention 实现无关。当前恢复此方案用于性能对比测试。
+
+已知硬件限制（BI-V100）：
+  cudnnFlashAttnForward 不支持 is_causal=True（报错）。
+  本实现使用 is_causal=False + 显式块对角 additive mask 规避此限制。
+  若 SDPA 仍分发到有问题的 kernel，回退到 patch_xformers_sdpa_batch.py。
+
+优点（vs 纯 matmul）：
+  SDPA 可分发到 Flash Attention kernel → O(L) 显存、更快的 CUDA kernel。
+
+缺点：
+  依赖硬件 kernel 行为，若 kernel 有 bug 则数值错误（需与 matmul 版对比验证）。
+
+Deploy:
+  python3 modified_scripts/patch_xformers_sdpa_batch_kernel.py
+"""
+
+from patch_utils import package_root, replace_once
+
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
+
+FALLBACK_METHOD = '''
+    def _run_sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: "XFormersMetadata",
+    ) -> torch.Tensor:
+        """批量 F.scaled_dot_product_attention fallback（可走硬件 kernel）。
+
+        构建块对角 causal mask，对整批序列一次 SDPA 调用。
+        SDPA 可分发到 Flash Attention / mem-efficient attention kernel。
+        is_causal=False + 显式 additive mask，规避 cudnnFlashAttnForward
+        不支持 is_causal=True 的限制。
+
+        块对角 mask（seq1 len=3，seq2 len=2）：
+                 s1,0  s1,1  s1,2  s2,0  s2,1
+          s1,0 [  0   -inf  -inf  -inf  -inf ]
+          s1,1 [  0    0   -inf  -inf  -inf ]
+          s1,2 [  0    0    0   -inf  -inf ]
+          s2,0 [-inf  -inf  -inf   0   -inf ]
+          s2,1 [-inf  -inf  -inf   0    0  ]
+
+        Args:
+            query : [1, total_prefill_tokens, num_heads,    head_dim]
+            key   : [1, total_prefill_tokens, num_kv_heads, head_dim]
+            value : [1, total_prefill_tokens, num_kv_heads, head_dim]
+        Returns:
+            [1, total_prefill_tokens, num_heads, head_dim]
+        """
+        import torch.nn.functional as F
+
+        assert attn_metadata.seq_lens is not None
+        orig_dtype = query.dtype
+        total_tokens = query.shape[1]
+
+        # ── 块对角 causal mask [T, T] ─────────────────────────────────────
+        mask = torch.full(
+            (total_tokens, total_tokens),
+            float("-inf"),
+            dtype=orig_dtype,
+            device=query.device,
+        )
+        start = 0
+        for seq_len in attn_metadata.seq_lens:
+            end = start + seq_len
+            mask[start:end, start:end] = torch.tril(
+                torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=query.device)
+            )
+            start = end
+
+        # ── [1, H, T, D] ──────────────────────────────────────────────────
+        q_all = query.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+        k_all = key.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+        v_all = value.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+
+        # ── GQA：展开 KV heads ────────────────────────────────────────────
+        if k_all.shape[1] != q_all.shape[1]:
+            n = q_all.shape[1] // k_all.shape[1]
+            k_all = k_all.repeat_interleave(n, dim=1).contiguous()
+            v_all = v_all.repeat_interleave(n, dim=1).contiguous()
+
+        # ── F.scaled_dot_product_attention（可走硬件 kernel）─────────────
+        # is_causal=False：避免 cudnnFlashAttnForward "not support causal mode"
+        # attn_mask 传 additive float mask（非 bool），SDPA 选择 math/kernel 路径
+        out = F.scaled_dot_product_attention(
+            q_all, k_all, v_all,
+            attn_mask=mask,
+            dropout_p=0.0,
+            is_causal=False,
+            scale=self.scale,
+        )
+        # [1, H, T, D] → [1, T, H, D]
+        return out.squeeze(0).permute(1, 0, 2).contiguous().unsqueeze(0)
+
+'''
+
+OLD_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            out = xops.memory_efficient_attention_forward(
+                query,
+                key,
+                value,
+                attn_bias=attn_bias[0],
+                p=0.0,
+                scale=self.scale,
+                op = self.attn_op
+                )
+            return out.view_as(original_query)\
+"""
+
+NEW_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            if self.head_size > 128:
+                out = self._run_sdpa_fallback(query, key, value, attn_metadata)
+            else:
+                out = xops.memory_efficient_attention_forward(
+                    query,
+                    key,
+                    value,
+                    attn_bias=attn_bias[0],
+                    p=0.0,
+                    scale=self.scale,
+                    op=self.attn_op,
+                )
+            return out.view_as(original_query)\
+"""
+
+INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
+
+
+def patch_file(path):
+    replace_once(
+        path,
+        INJECT_ANCHOR,
+        FALLBACK_METHOD + INJECT_ANCHOR,
+        required=True,
+        already_contains="def _run_sdpa_fallback(")
+    replace_once(
+        path,
+        OLD_XFORMER_BLOCK,
+        NEW_XFORMER_BLOCK,
+        required=True,
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
+
+
+def main():
+    print("=== patch_xformers_sdpa_batch_kernel (batch, F.sdpa + kernel dispatch) ===")
+    print(f"Target: {XFORMERS_PATH}")
+    patch_file(XFORMERS_PATH)
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_xformers_sdpa_seq.py

@@ -0,0 +1,427 @@
+"""
+策略：顺序（per-sequence）fallback — 纯 PyTorch 数学实现
+==========================================================
+逐条序列用 matmul + softmax 手写 attention，完全绕开所有硬件
+flash attention kernel（ixformer / cudnnFlashAttnForward）。
+
+背景：
+  Iluvatar cudnnFlashAttnForward 存在两个已知问题：
+    1. 不支持 is_causal=True（报错）
+    2. 使用 attn_mask 路径时数值结果不正确（静默错误，输出全为"!"）
+  与华为昇腾 910B4 上 llama.cpp --flash-attn off 修复同类问题的原理相同。
+  纯数学路径（matmul + softmax）在任何 PyTorch 后端上结果都正确。
+
+优点：
+  数值正确，不依赖任何硬件特定 attention kernel。
+  峰值显存 = max(seq_len)² × H × dtype_size，由 --max-model-len 控制。
+
+缺点：
+  并发请求的 prefill attention 串行执行。
+  O(L²) 显存（无 flash attention 的 O(L) 优化）。
+
+内存参考（fp16，H_local=6）：
+  max-model-len=4096  → 峰值 ~200 MB
+  max-model-len=8192  → 峰值 ~800 MB
+  max-model-len=16384 → 峰值 ~3.2 GB
+
+额外 patch（arg_utils.py）：
+  vllm 0.6.3 在 max_model_len > 32K 时会自动开启 chunked prefill（无命令行
+  关闭选项），原意是防止 profiling OOM。但 _run_sdpa_fallback 已通过 Q-tiling
+  解决了该问题，chunked prefill 反而会把推理路径从 _run_sdpa_fallback 切换到
+  _forward_prefix_pytorch，属于不必要的行为变更，因此一并禁用该自动逻辑。
+
+Deploy:
+  python3 modified_scripts/patch_xformers_sdpa_seq.py
+"""
+
+from patch_utils import package_root, replace_one_of, replace_once
+
+VLLM_ROOT = package_root("vllm")
+XFORMERS_PATH = VLLM_ROOT / "attention" / "backends" / "xformers.py"
+ARG_UTILS_PATH = VLLM_ROOT / "engine" / "arg_utils.py"
+LOGITS_PROC_PATH = (
+    VLLM_ROOT / "model_executor" / "layers" / "logits_processor.py")
+OUTLINES_DECODING_PATH = (
+    VLLM_ROOT / "model_executor" / "guided_decoding" /
+    "outlines_decoding.py")
+
+# _apply_logits_processors crashes when seq_groups is None (intermediate
+# chunked-prefill chunks on the driver rank). Add an early-return guard.
+_LP_OLD_BLOCK = """\
+def _apply_logits_processors(
+    logits: torch.Tensor,
+    sampling_metadata: SamplingMetadata,
+) -> torch.Tensor:
+    found_logits_processors = False\
+"""
+
+_LP_NEW_BLOCK = """\
+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\
+"""
+
+# Outlines' UNESCAPED_STRING accepts raw JSON control characters, including
+# newlines and tabs. The generated text can therefore satisfy the CFG while
+# still failing json.loads(). Use the RFC 8259 string character constraints.
+_JSON_STRING_OLD_BLOCK = """\
+| UNESCAPED_STRING
+| SIGNED_NUMBER      -> number
+| "true"             -> true
+| "false"            -> false
+| "null"             -> null
+
+array  : "[" [value ("," value)*] "]"
+object : "{" [pair ("," pair)*] "}"
+pair   : UNESCAPED_STRING ":" value
+
+%import common.UNESCAPED_STRING
+%import common.SIGNED_NUMBER
+%import common.WS
+
+%ignore WS\
+"""
+
+_JSON_STRING_V1_BLOCK = r'''| JSON_STRING
+| SIGNED_NUMBER      -> number
+| "true"             -> true
+| "false"            -> false
+| "null"             -> null
+
+array  : "[" [value ("," value)*] "]"
+object : "{" [pair ("," pair)*] "}"
+pair   : JSON_STRING ":" value
+
+JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
+%import common.SIGNED_NUMBER
+%import common.WS
+
+%ignore WS'''
+
+_JSON_STRING_NEW_BLOCK = r'''| JSON_STRING
+| SIGNED_NUMBER      -> number
+| "true"             -> true
+| "false"            -> false
+| "null"             -> null
+
+array  : "[" _ws [value (_ws "," _ws value)*] _ws "]"
+object : "{" _ws [pair (_ws "," _ws pair)*] _ws "}"
+pair   : JSON_STRING _ws ":" _ws value
+_ws    : JSON_WS?
+
+JSON_STRING: /"(\\["\\\/bfnrt]|\\u[0-9a-fA-F]{4}|[^"\\\x00-\x1f])*"/
+JSON_WS: /[ \t\r\n]{1,4}/
+%import common.SIGNED_NUMBER'''
+
+# vllm 0.6.3 自动开启 chunked prefill 的原始块
+_ARG_OLD_BLOCK = """\
+                if (is_gpu and not use_sliding_window and not use_spec_decode
+                        and not self.enable_lora
+                        and not self.enable_prompt_adapter):
+                    self.enable_chunked_prefill = True
+                    logger.warning(
+                        "Chunked prefill is enabled by default for models with "
+                        "max_model_len > 32K. Currently, chunked prefill might "
+                        "not work with some features or models. If you "
+                        "encounter any issues, please disable chunked prefill "
+                        "by setting --enable-chunked-prefill=False.")\
+"""
+
+_ARG_NEW_BLOCK = """\
+                if (is_gpu and not use_sliding_window and not use_spec_decode
+                        and not self.enable_lora
+                        and not self.enable_prompt_adapter):
+                    pass  # skip auto-enable: Q-tiling in _run_sdpa_fallback
+                          # handles long-context memory without chunked prefill\
+"""
+
+_MM_PREFIX_OLD_BLOCK = """\
+        if model_config.is_multimodal_model:
+            if self.enable_prefix_caching:
+                logger.warning(
+                    "--enable-prefix-caching is currently not "
+                    "supported for multimodal models and has been disabled.")
+            self.enable_prefix_caching = False\
+"""
+
+_MM_PREFIX_NEW_BLOCK = """\
+        if model_config.is_multimodal_model:
+            architectures = getattr(model_config.hf_config,
+                                    "architectures", []) or []
+            qwen36_native_vision = "Qwen3_5MoeForCausalLM" in architectures
+            if self.enable_prefix_caching and qwen36_native_vision:
+                logger.info(
+                    "Keeping prefix caching enabled for the Qwen3.6 native "
+                    "vision path.")
+            elif self.enable_prefix_caching:
+                logger.warning(
+                    "--enable-prefix-caching is currently not "
+                    "supported for multimodal models and has been disabled.")
+                self.enable_prefix_caching = False\
+"""
+
+FALLBACK_METHOD = '''
+    def _run_sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: "XFormersMetadata",
+    ) -> torch.Tensor:
+        """纯数学 causal attention fallback，带 Q-tiling 内存优化。
+
+        调用时机：kv_cache.numel()==0（profiling 阶段）。
+        此路径无 KV 缓存前缀，KV 长度 == query 长度。
+
+        内存优化（Q-tiling，与 Flash Attention 同思路）：
+          将 Q 分成 _Q_CHUNK 大小的子块逐块计算，每块峰值内存
+          O(_Q_CHUNK × q_len) 而非 O(q_len²)。
+          profiling 阶段序列可能达到 max_model_len（如 20K tokens），
+          不加 Q-tiling 会产生 9.6 GB 矩阵直接 OOM。
+
+        softmax 在 float32 下计算以防止 float16 溢出，结果转回原始 dtype。
+
+        Args:
+            query : [1, total_query_tokens, num_heads,    head_dim]
+            key   : [1, total_query_tokens, num_kv_heads, head_dim]
+            value : [1, total_query_tokens, num_kv_heads, head_dim]
+        Returns:
+            [1, total_query_tokens, num_heads, head_dim]
+        """
+        _Q_CHUNK = 256  # 与 _forward_prefix_pytorch 的 _ATTN_Q_CHUNK 保持一致
+
+        assert attn_metadata.seq_lens is not None
+        orig_dtype = query.dtype
+        num_seqs = len(attn_metadata.seq_lens)
+
+        # 推导每条序列的实际 query 长度。
+        # 正常 prefill 时 q_len == seq_len；如果将来遇到 chunked 场景，
+        # query_start_loc 记录的是真实 query token 数（非全序列长度）。
+        if (attn_metadata.query_start_loc is not None
+                and len(attn_metadata.query_start_loc) == num_seqs + 1):
+            q_lens = [
+                int(attn_metadata.query_start_loc[i + 1].item()) -
+                int(attn_metadata.query_start_loc[i].item())
+                for i in range(num_seqs)
+            ]
+        else:
+            q_lens = list(attn_metadata.seq_lens)
+
+        q_flat = query.squeeze(0)   # [T, H,   D]
+        k_flat = key.squeeze(0)     # [T, Hkv, D]
+        v_flat = value.squeeze(0)
+
+        output = torch.empty_like(q_flat)
+        seq_start = 0
+        for q_len in q_lens:
+            seq_end = seq_start + q_len
+
+            # 当前序列的完整 K/V（此路径无前缀，KV == Q）
+            k_s = k_flat[seq_start:seq_end].permute(1, 0, 2).float()  # [Hkv, q_len, D]
+            v_s = v_flat[seq_start:seq_end].permute(1, 0, 2).float()  # [Hkv, q_len, D]
+
+            # GQA：展开 KV heads 至与 query heads 一致
+            if k_s.shape[0] != self.num_heads:
+                n = self.num_heads // k_s.shape[0]
+                k_s = k_s.repeat_interleave(n, dim=0).contiguous()
+                v_s = v_s.repeat_interleave(n, dim=0).contiguous()
+
+            # k_pos 用于因果掩码
+            k_pos = torch.arange(q_len, device=query.device)
+
+            # Q-tiling：分块处理 query，峰值内存 O(_Q_CHUNK × q_len)
+            for qc_start in range(0, q_len, _Q_CHUNK):
+                qc_end = min(qc_start + _Q_CHUNK, q_len)
+
+                # [H, qc, D]
+                q_c = q_flat[seq_start + qc_start:seq_start + qc_end] \
+                      .permute(1, 0, 2).float()
+
+                # [H, qc, q_len]
+                attn_w = torch.matmul(q_c, k_s.transpose(-2, -1)) * self.scale
+
+                # 因果掩码：q_c 里位置 j 只能看 k_pos <= j（相对位置）
+                qc_q_pos = torch.arange(qc_start, qc_end, device=query.device)
+                mask = k_pos.unsqueeze(0) > qc_q_pos.unsqueeze(1)
+                attn_w = attn_w.masked_fill(mask.unsqueeze(0), float("-inf"))
+
+                attn_w = torch.softmax(attn_w, dim=-1)
+                out_c = torch.matmul(attn_w, v_s).to(orig_dtype)  # [H, qc, D]
+
+                output[seq_start + qc_start:seq_start + qc_end] = (
+                    out_c.permute(1, 0, 2))
+
+            seq_start = seq_end
+
+        return output.unsqueeze(0)  # [1, T, H, D]
+
+'''
+
+OLD_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            out = xops.memory_efficient_attention_forward(
+                query,
+                key,
+                value,
+                attn_bias=attn_bias[0],
+                p=0.0,
+                scale=self.scale,
+                op = self.attn_op
+                )
+            return out.view_as(original_query)\
+"""
+
+NEW_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            if self.head_size > 128:
+                out = self._run_sdpa_fallback(query, key, value, attn_metadata)
+            else:
+                out = xops.memory_efficient_attention_forward(
+                    query,
+                    key,
+                    value,
+                    attn_bias=attn_bias[0],
+                    p=0.0,
+                    scale=self.scale,
+                    op=self.attn_op,
+                )
+            return out.view_as(original_query)\
+"""
+
+INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
+
+_PREFIX_CALL_OLD_BLOCK = """\
+                out = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    self.kv_cache_dtype,
+                    key_cache,
+                    value_cache,
+                    prefill_meta.block_tables,
+                    prefill_meta.query_start_loc,
+                    prefill_meta.seq_lens_tensor,
+                    prefill_meta.context_lens_tensor,
+                    prefill_meta.max_query_len,
+                    self.alibi_slopes,
+                    self.sliding_window,
+                    k_scale,
+                    v_scale,
+                )\
+"""
+
+_PREFIX_CALL_NEW_BLOCK = """\
+                out = PagedAttention.forward_prefix(
+                    query,
+                    key,
+                    value,
+                    self.kv_cache_dtype,
+                    key_cache,
+                    value_cache,
+                    prefill_meta.block_tables,
+                    prefill_meta.query_start_loc,
+                    prefill_meta.seq_lens_tensor,
+                    prefill_meta.context_lens_tensor,
+                    prefill_meta.max_query_len,
+                    self.alibi_slopes,
+                    self.sliding_window,
+                    k_scale,
+                    v_scale,
+                    is_causal_decoder=(attn_type == AttentionType.DECODER),
+                )\
+"""
+
+
+def patch_file(path):
+    replace_once(
+        path,
+        INJECT_ANCHOR,
+        FALLBACK_METHOD + INJECT_ANCHOR,
+        required=True,
+        already_contains="def _run_sdpa_fallback(")
+    replace_once(
+        path,
+        OLD_XFORMER_BLOCK,
+        NEW_XFORMER_BLOCK,
+        required=True,
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
+    replace_once(
+        path,
+        _PREFIX_CALL_OLD_BLOCK,
+        _PREFIX_CALL_NEW_BLOCK,
+        required=True,
+        already_contains=(
+            "is_causal_decoder=(attn_type == AttentionType.DECODER)"))
+
+
+def patch_arg_utils(path):
+    replace_once(
+        path,
+        _ARG_OLD_BLOCK,
+        _ARG_NEW_BLOCK,
+        required=True,
+        already_contains="skip auto-enable: Q-tiling")
+    replace_once(
+        path,
+        _MM_PREFIX_OLD_BLOCK,
+        _MM_PREFIX_NEW_BLOCK,
+        required=True,
+        already_contains="Keeping prefix caching enabled for the Qwen3.6")
+
+
+def patch_logits_processor(path):
+    replace_once(
+        path,
+        _LP_OLD_BLOCK,
+        _LP_NEW_BLOCK,
+        required=True,
+        already_contains="intermediate chunked-prefill chunk")
+
+
+def patch_outlines_json_grammar(path):
+    replace_one_of(
+        path,
+        [
+            (_JSON_STRING_V1_BLOCK, _JSON_STRING_NEW_BLOCK),
+            (_JSON_STRING_OLD_BLOCK, _JSON_STRING_NEW_BLOCK),
+        ],
+        required=True,
+        already_contains="JSON_WS:")
+
+
+def main():
+    print("=== patch_xformers_sdpa_seq (sequential, pure-math) ===")
+    print(f"Target: {XFORMERS_PATH}")
+    patch_file(XFORMERS_PATH)
+
+    print("\n=== patch_arg_utils (disable chunked-prefill auto-enable) ===")
+    print(f"Target: {ARG_UTILS_PATH}")
+    patch_arg_utils(ARG_UTILS_PATH)
+
+    print("\n=== patch_logits_processor (seq_groups=None guard for chunked prefill) ===")
+    print(f"Target: {LOGITS_PROC_PATH}")
+    patch_logits_processor(LOGITS_PROC_PATH)
+
+    print("\n=== patch_outlines_json_grammar (reject raw control chars) ===")
+    print(f"Target: {OUTLINES_DECODING_PATH}")
+    patch_outlines_json_grammar(OUTLINES_DECODING_PATH)
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py

@@ -0,0 +1,166 @@
+"""
+策略：顺序（per-sequence）— F.scaled_dot_product_attention，可走硬件 kernel
+=============================================================================
+逐条序列调用 F.scaled_dot_product_attention，is_causal=False + 显式因果 mask。
+与 patch_xformers_sdpa_seq.py（纯 matmul）的区别：
+  SDPA 可分发到 Flash Attention / mem-efficient attention kernel，
+  而纯 matmul 固定走 cublas。
+
+硬件限制（BI-V100）：
+  cudnnFlashAttnForward 不支持 is_causal=True（直接报错）。
+  必须使用 is_causal=False + 显式 additive causal mask。
+  每条序列单独构造上三角 -inf mask，peak 显存 = max(seq_len)² × dtype，
+  比 batch 版的 total_tokens² 小得多。
+
+与 batch_kernel 的对比：
+  seq_kernel:   显存小，peak = max_single_seq²；并发 prefill 串行排队
+  batch_kernel: 显存大，peak = total_tokens²；并发 prefill 一次并行处理，
+                通过 --max-num-batched-tokens 控制 total_tokens 上限
+
+Deploy:
+  python3 modified_scripts/patch_xformers_sdpa_seq_kernel.py
+"""
+
+from patch_utils import package_root, replace_once
+
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
+
+FALLBACK_METHOD = '''
+    def _run_sdpa_fallback(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_metadata: "XFormersMetadata",
+    ) -> torch.Tensor:
+        """顺序 F.scaled_dot_product_attention fallback（可走硬件 kernel）。
+
+        逐条序列调用 SDPA，is_causal=False + 显式上三角 additive mask。
+        cudnnFlashAttnForward 不支持 is_causal=True，必须用显式 mask。
+        逐序列构造 mask，peak 显存 = max(seq_len)² × dtype（远小于 batch 版）。
+
+        Args:
+            query : [1, total_prefill_tokens, num_heads,    head_dim]
+            key   : [1, total_prefill_tokens, num_kv_heads, head_dim]
+            value : [1, total_prefill_tokens, num_kv_heads, head_dim]
+        Returns:
+            [1, total_prefill_tokens, num_heads, head_dim]
+        """
+        import torch.nn.functional as F
+
+        assert attn_metadata.seq_lens is not None
+        orig_dtype = query.dtype
+
+        q_flat = query.squeeze(0)   # [T, H,   D]
+        k_flat = key.squeeze(0)     # [T, Hkv, D]
+        v_flat = value.squeeze(0)
+
+        output = torch.empty_like(q_flat)
+        start = 0
+        for seq_len in attn_metadata.seq_lens:
+            end = start + seq_len
+            # [1, H, L, D]
+            q_s = q_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
+            k_s = k_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
+            v_s = v_flat[start:end].permute(1, 0, 2).contiguous().unsqueeze(0)
+
+            # GQA：展开 KV heads
+            if k_s.shape[1] != q_s.shape[1]:
+                n = q_s.shape[1] // k_s.shape[1]
+                k_s = k_s.repeat_interleave(n, dim=1).contiguous()
+                v_s = v_s.repeat_interleave(n, dim=1).contiguous()
+
+            # 逐序列因果 mask [L, L]，上三角 -inf
+            causal_mask = torch.tril(
+                torch.zeros(seq_len, seq_len, dtype=orig_dtype, device=q_s.device)
+            )
+            causal_mask = causal_mask.masked_fill(
+                torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool,
+                                     device=q_s.device), diagonal=1),
+                float("-inf"),
+            )
+
+            # is_causal=False + 显式 mask，规避 cudnnFlashAttnForward 不支持 is_causal=True
+            out_s = F.scaled_dot_product_attention(
+                q_s, k_s, v_s,
+                attn_mask=causal_mask,
+                dropout_p=0.0,
+                is_causal=False,
+                scale=self.scale,
+            )
+            # [1, H, L, D] → [L, H, D]
+            output[start:end] = out_s.squeeze(0).permute(1, 0, 2).to(orig_dtype)
+            start = end
+
+        return output.unsqueeze(0)  # [1, T, H, D]
+
+'''
+
+OLD_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            out = xops.memory_efficient_attention_forward(
+                query,
+                key,
+                value,
+                attn_bias=attn_bias[0],
+                p=0.0,
+                scale=self.scale,
+                op = self.attn_op
+                )
+            return out.view_as(original_query)\
+"""
+
+NEW_XFORMER_BLOCK = """\
+        self.attn_op = xops.fmha.flash.FwOp()
+        if self.alibi_slopes is None:
+            # Add the batch dimension.
+            query = query.unsqueeze(0)
+            key = key.unsqueeze(0)
+            value = value.unsqueeze(0)
+            if self.head_size > 128:
+                out = self._run_sdpa_fallback(query, key, value, attn_metadata)
+            else:
+                out = xops.memory_efficient_attention_forward(
+                    query,
+                    key,
+                    value,
+                    attn_bias=attn_bias[0],
+                    p=0.0,
+                    scale=self.scale,
+                    op=self.attn_op,
+                )
+            return out.view_as(original_query)\
+"""
+
+INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
+
+
+def patch_file(path):
+    replace_once(
+        path,
+        INJECT_ANCHOR,
+        FALLBACK_METHOD + INJECT_ANCHOR,
+        required=True,
+        already_contains="def _run_sdpa_fallback(")
+    replace_once(
+        path,
+        OLD_XFORMER_BLOCK,
+        NEW_XFORMER_BLOCK,
+        required=True,
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
+
+
+def main():
+    print("=== patch_xformers_sdpa_seq_kernel (seq, F.sdpa + kernel dispatch) ===")
+    print(f"Target: {XFORMERS_PATH}")
+    patch_file(XFORMERS_PATH)
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/SHA256SUMS

@@ -0,0 +1,12 @@
+534019b3c2ad2d2c65492b01a975874ee440026eda2e8666bc3c1dc8a0a0a6f6  corex_attn_head_rms_norm.so
+7e2aafd8dc755b0ee16c3b9bb812b95548fc042bbaa840dd9db7d2c51a10474c  corex_block_major_kv_transfer.so
+ad4ea7707bb2f2bfe04e07a7ad5fd58a647232be70a3056937a0d738c8254bff  corex_fused_paged_prefill.so
+1856c86e3100415061aa698a48bdeff3fe785994b45b4e72a42cd9158552a7d8  corex_gdn_beta_decay.so
+957c7518f5831299fc73f19a4ca2aa3c8231afe9ea7c979127b4f426cd9d6906  corex_gdn_causal_conv.so
+ec2d11fa82d9d0816a6da53e62605e962786fa20ecd5f62e50f9d43087fc4d67  corex_gdn_gated_norm.so
+27b7ae2ce4fe173336355d72a2678d043df4bd1ed85e9231a99bfb81885a6ce3  corex_gdn_packed_decode.so
+015b61046ad73d8f12d754f7a87d4f6cba33070af1c079879e15b71a94571670  corex_gdn_qk_map.so
+0eb120e89608bb5b64ca4356a5d3d362121806d081ccc1ccf346dac472a819ec  corex_moe_direct_routed.so
+d26f2fa39c3921a95793786601e90cf6ebadd06f1d752af541bf82c21acbc1c9  corex_moe_exact_reduce.so
+50b0b44c1da779bb2c03419ed549aee9bb922d1f9bab8b7f11a3d91cca0d21c3  corex_moe_weight_gather.so
+e944ec0528ed9b6cb74518de3c57e3730543a7bdebc872f993bfdc8424f13e6b  corex_paged_kv_gather.so

qwen3_6_scripts/prefix_prefill.py

@@ -1,900 +0,0 @@
-# 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/protocol.py

@@ -1,5 +1,6 @@
 # Adapted from
 # https://github.com/lm-sys/FastChat/blob/168ccc29d3f7edc50823016105c024fe2282732a/fastchat/protocol/openai_api_protocol.py
+import json
 import time
 from argparse import Namespace
 from typing import Any, Dict, List, Literal, Optional, Union
@@ -57,10 +58,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):
 
 class OpenAIBaseModel(BaseModel):
     # OpenAI API does not allow extra fields
-    # Real-world clients (replay, third-party SDKs) may send extra fields
-    # like service_tier, store, metadata, reasoning_effort, etc.
-    # "ignore" accepts the request and silently drops unknown fields.
-    model_config = ConfigDict(extra="ignore")
+    model_config = ConfigDict(extra="forbid")
 
 
 class ErrorResponse(OpenAIBaseModel):
@@ -143,6 +141,19 @@ class FunctionDefinition(OpenAIBaseModel):
     name: str
     description: Optional[str] = None
     parameters: Optional[Dict[str, Any]] = None
+    # OpenAI clients commonly serialize strict=false explicitly. It is a
+    # semantic no-op, so accept it but keep it out of the tokenizer template.
+    # strict=true requires constrained tool decoding that this runtime does not
+    # provide and must not be silently degraded to ordinary auto tool choice.
+    strict: Optional[bool] = Field(default=None, exclude=True)
+
+    @model_validator(mode="after")
+    def reject_unsupported_strict_tools(self):
+        if self.strict is True:
+            raise ValueError(
+                "Function tools with strict=true are not supported by this "
+                "runtime.")
+        return self
 
 
 class ChatCompletionToolsParam(OpenAIBaseModel):
@@ -169,10 +180,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
     logprobs: Optional[bool] = False
     top_logprobs: Optional[int] = 0
     max_tokens: Optional[int] = None
-    # OpenAI newer API uses max_completion_tokens as alias for max_tokens.
-    # CCCL namespace_wrapped.cu pattern: accept alternate names for same concept.
-    # Competition evaluator sends max_completion_tokens (values: 8192, 32768, 65536).
-    max_completion_tokens: Optional[int] = None
     n: Optional[int] = 1
     presence_penalty: Optional[float] = 0.0
     response_format: Optional[ResponseFormat] = None
@@ -184,15 +191,12 @@ class ChatCompletionRequest(OpenAIBaseModel):
     top_p: Optional[float] = 1.0
     tools: Optional[List[ChatCompletionToolsParam]] = None
     tool_choice: Optional[Union[Literal["none"], Literal["auto"],
-                                Literal["required"],
                                 ChatCompletionNamedToolChoiceParam]] = "none"
+    thinking: Optional[Union[bool, str, Dict[str, Any]]] = None
 
     # NOTE this will be ignored by VLLM -- the model determines the behavior
     parallel_tool_calls: Optional[bool] = False
     user: Optional[str] = None
-    # Qwen3/OpenAI thinking/reasoning control.
-    # Competition evaluator sends thinking={enable:true/false}.
-    thinking: Optional[dict] = None
 
     # doc: begin-chat-completion-sampling-params
     best_of: Optional[int] = None
@@ -209,6 +213,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
     spaces_between_special_tokens: bool = True
     truncate_prompt_tokens: Optional[Annotated[int, Field(ge=1)]] = None
     prompt_logprobs: Optional[int] = None
+    bi100_prompt_logprobs_sample_positions: Optional[List[int]] = None
     # doc: end-chat-completion-sampling-params
 
     # doc: begin-chat-completion-extra-params
@@ -309,8 +314,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
         max_tokens = self.max_tokens
         if max_tokens is None:
             max_tokens = default_max_tokens
-        if default_max_tokens > 0:
-            max_tokens = min(max_tokens, default_max_tokens)
 
         n = self.n if self.n is not None else 1
         temperature = self.temperature if self.temperature is not None else 0.0
@@ -327,10 +330,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
         max_tokens = self.max_tokens
         if max_tokens is None:
             max_tokens = default_max_tokens
-        # Clamp to available context space so requests with max_tokens ≥
-        # max_model_len don't get rejected with HTTP 400.
-        if default_max_tokens > 0:
-            max_tokens = min(max_tokens, default_max_tokens)
 
         prompt_logprobs = self.prompt_logprobs
         if prompt_logprobs is None and self.echo:
@@ -340,7 +339,10 @@ class ChatCompletionRequest(OpenAIBaseModel):
         guided_json_from_schema = None
         if self.response_format is not None:
             if self.response_format.type == "json_object":
-                guided_json_object = True
+                # The generic CFG backend has a stateful first-request bug in
+                # this vLLM/Outlines build. A generic object schema has the
+                # same API semantics and uses the stable regex backend.
+                guided_json_from_schema = {"type": "object"}
             elif (self.response_format.type == "json_schema"
                   and self.response_format.json_schema is not None
                   and self.response_format.json_schema.json_schema is not None):
@@ -373,6 +375,8 @@ class ChatCompletionRequest(OpenAIBaseModel):
             stop_token_ids=self.stop_token_ids,
             logprobs=self.top_logprobs if self.logprobs else None,
             prompt_logprobs=prompt_logprobs,
+            prompt_logprob_positions=(
+                self.bi100_prompt_logprobs_sample_positions),
             ignore_eos=self.ignore_eos,
             max_tokens=max_tokens,
             min_tokens=self.min_tokens,
@@ -414,52 +418,6 @@ class ChatCompletionRequest(OpenAIBaseModel):
         reasoning_content is intentionally kept — chat_utils.py wraps it as
         <think>...</think> for multi-turn reasoning history.
         """
-        # Map max_completion_tokens → max_tokens (OpenAI API v2 name)
-        if data.get("max_completion_tokens") is not None and data.get("max_tokens") is None:
-            data["max_tokens"] = data["max_completion_tokens"]
-
-        # n > max_num_seqs: clamp handled in serving_chat.py via scheduler check.
-        # With max_num_seqs=2, n=2 should work. n>2 will be clamped there.
-
-        # Map thinking parameter → chat_template_kwargs.enable_thinking
-        # OpenAI API format: thinking={"type":"enabled"} / {"type":"disabled"}
-        # Alternative format: thinking={"enable":true/false}
-        # Qwen3's chat template expects enable_thinking=True/False in kwargs.
-        thinking = data.get("thinking")
-        thinking_explicitly_set = False
-        if isinstance(thinking, dict):
-            # Try OpenAI format first: {"type": "enabled"/"disabled"}
-            thinking_type = thinking.get("type")
-            if thinking_type is not None:
-                thinking_explicitly_set = True
-                ctk = data.get("chat_template_kwargs") or {}
-                ctk["enable_thinking"] = (thinking_type == "enabled"
-                                          or thinking_type is True)
-                data["chat_template_kwargs"] = ctk
-            else:
-                # Fallback: {"enable": true/false}
-                enable = thinking.get("enable")
-                if enable is not None:
-                    thinking_explicitly_set = True
-                    ctk = data.get("chat_template_kwargs") or {}
-                    ctk["enable_thinking"] = bool(enable)
-                    data["chat_template_kwargs"] = ctk
-
-        # CRITICAL: When tools are present with tool_choice=auto and thinking
-        # is NOT explicitly requested, disable thinking to preserve token budget
-        # for tool call XML generation. Without this, the model spends all
-        # tokens on <think>...</think> and finishes before emitting <tool_call>.
-        # This matches the competition reference (sub168: d03 in 2.12s).
-        if not thinking_explicitly_set:
-            has_tools = data.get("tools") is not None and len(data.get("tools", [])) > 0
-            tc = data.get("tool_choice")
-            tool_choice_active = (tc == "auto" or (tc is None and has_tools)
-                                  or isinstance(tc, dict))
-            if has_tools and tool_choice_active:
-                ctk = data.get("chat_template_kwargs") or {}
-                ctk["enable_thinking"] = False
-                data["chat_template_kwargs"] = ctk
-
         messages = data.get("messages")
         if not isinstance(messages, list):
             return data
@@ -468,24 +426,131 @@ class ChatCompletionRequest(OpenAIBaseModel):
             if not isinstance(msg, dict):
                 normalized.append(msg)
                 continue
+            tool_calls = msg.get("tool_calls")
+            if isinstance(tool_calls, list):
+                normalized_calls = []
+                for call in tool_calls:
+                    if not isinstance(call, dict):
+                        normalized_calls.append(call)
+                        continue
+                    function = call.get("function")
+                    if not isinstance(function, dict):
+                        normalized_calls.append(call)
+                        continue
+                    arguments = function.get("arguments")
+                    if isinstance(arguments, dict):
+                        arguments = json.dumps(
+                            arguments,
+                            ensure_ascii=False,
+                            separators=(",", ":"),
+                        )
+                    elif isinstance(arguments, str):
+                        try:
+                            decoded_arguments = json.loads(arguments)
+                        except json.JSONDecodeError as exc:
+                            raise ValueError(
+                                "Tool call arguments are not valid JSON."
+                            ) from exc
+                        if not isinstance(decoded_arguments, dict):
+                            raise ValueError(
+                                "Tool call arguments must decode to a JSON "
+                                "object.")
+                    elif arguments is not None:
+                        raise ValueError(
+                            "Tool call arguments must be a JSON object or a "
+                            "JSON-encoded object string.")
+                    if arguments is not None:
+                        function = {**function, "arguments": arguments}
+                        call = {**call, "function": function}
+                    normalized_calls.append(call)
+                msg = {**msg, "tool_calls": normalized_calls}
             if msg.get("content") is None:
-                # Allow tool_calls messages and tool-role messages without content.
-                # CCCL namespace pattern: accept valid alternate message formats.
-                if msg.get("reasoning_content") is not None:
-                    msg = {**msg, "content": ""}
-                elif msg.get("tool_calls") is not None:
-                    msg = {**msg, "content": ""}
-                elif msg.get("role") == "tool":
-                    msg = {**msg, "content": ""}
-                else:
+                if (msg.get("reasoning_content") is None
+                        and not msg.get("tool_calls")):
                     raise ValueError(
-                        "Each message must have at least one of 'content', "
-                        "'reasoning_content', or 'tool_calls'.")
-                    
+                        "Each message must have at least one of 'content' or "
+                        "'reasoning_content', or contain 'tool_calls'.")
+                msg = {**msg, "content": ""}
+            if (msg.get("role") == "system"
+                    and isinstance(msg.get("content"), list)):
+                content_parts = msg["content"]
+                if all(
+                        isinstance(part, dict)
+                        and part.get("type") == "text"
+                        and isinstance(part.get("text"), str)
+                        for part in content_parts):
+                    # Match chat_utils' existing text-part semantics before
+                    # combining multiple system messages for Qwen.
+                    msg = {
+                        **msg,
+                        "content": "\n".join(
+                            part["text"] for part in content_parts),
+                    }
             normalized.append(msg)
+
+        # Qwen's tokenizer template accepts at most one system message and
+        # requires it to be first. OpenAI-compatible clients may send several
+        # system messages, including after conversation history. Preserve
+        # their order and semantics by merging text content at the beginning.
+        system_messages = [
+            msg for msg in normalized
+            if isinstance(msg, dict) and msg.get("role") == "system"
+        ]
+        if system_messages:
+            system_contents = [
+                msg.get("content") for msg in system_messages
+            ]
+            if all(isinstance(content, str)
+                   for content in system_contents):
+                merged_system = {
+                    **system_messages[0],
+                    "content": "\n\n".join(system_contents),
+                }
+                normalized = [merged_system] + [
+                    msg for msg in normalized
+                    if not (isinstance(msg, dict)
+                            and msg.get("role") == "system")
+                ]
         data = {**data, "messages": normalized}
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def normalize_thinking(cls, data):
+        thinking = data.get("thinking")
+        if thinking is None:
+            return data
+
+        enable_thinking: Optional[bool] = None
+        if thinking is False:
+            enable_thinking = False
+        elif thinking is True:
+            enable_thinking = True
+        elif isinstance(thinking, str):
+            lowered = thinking.lower()
+            if lowered == "disabled":
+                enable_thinking = False
+            elif lowered == "enabled":
+                enable_thinking = True
+        elif isinstance(thinking, dict):
+            thinking_type = thinking.get("type")
+            if isinstance(thinking_type, str):
+                lowered = thinking_type.lower()
+                if lowered == "disabled":
+                    enable_thinking = False
+                elif lowered == "enabled":
+                    enable_thinking = True
+
+        if enable_thinking is None:
+            raise ValueError(
+                "`thinking` must be false, \"disabled\", true, \"enabled\", "
+                "or an object with type \"disabled\"/\"enabled\".")
+
+        chat_template_kwargs = dict(data.get("chat_template_kwargs") or {})
+        chat_template_kwargs["enable_thinking"] = enable_thinking
+        data = {**data, "chat_template_kwargs": chat_template_kwargs}
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def validate_stream_options(cls, data):
@@ -517,6 +582,38 @@ class ChatCompletionRequest(OpenAIBaseModel):
 
         return data
 
+    @model_validator(mode="before")
+    @classmethod
+    def validate_bi100_prompt_logprob_sample(cls, data):
+        positions = data.get("bi100_prompt_logprobs_sample_positions")
+        if positions is None:
+            return data
+        if (
+            not isinstance(positions, list)
+            or not positions
+            or len(positions) > 4096
+            or any(
+                not isinstance(position, int)
+                or isinstance(position, bool)
+                or position <= 0
+                or position >= 262144
+                for position in positions
+            )
+            or positions != sorted(set(positions))
+        ):
+            raise ValueError(
+                "`bi100_prompt_logprobs_sample_positions` must be a sorted "
+                "unique list of prompt positions in [1, 262143].")
+        if data.get("stream"):
+            raise ValueError(
+                "BI100 sampled prompt logprobs require `stream=False`.")
+        if not isinstance(data.get("prompt_logprobs"), int) \
+                or data["prompt_logprobs"] <= 0:
+            raise ValueError(
+                "BI100 sampled prompt logprobs require positive "
+                "`prompt_logprobs`.")
+        return data
+
     @model_validator(mode="before")
     @classmethod
     def check_guided_decoding_count(cls, data):
@@ -533,8 +630,8 @@ class ChatCompletionRequest(OpenAIBaseModel):
             raise ValueError(
                 "You can only use one kind of guided decoding "
                 "('guided_json', 'guided_regex' or 'guided_choice').")
-        # you can only either use guided decoding or tools, not both
-        if guide_count > 1 and data.get("tool_choice",
+        # you can only either use guided decoding or a forced tool, not both
+        if guide_count > 0 and data.get("tool_choice",
                                         "none") not in ("none", "auto"):
             raise ValueError(
                 "You can only either use guided decoding or tools, not both.")
@@ -551,11 +648,7 @@ class ChatCompletionRequest(OpenAIBaseModel):
 
         # if "tool_choice" is specified -- validation
         if "tool_choice" in data:
-
-            # "none" means don't use any tools — valid per OpenAI spec,
-            # just strip tool_choice and let vLLM ignore tools.
             if data["tool_choice"] == "none":
-                del data["tool_choice"]
                 return data
 
             # ensure that if "tool choice" is specified, tools are present
@@ -564,12 +657,12 @@ class ChatCompletionRequest(OpenAIBaseModel):
                     "When using `tool_choice`, `tools` must be set.")
 
             # make sure that tool choice is either a named tool
-            # OR that it's set to "auto"
-            if data["tool_choice"] not in ("auto", "required", "none") \
-                    and not isinstance(data["tool_choice"], dict):
+            # OR that it's set to "auto"/"none"
+            if data["tool_choice"] != "auto" and not isinstance(
+                    data["tool_choice"], dict):
                 raise ValueError(
-                    "`tool_choice` must be a named tool, \"auto\", "
-                    "\"required\", or \"none\".")
+                    "`tool_choice` must be a named tool, \"auto\", or "
+                    "\"none\".")
 
             # ensure that if "tool_choice" is specified as an object,
             # it matches a valid tool
@@ -731,7 +824,8 @@ class CompletionRequest(OpenAIBaseModel):
         guided_json_from_schema = None
         if self.response_format is not None:
             if self.response_format.type == "json_object":
-                guided_json_object = True
+                # Keep CompletionRequest aligned with ChatCompletionRequest.
+                guided_json_from_schema = {"type": "object"}
             elif (self.response_format.type == "json_schema"
                   and self.response_format.json_schema is not None
                   and self.response_format.json_schema.json_schema is not None):
@@ -1081,6 +1175,7 @@ class TokenizeChatRequest(OpenAIBaseModel):
     add_generation_prompt: bool = Field(default=True)
     continue_final_message: bool = Field(default=False)
     add_special_tokens: bool = Field(default=False)
+    chat_template_kwargs: Optional[Dict[str, Any]] = Field(default=None)
 
     @model_validator(mode="before")
     @classmethod

qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py

@@ -1,6 +1,7 @@
 # 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
 
+import os
 from typing import Optional, List
 
 from ...configuration_utils import PretrainedConfig as PreTrainedConfig
@@ -18,22 +19,62 @@ def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
             f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
         )
 
+
+HYBRID_KV_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
+HYBRID_KV_ACCOUNTING_CONFIG = "bi100_hybrid_kv_accounting_mode"
+LEGACY_KV_ACCOUNTING = "legacy40"
+FULL_ATTENTION_KV_ACCOUNTING = "full_attention"
+
+
+def _hybrid_kv_accounting_mode(environ=None, serialized_mode=None):
+    source = os.environ if environ is None else environ
+    environment_mode = source.get(HYBRID_KV_ACCOUNTING_ENV)
+    if (environment_mode is not None and serialized_mode is not None
+            and environment_mode != serialized_mode):
+        raise RuntimeError(
+            f"{HYBRID_KV_ACCOUNTING_ENV}={environment_mode!r} conflicts "
+            f"with serialized {HYBRID_KV_ACCOUNTING_CONFIG}="
+            f"{serialized_mode!r}")
+    mode = environment_mode or serialized_mode or LEGACY_KV_ACCOUNTING
+    if mode not in (LEGACY_KV_ACCOUNTING, FULL_ATTENTION_KV_ACCOUNTING):
+        raise RuntimeError(
+            f"{HYBRID_KV_ACCOUNTING_ENV} must be "
+            f"'{LEGACY_KV_ACCOUNTING}' or "
+            f"'{FULL_ATTENTION_KV_ACCOUNTING}', got {mode!r}")
+    return mode
+
+
+def _vllm_layers_block_type(
+    layer_types,
+    environ=None,
+    serialized_mode=None,
+):
+    """Expose hybrid-layer ownership in the form vLLM 0.6.3 consumes."""
+    mode = _hybrid_kv_accounting_mode(environ, serialized_mode)
+    if mode == LEGACY_KV_ACCOUNTING:
+        return ["attention"] * len(layer_types)
+    return [
+        "attention" if layer_type == "full_attention" else layer_type
+        for layer_type in layer_types
+    ]
+
 try:
     from typing import TypedDict
+except ImportError:
+    RopeParameters = dict
+else:
     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.
+    Configuration for the text backbone of Qwen3.5 / Qwen3.6-35B-A3B models.
     model_type is "qwen3_5_text" (used internally by the nested config).
     """
 
@@ -143,7 +184,7 @@ class Qwen3_5VisionConfig(PreTrainedConfig):
 
 class Qwen3_5Config(PreTrainedConfig):
     r"""
-    Top-level configuration for Qwen3.5 / Qwen3.6-27B.
+    Top-level configuration for Qwen3.5 / Qwen3.6-35B-A3B.
     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.
@@ -163,6 +204,8 @@ class Qwen3_5Config(PreTrainedConfig):
         tie_word_embeddings=False,
         **kwargs,
     ):
+        serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None)
+        serialized_layers = kwargs.pop("layers_block_type", None)
         if isinstance(text_config, dict):
             self.text_config = Qwen3_5TextConfig(**text_config)
         elif text_config is None:
@@ -183,6 +226,17 @@ class Qwen3_5Config(PreTrainedConfig):
         self.vision_end_token_id = vision_end_token_id
         self.tie_word_embeddings = tie_word_embeddings
         super().__init__(**kwargs)
+        mode = _hybrid_kv_accounting_mode(
+            serialized_mode=serialized_mode)
+        layers_block_type = _vllm_layers_block_type(
+            self.text_config.layer_types, serialized_mode=mode)
+        if (serialized_layers is not None
+                and list(serialized_layers) != layers_block_type):
+            raise RuntimeError(
+                "serialized layers_block_type conflicts with "
+                f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}")
+        setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode)
+        self.layers_block_type = layers_block_type
 
 
 __all__ = ["Qwen3_5Config", "Qwen3_5TextConfig", "Qwen3_5VisionConfig"]

qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py

@@ -4,6 +4,7 @@
 # Removes ignore_keys_at_rope_validation / base_model_tp_plan / base_model_pp_plan
 # which are 5.x-only and irrelevant for vLLM inference.
 
+import os
 from typing import Optional
 
 from ...configuration_utils import PretrainedConfig as PreTrainedConfig
@@ -18,15 +19,55 @@ def layer_type_validation(layer_types, num_hidden_layers=None, attention=True):
             f"num_hidden_layers ({num_hidden_layers}) != len(layer_types) ({len(layer_types)})"
         )
 
+
+HYBRID_KV_ACCOUNTING_ENV = "BI100_HYBRID_KV_ACCOUNTING"
+HYBRID_KV_ACCOUNTING_CONFIG = "bi100_hybrid_kv_accounting_mode"
+LEGACY_KV_ACCOUNTING = "legacy40"
+FULL_ATTENTION_KV_ACCOUNTING = "full_attention"
+
+
+def _hybrid_kv_accounting_mode(environ=None, serialized_mode=None):
+    source = os.environ if environ is None else environ
+    environment_mode = source.get(HYBRID_KV_ACCOUNTING_ENV)
+    if (environment_mode is not None and serialized_mode is not None
+            and environment_mode != serialized_mode):
+        raise RuntimeError(
+            f"{HYBRID_KV_ACCOUNTING_ENV}={environment_mode!r} conflicts "
+            f"with serialized {HYBRID_KV_ACCOUNTING_CONFIG}="
+            f"{serialized_mode!r}")
+    mode = environment_mode or serialized_mode or LEGACY_KV_ACCOUNTING
+    if mode not in (LEGACY_KV_ACCOUNTING, FULL_ATTENTION_KV_ACCOUNTING):
+        raise RuntimeError(
+            f"{HYBRID_KV_ACCOUNTING_ENV} must be "
+            f"'{LEGACY_KV_ACCOUNTING}' or "
+            f"'{FULL_ATTENTION_KV_ACCOUNTING}', got {mode!r}")
+    return mode
+
+
+def _vllm_layers_block_type(
+    layer_types,
+    environ=None,
+    serialized_mode=None,
+):
+    """Expose hybrid-layer ownership in the form vLLM 0.6.3 consumes."""
+    mode = _hybrid_kv_accounting_mode(environ, serialized_mode)
+    if mode == LEGACY_KV_ACCOUNTING:
+        return ["attention"] * len(layer_types)
+    return [
+        "attention" if layer_type == "full_attention" else layer_type
+        for layer_type in layer_types
+    ]
+
 try:
     from typing import TypedDict
+except ImportError:
+    RopeParameters = dict
+else:
     class RopeParameters(TypedDict, total=False):
         rope_theta: float
         rope_type: str
         partial_rotary_factor: float
         factor: float
-except Exception:
-    RopeParameters = dict
 
 # --- End stubs ---
 
@@ -173,6 +214,8 @@ class Qwen3_5MoeConfig(PreTrainedConfig):
         tie_word_embeddings=False,
         **kwargs,
     ):
+        serialized_mode = kwargs.pop(HYBRID_KV_ACCOUNTING_CONFIG, None)
+        serialized_layers = kwargs.pop("layers_block_type", None)
         if isinstance(text_config, dict):
             self.text_config = Qwen3_5MoeTextConfig(**text_config)
         elif text_config is None:
@@ -193,6 +236,17 @@ class Qwen3_5MoeConfig(PreTrainedConfig):
         self.vision_end_token_id = vision_end_token_id
         self.tie_word_embeddings = tie_word_embeddings
         super().__init__(**kwargs)
+        mode = _hybrid_kv_accounting_mode(
+            serialized_mode=serialized_mode)
+        layers_block_type = _vllm_layers_block_type(
+            self.text_config.layer_types, serialized_mode=mode)
+        if (serialized_layers is not None
+                and list(serialized_layers) != layers_block_type):
+            raise RuntimeError(
+                "serialized layers_block_type conflicts with "
+                f"{HYBRID_KV_ACCOUNTING_CONFIG}={mode!r}")
+        setattr(self, HYBRID_KV_ACCOUNTING_CONFIG, mode)
+        self.layers_block_type = layers_block_type
 
 
 __all__ = ["Qwen3_5MoeConfig", "Qwen3_5MoeTextConfig"]

qwen3_6_scripts/qwen3coder_tool_parser.py

@@ -77,28 +77,6 @@ class Qwen3CoderToolParser(ToolParser):
         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]}"
@@ -192,11 +170,21 @@ class Qwen3CoderToolParser(ToolParser):
                 try:
                     return json.loads(param_value)
                 except (json.JSONDecodeError, TypeError, ValueError):
-                    pass
+                    logger.debug(
+                        "Could not JSON-decode parameter '%s' for tool '%s'; "
+                        "falling back to literal evaluation.",
+                        param_name,
+                        func_name,
+                        exc_info=True)
             try:
                 return ast.literal_eval(param_value)
             except (ValueError, SyntaxError, TypeError):
-                pass
+                logger.debug(
+                    "Could not literal-eval parameter '%s' for tool '%s'; "
+                    "returning string value.",
+                    param_name,
+                    func_name,
+                    exc_info=True)
             return param_value
 
     def _parse_xml_function_call(
@@ -464,8 +452,8 @@ class Qwen3CoderToolParser(ToolParser):
                 serialized = json.dumps(converted, ensure_ascii=False)
 
                 sep = "" if self.param_count == 0 else ", "
-                json_fragments.append(
-                    f'{sep}"{current_param_name}": {serialized}')
+                key = json.dumps(current_param_name, ensure_ascii=False)
+                json_fragments.append(f"{sep}{key}: {serialized}")
                 self.param_count += 1
 
             if json_fragments:

qwen3_6_scripts/reasoning/__init__.py

@@ -1,5 +1,5 @@
 """
-Reasoning parser module for vLLM 0.6.3 (BI-V100 / Qwen3.6-27B adaptation).
+Reasoning parser module for vLLM 0.6.3 (BI-V100 / Qwen3.6-35B-A3B adaptation).
 
 Usage: --reasoning-parser qwen3
 """

qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py

@@ -38,16 +38,18 @@ class Qwen3ReasoningParser(BaseThinkingReasoningParser):
         parts = model_output.partition(self.start_token)
         model_output = parts[2] if parts[1] else parts[0]
 
+        if not self.thinking_enabled:
+            if self.end_token in model_output:
+                _, _, content = model_output.partition(self.end_token)
+                return None, content or ""
+            return None, model_output
+
         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
+        return reasoning, content or None
 
     def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int:
         token_ids = list(token_ids)

qwen3_6_scripts/registry.py

@@ -1,456 +0,0 @@
-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/scheduler.py

@@ -17,6 +17,28 @@ from vllm.sequence import (Sequence, SequenceData, SequenceGroup,
                            SequenceStatus)
 from vllm.utils import Device, PyObjectCache
 
+try:
+    from vllm.gdn_prefix import (GdnPrefixKey, GdnPrefixStatePolicy,
+                                 cap_prefill_end_at_capture_boundary,
+                                 canonical_direct_segment_offsets,
+                                 capture_points_for_step,
+                                 final_capture_key,
+                                 gdn_cache_policy_from_env,
+                                 gdn_restore_alignment,
+                                 gdn_restore_mode_from_env,
+                                 keys_from_block_hashes,
+                                 restore_key_is_eligible,
+                                 strict_prefix_block_count)
+except ImportError:  # Local source-tree tests.
+    from qwen3_6_scripts.gdn_prefix import (
+        GdnPrefixKey, GdnPrefixStatePolicy,
+        cap_prefill_end_at_capture_boundary,
+        canonical_direct_segment_offsets, capture_points_for_step,
+        final_capture_key, gdn_cache_policy_from_env,
+        gdn_restore_alignment, gdn_restore_mode_from_env,
+        keys_from_block_hashes, restore_key_is_eligible,
+        strict_prefix_block_count)
+
 logger = init_logger(__name__)
 
 # Test-only. If configured, decode is preempted with
@@ -27,6 +49,47 @@ ARTIFICIAL_PREEMPTION_PROB = 0.5
 ARTIFICIAL_PREEMPTION_MAX_CNT = 500
 
 
+def _plan_gdn_prefix_fast_forward(
+        restore_key: Optional[GdnPrefixKey], num_computed_tokens: int,
+        prompt_len: int, nominal_chunk_size: int,
+        remaining_token_budget: int, block_size: int,
+        logical_chunk_alignment: Optional[int] = None) -> Tuple[int, int]:
+    """Return logical progress and physical query tokens for a direct hit.
+
+    The scheduler normally uses one value for both quantities. A GDN prefix
+    state makes it safe to advance over a much larger logical prefix while
+    sending only the suffix after that checkpoint to the model runner.
+    """
+    fallback = (nominal_chunk_size, nominal_chunk_size)
+    if (restore_key is None or num_computed_tokens != 0 or prompt_len <= 0
+            or nominal_chunk_size <= 0 or remaining_token_budget <= 0
+            or block_size <= 0):
+        return fallback
+
+    checkpoint_tokens = restore_key[0] * block_size
+    logical_limit = checkpoint_tokens + remaining_token_budget
+    if logical_chunk_alignment is not None:
+        if (logical_chunk_alignment <= 0
+                or logical_chunk_alignment % block_size != 0):
+            raise ValueError("logical_chunk_alignment must be a positive "
+                             "multiple of block_size")
+        next_boundary = (
+            checkpoint_tokens // logical_chunk_alignment + 1
+        ) * logical_chunk_alignment
+        logical_limit = min(logical_limit, next_boundary)
+
+    logical_chunk_size = min(prompt_len, logical_limit)
+    physical_query_tokens = logical_chunk_size - checkpoint_tokens
+    if (physical_query_tokens <= 0
+            or physical_query_tokens > remaining_token_budget):
+        return fallback
+    if (logical_chunk_size <= nominal_chunk_size
+            and (logical_chunk_alignment is None
+                 or physical_query_tokens >= nominal_chunk_size)):
+        return fallback
+    return logical_chunk_size, physical_query_tokens
+
+
 class PreemptionMode(enum.Enum):
     """Preemption modes.
 
@@ -56,6 +119,9 @@ class SchedulingBudget:
     _request_ids_num_batched_tokens: Set[str] = field(default_factory=set)
     _request_ids_num_curr_seqs: Set[str] = field(default_factory=set)
     _num_batched_tokens: int = 0
+    _num_scheduled_tokens: int = 0
+    _request_num_scheduled_tokens: Dict[str, int] = field(
+        default_factory=dict)
     _num_curr_seqs: int = 0
 
     def can_schedule(self, *, num_new_tokens: int, num_new_seqs: int):
@@ -67,18 +133,26 @@ class SchedulingBudget:
     def remaining_token_budget(self):
         return self.token_budget - self.num_batched_tokens
 
-    def add_num_batched_tokens(self, req_id: str, num_batched_tokens: int):
+    def add_num_batched_tokens(
+            self, req_id: str, num_batched_tokens: int,
+            num_scheduled_tokens: Optional[int] = None):
         if req_id in self._request_ids_num_batched_tokens:
             return
 
+        if num_scheduled_tokens is None:
+            num_scheduled_tokens = num_batched_tokens
         self._request_ids_num_batched_tokens.add(req_id)
         self._num_batched_tokens += num_batched_tokens
+        self._num_scheduled_tokens += num_scheduled_tokens
+        self._request_num_scheduled_tokens[req_id] = num_scheduled_tokens
 
     def subtract_num_batched_tokens(self, req_id: str,
                                     num_batched_tokens: int):
         if req_id in self._request_ids_num_batched_tokens:
             self._request_ids_num_batched_tokens.remove(req_id)
             self._num_batched_tokens -= num_batched_tokens
+            self._num_scheduled_tokens -= (
+                self._request_num_scheduled_tokens.pop(req_id))
 
     def add_num_seqs(self, req_id: str, num_curr_seqs: int):
         if req_id in self._request_ids_num_curr_seqs:
@@ -96,6 +170,10 @@ class SchedulingBudget:
     def num_batched_tokens(self):
         return self._num_batched_tokens
 
+    @property
+    def num_scheduled_tokens(self):
+        return self._num_scheduled_tokens
+
     @property
     def num_curr_seqs(self):
         return self._num_curr_seqs
@@ -135,7 +213,8 @@ class SchedulerOutputs:
     preempted: int
 
     def __post_init__(self):
-        # Swap in and swap out should never happen at the same time.
+        # Request-level preemption cannot swap both ways in one step. The
+        # content-addressed CPU tier appends its ordered maps after creation.
         assert not (self.blocks_to_swap_in and self.blocks_to_swap_out)
 
         self.num_loras: int = len(self.lora_requests)
@@ -351,6 +430,19 @@ class Scheduler:
         # can and must be released after the current step.
         # This is used to evict the finished requests from the Mamba cache.
         self._finished_requests_ids: List[str] = list()
+        self._gdn_prefix_policy = GdnPrefixStatePolicy(
+            gdn_cache_policy_from_env())
+        self._gdn_restore_mode = gdn_restore_mode_from_env()
+        try:
+            self._gdn_replay_alignment = gdn_restore_alignment(
+                self._gdn_restore_mode, self.cache_config.block_size,
+                scheduler_config.max_num_batched_tokens)
+        except ValueError as exc:
+            raise RuntimeError(str(exc)) from exc
+        self._gdn_request_restore_keys: Dict[
+            str, Optional[GdnPrefixKey]] = {}
+        self._gdn_request_capture_targets: Dict[
+            str, Tuple[GdnPrefixKey, ...]] = {}
         # Time at previous scheduling step
         self.prev_time = 0.0
         # Did we schedule a prompt at previous step?
@@ -423,6 +515,46 @@ class Scheduler:
         # Only for testing purposes.
         self.swapped.append(seq_group)
 
+    def _cap_gdn_capture_boundary(
+            self, seq_group: SequenceGroup, token_chunk_size: int,
+            physical_query_tokens: int) -> Tuple[int, int]:
+        """Align admission64 capture state with a physical model forward."""
+        targets = self._gdn_request_capture_targets.get(
+            seq_group.request_id, ())
+        if (self._gdn_prefix_policy.policy != "admission64"
+                or not targets):
+            return token_chunk_size, physical_query_tokens
+        if token_chunk_size <= 0 or physical_query_tokens <= 0:
+            raise RuntimeError("GDN prefill token counts must be positive")
+
+        seqs = seq_group.get_seqs()
+        if len(seqs) != 1 or not seq_group.is_prefill():
+            raise RuntimeError(
+                "GDN capture boundary requires one prefill sequence")
+        num_computed_tokens = seqs[0].data.get_num_computed_tokens()
+        logical_end_tokens = num_computed_tokens + token_chunk_size
+        logical_start_tokens = logical_end_tokens - physical_query_tokens
+        if logical_start_tokens < num_computed_tokens:
+            raise RuntimeError(
+                "GDN physical query starts before scheduler progress")
+
+        capped_end_tokens = cap_prefill_end_at_capture_boundary(
+            logical_start_tokens, logical_end_tokens, targets,
+            self.cache_config.block_size)
+        if capped_end_tokens == logical_end_tokens:
+            return token_chunk_size, physical_query_tokens
+
+        capped_chunk_size = capped_end_tokens - num_computed_tokens
+        capped_query_tokens = capped_end_tokens - logical_start_tokens
+        if capped_chunk_size <= 0 or capped_query_tokens <= 0:
+            raise RuntimeError("GDN capture boundary produced an empty step")
+        logger.info(
+            "[BI100 GDN CAPTURE BOUNDARY] request=%s logical_start=%d "
+            "logical_end=%d capped_end=%d physical_query_tokens=%d",
+            seq_group.request_id, logical_start_tokens, logical_end_tokens,
+            capped_end_tokens, capped_query_tokens)
+        return capped_chunk_size, capped_query_tokens
+
     def abort_seq_group(self, request_id: Union[str, Iterable[str]]) -> None:
         """Aborts a sequence group with the given ID.
 
@@ -469,10 +601,16 @@ class Scheduler:
     ) -> None:
         """
         Free a sequence group from a cross-attention block table.
-        Has no effect on decoder-only models.
+        Also release any request-local multimodal cache namespace.
         """
-        if seq_group.is_encoder_decoder():
-            self.block_manager.free_cross(seq_group)
+        try:
+            if seq_group.is_encoder_decoder():
+                self.block_manager.free_cross(seq_group)
+        finally:
+            release_namespace = getattr(
+                self.block_manager, "release_request_cache_namespace", None)
+            if release_namespace is not None:
+                release_namespace(seq_group.request_id)
 
     def has_unfinished_seqs(self) -> bool:
         return len(self.waiting) != 0 or len(self.running) != 0 or len(
@@ -548,6 +686,9 @@ class Scheduler:
             if num_running_tokens == 0:
                 # No budget => Stop
                 break
+            if enable_chunking and seq_group.is_prefill():
+                num_running_tokens, _ = self._cap_gdn_capture_boundary(
+                    seq_group, num_running_tokens, num_running_tokens)
 
             running_queue.popleft()
 
@@ -949,6 +1090,83 @@ class Scheduler:
             waiting_queue.popleft()
             self._allocate_and_set_running(seq_group)
 
+            budget_token_count = num_new_tokens
+            if (enable_chunking
+                    and self.cache_config.enable_prefix_caching
+                    and len(waiting_seqs) == 1):
+                prompt_seq = waiting_seqs[0]
+                computed_block_nums = list(
+                    self.block_manager.get_common_computed_block_ids(
+                        waiting_seqs))
+                block_hashes = self.block_manager.get_content_hashes(prompt_seq)
+                max_live_blocks = min(
+                    len(computed_block_nums), len(block_hashes),
+                    strict_prefix_block_count(
+                        prompt_seq.data.get_len(),
+                        self.cache_config.block_size))
+                live_keys = keys_from_block_hashes(
+                    block_hashes[:max_live_blocks])
+                direct_final_key = final_capture_key(
+                    block_hashes, prompt_seq.data.get_len(),
+                    self.cache_config.block_size, "direct",
+                    self.cache_config.block_size)
+                live_keys = [
+                    key for key in live_keys
+                    if restore_key_is_eligible(
+                        key, prompt_seq.data.get_len(),
+                        self.cache_config.block_size,
+                        self._gdn_restore_mode,
+                        self._gdn_replay_alignment,
+                        direct_final_key=(
+                            direct_final_key
+                            if self._gdn_restore_mode == "hybrid64" else None))
+                ]
+                restore_key = self._gdn_prefix_policy.select_restore(
+                    live_keys, len(live_keys))
+                self._gdn_request_restore_keys[
+                    seq_group.request_id] = restore_key
+
+                capture_targets = []
+                branch_key = self._gdn_prefix_policy.repeated_branch_candidate(
+                    live_keys, len(live_keys))
+                if branch_key is not None:
+                    capture_targets.append(branch_key)
+                final_key = final_capture_key(
+                    block_hashes, prompt_seq.data.get_len(),
+                    self.cache_config.block_size, self._gdn_restore_mode,
+                    self._gdn_replay_alignment)
+                if (final_key is not None
+                        and final_key not in capture_targets
+                        and self._gdn_prefix_policy.should_capture_final(
+                            final_key)):
+                    capture_targets.append(final_key)
+                self._gdn_request_capture_targets[seq_group.request_id] = tuple(
+                    capture_targets)
+
+                num_new_tokens, budget_token_count = (
+                    _plan_gdn_prefix_fast_forward(
+                        restore_key,
+                        prompt_seq.data.get_num_computed_tokens(),
+                        prompt_seq.data.get_len(),
+                        num_new_tokens,
+                        budget.remaining_token_budget(),
+                        self.cache_config.block_size,
+                        logical_chunk_alignment=(
+                            self.scheduler_config.max_num_batched_tokens
+                            if self._gdn_restore_mode == "hybrid64" else None)))
+                if budget_token_count != num_new_tokens:
+                    logger.info(
+                        "[BI100 GDN FAST-FORWARD] request=%s "
+                        "checkpoint_tokens=%d logical_tokens=%d "
+                        "physical_query_tokens=%d",
+                        seq_group.request_id,
+                        num_new_tokens - budget_token_count,
+                        num_new_tokens,
+                        budget_token_count)
+                num_new_tokens, budget_token_count = (
+                    self._cap_gdn_capture_boundary(
+                        seq_group, num_new_tokens, budget_token_count))
+
             if enable_chunking and self.scheduler_config.is_multi_step:
                 blocks_to_copy: List[Tuple[int, int]] = []
                 # init_multi_step_from_lookahead_slots happens in append_slots
@@ -969,7 +1187,10 @@ class Scheduler:
             seq_groups.append(
                 ScheduledSequenceGroup(seq_group=seq_group,
                                        token_chunk_size=num_new_tokens))
-            budget.add_num_batched_tokens(seq_group.request_id, num_new_tokens)
+            budget.add_num_batched_tokens(
+                seq_group.request_id,
+                budget_token_count,
+                num_scheduled_tokens=num_new_tokens)
             budget.add_num_seqs(seq_group.request_id, num_new_seqs)
 
         # Queue requests that couldn't be scheduled.
@@ -1077,7 +1298,7 @@ class Scheduler:
         return SchedulerOutputs(
             scheduled_seq_groups=scheduled_seq_groups,
             num_prefill_groups=num_prefill_groups,
-            num_batched_tokens=budget.num_batched_tokens,
+            num_batched_tokens=budget.num_scheduled_tokens,
             blocks_to_swap_in=swapped_in.blocks_to_swap_in,
             blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
             blocks_to_copy=blocks_to_copy,
@@ -1158,7 +1379,7 @@ class Scheduler:
             num_prefill_groups=(len(prefills.seq_groups) +
                                 len(swapped_in.prefill_seq_groups) +
                                 len(running_scheduled.prefill_seq_groups)),
-            num_batched_tokens=budget.num_batched_tokens,
+            num_batched_tokens=budget.num_scheduled_tokens,
             blocks_to_swap_in=swapped_in.blocks_to_swap_in,
             blocks_to_swap_out=running_scheduled.blocks_to_swap_out,
             blocks_to_copy=running_scheduled.blocks_to_copy +
@@ -1217,7 +1438,25 @@ class Scheduler:
         # such as self.running, self.swapped, and self.waiting.
         scheduler_start_time = time.perf_counter()
 
+        begin_prefix_cache_step = getattr(
+            self.block_manager, "begin_prefix_cache_step", None)
+        if callable(begin_prefix_cache_step):
+            begin_prefix_cache_step()
         scheduler_outputs: SchedulerOutputs = self._schedule()
+        drain_prefix_swaps = getattr(
+            self.block_manager, "get_and_reset_prefix_swaps", None)
+        if callable(drain_prefix_swaps):
+            prefix_swap_in, prefix_swap_out = drain_prefix_swaps()
+            if prefix_swap_in or prefix_swap_out:
+                if (scheduler_outputs.blocks_to_swap_in
+                        or scheduler_outputs.blocks_to_swap_out):
+                    raise RuntimeError(
+                        "content-addressed CPU KV transfers cannot share a "
+                        "scheduler step with request-level preemption swap")
+                # Both directions are valid for this tier: a GPU victim is
+                # preserved before that same physical slot is reused by H2D.
+                scheduler_outputs.blocks_to_swap_in.extend(prefix_swap_in)
+                scheduler_outputs.blocks_to_swap_out.extend(prefix_swap_out)
         now = time.time()
 
         if not self.cache_config.enable_prefix_caching:
@@ -1262,28 +1501,111 @@ class Scheduler:
                 block_tables[seq_id] = self.block_manager.get_block_table(seq)
                 self.block_manager.access_all_blocks_in_seq(seq, now)
 
+            common_computed_block_nums = []
             if self.cache_config.enable_prefix_caching:
-                common_computed_block_nums = (
+                raw_computed_block_nums = list(
                     self.block_manager.get_common_computed_block_ids(
                         seq_group.get_seqs(status=SequenceStatus.RUNNING)))
+                if not seq_group.is_prefill():
+                    common_computed_block_nums = raw_computed_block_nums
 
             do_sample = True
             is_prompt = seq_group.is_prefill()
             # We should send the metadata to workers when the first prefill
             # is sent. Subsequent requests could be chunked prefill or decode.
             is_first_prefill = False
+            gdn_restore_key = None
+            gdn_capture_points = None
+            gdn_evict_keys = None
+            gdn_segment_offsets = None
             if is_prompt:
+                gdn_capture_points = []
+                gdn_evict_keys = []
+                gdn_segment_offsets = []
                 seqs = seq_group.get_seqs()
                 # Prefill has only 1 sequence.
                 assert len(seqs) == 1
                 num_computed_tokens = seqs[0].data.get_num_computed_tokens()
                 is_first_prefill = num_computed_tokens == 0
-                if (is_first_prefill
-                        and self.cache_config.enable_prefix_caching
-                        and seq_group.metrics is not None):
-                    seq_group.metrics.num_cached_tokens = (
-                        len(common_computed_block_nums)
-                        * self.cache_config.block_size)
+                logical_end_tokens = min(
+                    seqs[0].data.get_len(),
+                    num_computed_tokens + token_chunk_size)
+                if self.cache_config.enable_prefix_caching:
+                    restore_key = self._gdn_request_restore_keys.get(
+                        seq_group.request_id)
+                    if is_first_prefill and restore_key is not None:
+                        gdn_restore_key = restore_key
+                        common_computed_block_nums = raw_computed_block_nums[
+                            :restore_key[0]]
+                        if len(common_computed_block_nums) != restore_key[0]:
+                            raise RuntimeError(
+                                "GDN restore key exceeds the live KV prefix")
+                    else:
+                        # Once this request has started, the request-local Mamba
+                        # state is authoritative. Never let a longer raw KV hit
+                        # skip ahead without a matching recurrent state.
+                        max_context_blocks = (num_computed_tokens
+                                              // self.cache_config.block_size)
+                        common_computed_block_nums = raw_computed_block_nums[
+                            :max_context_blocks]
+
+                    restore_tokens = (
+                        restore_key[0] * self.cache_config.block_size
+                        if is_first_prefill and restore_key is not None else 0)
+                    if seq_group.metrics is not None and restore_tokens:
+                        seq_group.metrics.num_cached_tokens = max(
+                            seq_group.metrics.num_cached_tokens or 0,
+                            restore_tokens)
+
+                    capture_targets = list(
+                        self._gdn_request_capture_targets.get(
+                            seq_group.request_id, ()))
+                    if self._gdn_prefix_policy.policy == "fine32":
+                        step_key = final_capture_key(
+                            self.block_manager.get_content_hashes(seqs[0]),
+                            logical_end_tokens, self.cache_config.block_size,
+                            self._gdn_restore_mode,
+                            self._gdn_replay_alignment)
+                        capture_targets = ([step_key]
+                                           if step_key is not None else [])
+                    if self._gdn_prefix_policy.policy != "off":
+                        physical_context_tokens = (
+                            restore_tokens if is_first_prefill
+                            else num_computed_tokens)
+                        if (self._gdn_restore_mode == "hybrid64"
+                                and self._gdn_prefix_policy.policy
+                                == "admission64"):
+                            gdn_segment_offsets = list(
+                                canonical_direct_segment_offsets(
+                                    self.block_manager.get_content_hashes(
+                                        seqs[0]),
+                                    physical_context_tokens,
+                                    logical_end_tokens,
+                                    self.cache_config.block_size,
+                                    self.scheduler_config.max_num_batched_tokens))
+                        gdn_capture_points = list(capture_points_for_step(
+                            capture_targets, physical_context_tokens,
+                            logical_end_tokens, self.cache_config.block_size))
+                        gdn_evict_keys = list(
+                            self._gdn_prefix_policy.admit(
+                                key for _, key in gdn_capture_points))
+                    trace_update = getattr(
+                        self.block_manager, "_bi100_update_cache_trace", None)
+                    if callable(trace_update):
+                        capture_actions = []
+                        for _, key in gdn_capture_points:
+                            if self._gdn_prefix_policy.policy == "fine32":
+                                reason = "fine32_chunk"
+                            elif (capture_targets
+                                  and key == capture_targets[-1]):
+                                reason = "final_prefill"
+                            else:
+                                reason = "repeated_branch"
+                            capture_actions.append((key, reason))
+                        trace_update(
+                            seqs[0], len(raw_computed_block_nums),
+                            gdn_restore_key, capture_actions,
+                            gdn_evict_keys, self._gdn_prefix_policy.policy)
                 # In the next iteration, all prompt tokens are not computed.
                 # It means the prefill is chunked, and we don't need sampling.
                 # NOTE: We use get_len instead of get_prompt_len because when
@@ -1293,6 +1615,12 @@ class Scheduler:
                         seqs[0].data.get_len()):
                     do_sample = False
 
+                if logical_end_tokens >= seqs[0].data.get_len():
+                    self._gdn_request_restore_keys.pop(seq_group.request_id,
+                                                       None)
+                    self._gdn_request_capture_targets.pop(seq_group.request_id,
+                                                          None)
+
             # It assumes the scheduled_seq_groups is ordered by
             # prefill < decoding.
             if is_first_prefill or not self.scheduler_config.send_delta_data:
@@ -1318,6 +1646,10 @@ class Scheduler:
                     if scheduler_outputs.num_prefill_groups > 0 else None,
                     mm_processor_kwargs=seq_group.mm_processor_kwargs,
                     prompt_adapter_request=seq_group.prompt_adapter_request,
+                    gdn_restore_key=gdn_restore_key,
+                    gdn_capture_points=gdn_capture_points,
+                    gdn_evict_keys=gdn_evict_keys,
+                    gdn_segment_offsets=gdn_segment_offsets,
                 )
             else:
                 # When SPMD mode is enabled, we only send delta data except for
@@ -1333,6 +1665,10 @@ class Scheduler:
                     do_sample=do_sample,
                     token_chunk_size=token_chunk_size,
                     computed_block_nums=common_computed_block_nums,
+                    gdn_restore_key=gdn_restore_key,
+                    gdn_capture_points=gdn_capture_points,
+                    gdn_evict_keys=gdn_evict_keys,
+                    gdn_segment_offsets=gdn_segment_offsets,
                 )
             seq_group_metadata_list.append(seq_group_metadata)
 
@@ -1387,6 +1723,9 @@ class Scheduler:
             # Free cross-attention block table, if it exists
             self._free_seq_group_cross_attn_blocks(seq_group)
 
+            self._gdn_request_restore_keys.pop(seq_group.request_id, None)
+            self._gdn_request_capture_targets.pop(seq_group.request_id, None)
+
             # Add the finished requests to the finished requests list.
             # This list will be used to update the Mamba cache in the
             # next step.

qwen3_6_scripts/sequence.py

@@ -941,6 +941,12 @@ class SequenceGroupMetadataDelta(
     computed_block_nums: Optional[List[int]] = None
     state: Optional[SequenceGroupState] = msgspec.field(
         default_factory=lambda: SequenceGroupState())
+    # BI100 hybrid prefix-cache actions. Fields are appended for msgspec wire
+    # compatibility with the pre-existing array-like structure.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
 
 class SequenceGroupMetadata(
@@ -1006,6 +1012,12 @@ class SequenceGroupMetadata(
     # Zero means speculative decoding is disabled for some reasons.
     # TODO: We should maintain this states out of the sequence group.
     num_speculative_tokens: Optional[int] = None
+    # BI100 hybrid prefix-cache actions. These are internal scheduler-to-worker
+    # metadata and never surface through the OpenAI API.
+    gdn_restore_key: Optional[Tuple[int, bytes]] = None
+    gdn_capture_points: Optional[List[Tuple[int, Tuple[int, bytes]]]] = None
+    gdn_evict_keys: Optional[List[Tuple[int, bytes]]] = None
+    gdn_segment_offsets: Optional[List[int]] = None
 
     def __post_init__(self):
         if self.seq_data is not None and self.token_chunk_size is None:
@@ -1052,6 +1064,14 @@ class SequenceGroupMetadata(
         self.token_chunk_size = sequence_group_metadata_delta.token_chunk_size
         self.do_sample = sequence_group_metadata_delta.do_sample
         self.is_prompt = sequence_group_metadata_delta.is_prompt
+        self.computed_block_nums = (
+            sequence_group_metadata_delta.computed_block_nums)
+        self.gdn_restore_key = sequence_group_metadata_delta.gdn_restore_key
+        self.gdn_capture_points = (
+            sequence_group_metadata_delta.gdn_capture_points)
+        self.gdn_evict_keys = sequence_group_metadata_delta.gdn_evict_keys
+        self.gdn_segment_offsets = (
+            sequence_group_metadata_delta.gdn_segment_offsets)
 
     def finish_step(self) -> None:
         assert self.state is not None

qwen3_6_scripts/serving_chat.py

@@ -45,6 +45,157 @@ from vllm.utils import iterate_with_cancellation, random_uuid
 logger = init_logger(__name__)
 
 
+def _serialize_tool_arguments(arguments) -> str:
+    if arguments is None:
+        return "{}"
+    if isinstance(arguments, str):
+        return arguments
+    if isinstance(arguments, (dict, list)):
+        return json.dumps(arguments, ensure_ascii=False)
+    return json.dumps(arguments, ensure_ascii=False)
+
+
+def _tool_arguments_are_json_object(arguments: str) -> bool:
+    try:
+        value = json.loads(arguments)
+    except (json.JSONDecodeError, TypeError, ValueError):
+        return False
+    return isinstance(value, dict)
+
+
+def _reclassify_named_guided_json(
+    reasoning_text: Optional[str],
+    output_text: str,
+) -> tuple[Optional[str], str]:
+    """Recover guided JSON misclassified as unterminated reasoning."""
+    if (not output_text and reasoning_text is not None
+            and _tool_arguments_are_json_object(reasoning_text)):
+        return None, reasoning_text
+    return reasoning_text, output_text
+
+
+def _select_named_tool_arguments(
+    output_text: str,
+    expected_name: str,
+    parsed_tool_calls: Optional[List[ToolCall]],
+) -> str:
+    """Use parser output only to repair a malformed named-tool payload."""
+    if _tool_arguments_are_json_object(output_text):
+        return output_text
+    if not parsed_tool_calls or len(parsed_tool_calls) != 1:
+        return output_text
+    call = parsed_tool_calls[0]
+    function = getattr(call, "function", None)
+    if function is None or getattr(function, "name", None) != expected_name:
+        return output_text
+    arguments = _serialize_tool_arguments(
+        getattr(function, "arguments", None))
+    if not _tool_arguments_are_json_object(arguments):
+        return output_text
+    return arguments
+
+
+def _named_tool_delta_payload(name: str, arguments: str, index: int,
+                              call_id: str, first_delta: bool
+                              ) -> Dict[str, object]:
+    function: Dict[str, object] = {"arguments": arguments}
+    payload: Dict[str, object] = {"index": index, "function": function}
+    if first_delta:
+        function["name"] = name
+        payload["id"] = call_id
+        payload["type"] = "function"
+    return payload
+
+
+def _consume_named_tool_header_slot(header_sent: List[bool],
+                                    index: int) -> bool:
+    first_delta = not header_sent[index]
+    header_sent[index] = True
+    return first_delta
+
+
+def _sequential_greedy_fanout_count(
+    request: ChatCompletionRequest,
+    max_num_seqs: int,
+) -> int:
+    """Return the supported deterministic fan-out width, or zero."""
+    n = request.n if request.n is not None else 1
+    if (
+        max_num_seqs == 1
+        and n == 2
+        and request.temperature == 0
+        and not request.stream
+        and not request.use_beam_search
+        and request.best_of is None
+        and request.prompt_logprobs is None
+    ):
+        return n
+    return 0
+
+
+def _merge_sequential_chat_responses(
+    responses: List[ChatCompletionResponse],
+    request_id: str,
+    created_time: int,
+) -> ChatCompletionResponse:
+    if len(responses) != 2:
+        raise ValueError("deterministic fan-out requires exactly two responses")
+
+    first = responses[0]
+    if any(response.model != first.model for response in responses):
+        raise ValueError("fan-out response models differ")
+    if any(len(response.choices) != 1 for response in responses):
+        raise ValueError("fan-out child response must contain one choice")
+    if any(
+        response.usage.prompt_tokens != first.usage.prompt_tokens
+        for response in responses
+    ):
+        raise ValueError("fan-out prompt token counts differ")
+    if any(
+        response.usage.completion_tokens is None for response in responses
+    ):
+        raise ValueError("fan-out completion token count is missing")
+
+    choices = [
+        response.choices[0].model_copy(
+            deep=True,
+            update={"index": index},
+        )
+        for index, response in enumerate(responses)
+    ]
+    completion_tokens = sum(
+        response.usage.completion_tokens or 0 for response in responses
+    )
+    reasoning_counts = [
+        response.usage.reasoning_tokens for response in responses
+    ]
+    reasoning_tokens = (
+        None
+        if all(value is None for value in reasoning_counts)
+        else sum(value or 0 for value in reasoning_counts)
+    )
+    prompt_details = first.usage.prompt_tokens_details
+    usage = UsageInfo(
+        prompt_tokens=first.usage.prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=first.usage.prompt_tokens + completion_tokens,
+        reasoning_tokens=reasoning_tokens,
+        prompt_tokens_details=(
+            prompt_details.model_copy(deep=True)
+            if prompt_details is not None
+            else None
+        ),
+    )
+    return ChatCompletionResponse(
+        id=request_id,
+        created=created_time,
+        model=first.model,
+        choices=choices,
+        usage=usage,
+        prompt_logprobs=first.prompt_logprobs,
+    )
+
+
 class OpenAIServingChat(OpenAIServing):
 
     def __init__(self,
@@ -118,6 +269,10 @@ class OpenAIServingChat(OpenAIServing):
         ChatCompletion API.
 
         """
+        if not request.messages:
+            return self.create_error_response(
+                "messages must contain at least one message")
+
         error_check_ret = await self._check_model(request)
         if error_check_ret is not None:
             logger.error("Error with model %s", error_check_ret)
@@ -129,6 +284,23 @@ class OpenAIServingChat(OpenAIServing):
         if self.engine_client.errored:
             raise self.engine_client.dead_error
 
+        # The fixed competition command uses max_num_seqs=1. Native vLLM
+        # cannot schedule n=2 in that configuration and also rejects greedy
+        # n>1. Two greedy choices are identical by definition, so execute two
+        # isolated n=1 requests and merge only this exact deterministic shape.
+        if request.n is not None and request.n > 1:
+            scheduler_config = await self.engine_client.get_scheduler_config()
+            max_num_seqs = scheduler_config.max_num_seqs
+            if request.n > max_num_seqs:
+                fanout_count = _sequential_greedy_fanout_count(
+                    request, max_num_seqs)
+                if fanout_count:
+                    return await self._create_sequential_greedy_fanout(
+                        request, raw_request, fanout_count)
+                return self.create_error_response(
+                    f"n={request.n} exceeds max_num_seqs={max_num_seqs}. "
+                    f"Use n<={max_num_seqs} or omit n.")
+
         try:
             (
                 lora_request,
@@ -179,27 +351,11 @@ class OpenAIServingChat(OpenAIServing):
             logger.exception("Error in loading multi-modal data")
             return self.create_error_response(str(e))
 
-        # n > max_num_seqs deadlock guard: scheduler uses break (not continue)
-        # when can_schedule(num_new_seqs=n) fails, so an n that exceeds
-        # max_num_seqs permanently blocks the entire waiting queue with no error.
-        # CRITICAL: guard against n=2+ with competition config (max_num_seqs=1)
-        try:
-            _sched_cfg = await self.engine_client.get_scheduler_config()
-            _max_seqs = _sched_cfg.max_num_seqs
-        except Exception:
-            _max_seqs = 1  # BI-V100 safety: default to 1 if config unavailable
-        if request.n is not None and request.n > _max_seqs:
-            # Clamp n to max_seqs instead of rejecting — this way t2_n_2
-            # returns 200 with fewer choices instead of crashing the service.
-            logger.warning(
-                "n=%d exceeds max_num_seqs=%d, clamping to %d",
-                request.n, _max_seqs, _max_seqs)
-            request.n = _max_seqs
-
         # validation for OpenAI tools
-        # tool_choice = "required" → treat as "auto" for compatibility
+        # tool_choice = "required" is not supported
         if request.tool_choice == "required":
-            request.tool_choice = "auto"
+            return self.create_error_response(
+                "tool_choice = \"required\" is not supported!")
 
         if not is_mistral_tokenizer and request.tool_choice == "auto" and not (
                 self.enable_auto_tools and self.tool_parser is not None):
@@ -310,6 +466,58 @@ class OpenAIServingChat(OpenAIServing):
             # TODO: Use a vllm-specific Validation Error
             return self.create_error_response(str(e))
 
+    async def _create_sequential_greedy_fanout(
+        self,
+        request: ChatCompletionRequest,
+        raw_request: Optional[Request],
+        fanout_count: int,
+    ) -> Union[ChatCompletionResponse, ErrorResponse]:
+        request_id = f"chat-{random_uuid()}"
+        created_time = int(time.time())
+        responses: List[ChatCompletionResponse] = []
+
+        for _ in range(fanout_count):
+            child_request = request.model_copy(
+                deep=True,
+                update={"n": 1},
+            )
+            child_response = await self.create_chat_completion(
+                child_request, raw_request)
+            if isinstance(child_response, ErrorResponse):
+                return child_response
+            if not isinstance(child_response, ChatCompletionResponse):
+                logger.error(
+                    "Sequential greedy fan-out unexpectedly returned a stream")
+                return self.create_error_response(
+                    "Failed to aggregate deterministic n=2 completion")
+            responses.append(child_response)
+
+        try:
+            response = _merge_sequential_chat_responses(
+                responses,
+                request_id,
+                created_time,
+            )
+        except ValueError as error:
+            logger.error(
+                "Sequential greedy fan-out aggregation failed: %s",
+                type(error).__name__,
+            )
+            return self.create_error_response(
+                "Failed to aggregate deterministic n=2 completion")
+
+        if raw_request is not None:
+            metadata = RequestResponseMetadata(
+                request_id=request_id,
+                final_usage_info=response.usage,
+            )
+            raw_request.state.request_metadata = metadata
+        logger.info(
+            "[BI100 N_FANOUT] choices=%d mode=sequential_greedy",
+            fanout_count,
+        )
+        return response
+
     def get_chat_request_role(self, request: ChatCompletionRequest) -> str:
         if request.add_generation_prompt:
             return self.response_role
@@ -341,6 +549,10 @@ class OpenAIServingChat(OpenAIServing):
             tool_choice_function_name = request.tool_choice.function.name
         else:
             tool_choice_function_name = None
+        named_tool_call_ids = (
+            [f"chatcmpl-tool-{random_uuid()}" for _ in range(num_choices)]
+            if tool_choice_function_name else [])
+        named_tool_header_sent = [False] * num_choices
 
         # Determine whether tools are in use with "auto" tool choice
         tool_choice_auto = (
@@ -354,7 +566,7 @@ class OpenAIServingChat(OpenAIServing):
         # parsing and reasoning parsing (both require full-history context).
         if tool_choice_auto or use_reasoning:
             previous_texts = [""] * num_choices
-            all_previous_token_ids = [[] for _ in range(num_choices)]
+            all_previous_token_ids = [[]] * num_choices
         else:
             previous_texts, all_previous_token_ids = None, None
 
@@ -363,8 +575,7 @@ class OpenAIServingChat(OpenAIServing):
             if tool_choice_auto and self.tool_parser:
                 tool_parsers: List[Optional[ToolParser]] = [
                     self.tool_parser(tokenizer)
-                    for _ in range(num_choices)
-                ]
+                ] * num_choices
             else:
                 tool_parsers = [None] * num_choices
         except RuntimeError as e:
@@ -557,11 +768,16 @@ class OpenAIServingChat(OpenAIServing):
 
                     # handle streaming deltas for tools with named tool_choice
                     if tool_choice_function_name:
+                        first_named_delta = _consume_named_tool_header_slot(
+                            named_tool_header_sent, i)
                         delta_message = DeltaMessage(tool_calls=[
-                            DeltaToolCall(function=DeltaFunctionCall(
-                                name=tool_choice_function_name,
-                                arguments=delta_text),
-                                          index=i)
+                            DeltaToolCall(**_named_tool_delta_payload(
+                                tool_choice_function_name,
+                                delta_text,
+                                i,
+                                named_tool_call_ids[i],
+                                first_named_delta,
+                            ))
                         ])
 
                     # handle reasoning: route through reasoning parser while
@@ -665,7 +881,7 @@ class OpenAIServingChat(OpenAIServing):
                                 delta_message, output) and tool_parser:
                             # get the expected call based on partial JSON
                             # parsing which "autocompletes" the JSON
-                            expected_call = json.dumps(
+                            expected_call = _serialize_tool_arguments(
                                 tool_parser.prev_tool_call_arr[index].get(
                                     "arguments", {}))
 
@@ -698,8 +914,9 @@ class OpenAIServingChat(OpenAIServing):
                             index=i,
                             delta=delta_message,
                             logprobs=logprobs,
-                            finish_reason=output.finish_reason
-                            if not auto_tools_called else "tool_calls",
+                            finish_reason=("tool_calls" if (
+                                auto_tools_called or tool_choice_function_name)
+                                else output.finish_reason),
                             stop_reason=output.stop_reason)
                         chunk = ChatCompletionStreamResponse(
                             id=request_id,
@@ -728,7 +945,7 @@ class OpenAIServingChat(OpenAIServing):
             # is sent, send the usage
             if (request.stream_options
                     and request.stream_options.include_usage):
-                completion_tokens = previous_num_tokens[i]
+                completion_tokens = sum(previous_num_tokens)
                 total_reasoning = sum(reasoning_token_counts) if use_reasoning else None
                 final_usage = UsageInfo(
                     prompt_tokens=num_prompt_tokens,
@@ -870,18 +1087,13 @@ class OpenAIServingChat(OpenAIServing):
                 reasoning_text, extracted = r_parser.extract_reasoning(
                     output.text, request)
                 output_text = extracted or ""
+                if isinstance(request.tool_choice,
+                              ChatCompletionNamedToolChoiceParam):
+                    reasoning_text, output_text = \
+                        _reclassify_named_guided_json(
+                            reasoning_text, output_text)
 
-            # Content fallback: if reasoning exists but content is empty,
-            # use the last sentence of reasoning as content.
-            # This ONLY applies to non-tool-call paths.
-            # For tool calls, output_text must be preserved as-is for parsing.
-            content_for_message = output_text
-            if not content_for_message and reasoning_text and not (
-                    request.tools and request.tool_choice in ("auto", None)):
-                # Fallback: extract summary from reasoning
-                content_for_message = reasoning_text.strip().split('\n')[-1]
-                if not content_for_message:
-                    content_for_message = reasoning_text[:200]
+            named_tool_called = False
 
             # if auto tools are not enabled, and a named tool choice using
             #   outlines is not being used
@@ -891,12 +1103,31 @@ class OpenAIServingChat(OpenAIServing):
                         ChatCompletionNamedToolChoiceParam):
                 message = ChatMessage(role=role,
                                       reasoning_content=reasoning_text,
-                                      content=content_for_message)
+                                      content=output_text)
 
             # if the request uses tools and specified a tool choice
             elif request.tool_choice and type(
                     request.tool_choice) is ChatCompletionNamedToolChoiceParam:
 
+                named_tool_called = True
+                parsed_named_tool_calls: Optional[List[ToolCall]] = None
+                if (not _tool_arguments_are_json_object(output_text)
+                        and self.tool_parser is not None):
+                    try:
+                        named_tool_info = self.tool_parser(
+                            tokenizer).extract_tool_calls(
+                                output_text, request=request)
+                        if named_tool_info.tools_called:
+                            parsed_named_tool_calls = named_tool_info.tool_calls
+                    except RuntimeError as e:
+                        logger.warning(
+                            "Named tool parser unavailable; preserving raw "
+                            "arguments: %s", type(e).__name__)
+                named_arguments = _select_named_tool_arguments(
+                    output_text,
+                    request.tool_choice.function.name,
+                    parsed_named_tool_calls,
+                )
                 message = ChatMessage(
                     role=role,
                     reasoning_content=reasoning_text,
@@ -904,7 +1135,7 @@ class OpenAIServingChat(OpenAIServing):
                     tool_calls=[
                         ToolCall(function=FunctionCall(
                             name=request.tool_choice.function.name,
-                            arguments=output_text))
+                            arguments=named_arguments))
                     ])
 
             # if the request doesn't use tool choice
@@ -913,7 +1144,7 @@ class OpenAIServingChat(OpenAIServing):
 
                 message = ChatMessage(role=role,
                                       reasoning_content=reasoning_text,
-                                      content=content_for_message)
+                                      content=output_text)
 
             # handle when there are tools and tool choice is auto
             elif request.tools and (
@@ -940,7 +1171,7 @@ class OpenAIServingChat(OpenAIServing):
                 else:
                     message = ChatMessage(role=role,
                                           reasoning_content=reasoning_text,
-                                          content=content_for_message)
+                                          content=output_text)
 
             # undetermined case that is still important to handle
             else:
@@ -950,13 +1181,14 @@ class OpenAIServingChat(OpenAIServing):
                     "completion.")
                 message = ChatMessage(role=role,
                                       reasoning_content=reasoning_text,
-                                      content=content_for_message)
+                                      content=output_text)
 
             choice_data = ChatCompletionResponseChoice(
                 index=output.index,
                 message=message,
                 logprobs=logprobs,
-                finish_reason="tool_calls" if auto_tools_called else
+                finish_reason="tool_calls" if (
+                    auto_tools_called or named_tool_called) else
                 output.finish_reason if output.finish_reason else "stop",
                 stop_reason=output.stop_reason)
             choices.append(choice_data)
@@ -1000,13 +1232,30 @@ class OpenAIServingChat(OpenAIServing):
 
         request_metadata.final_usage_info = usage
 
+        prompt_logprobs = final_res.prompt_logprobs
+        sample_positions = request.bi100_prompt_logprobs_sample_positions
+        if sample_positions is not None:
+            if num_cached_tokens not in (None, 0):
+                return self.create_error_response(
+                    "BI100 sampled prompt logprobs require a cold request.")
+            if prompt_logprobs is None or (
+                    sample_positions
+                    and sample_positions[-1] >= len(prompt_logprobs)):
+                return self.create_error_response(
+                    "BI100 prompt-logprob sample positions exceed the prompt.")
+            selected = set(sample_positions)
+            prompt_logprobs = [
+                row if position in selected else None
+                for position, row in enumerate(prompt_logprobs)
+            ]
+
         response = ChatCompletionResponse(
             id=request_id,
             created=created_time,
             model=model_name,
             choices=choices,
             usage=usage,
-            prompt_logprobs=final_res.prompt_logprobs,
+            prompt_logprobs=prompt_logprobs,
         )
 
         return response

qwen3_6_scripts/serving_tokenization.py

@@ -0,0 +1,159 @@
+from typing import List, Optional, Union
+
+from vllm.config import ModelConfig
+from vllm.engine.protocol import EngineClient
+from vllm.entrypoints.chat_utils import (apply_hf_chat_template,
+                                         apply_mistral_chat_template,
+                                         load_chat_template,
+                                         parse_chat_messages_futures)
+from vllm.entrypoints.logger import RequestLogger
+# yapf conflicts with isort for this block
+# yapf: disable
+from vllm.entrypoints.openai.protocol import (DetokenizeRequest,
+                                              DetokenizeResponse,
+                                              ErrorResponse,
+                                              TokenizeChatRequest,
+                                              TokenizeRequest,
+                                              TokenizeResponse)
+# yapf: enable
+from vllm.entrypoints.openai.serving_engine import (BaseModelPath,
+                                                    LoRAModulePath,
+                                                    OpenAIServing)
+from vllm.logger import init_logger
+from vllm.transformers_utils.tokenizer import MistralTokenizer
+from vllm.utils import random_uuid
+
+logger = init_logger(__name__)
+
+
+class OpenAIServingTokenization(OpenAIServing):
+
+    def __init__(
+        self,
+        engine_client: EngineClient,
+        model_config: ModelConfig,
+        base_model_paths: List[BaseModelPath],
+        *,
+        lora_modules: Optional[List[LoRAModulePath]],
+        request_logger: Optional[RequestLogger],
+        chat_template: Optional[str],
+    ):
+        super().__init__(engine_client=engine_client,
+                         model_config=model_config,
+                         base_model_paths=base_model_paths,
+                         lora_modules=lora_modules,
+                         prompt_adapters=None,
+                         request_logger=request_logger)
+
+        # If this is None we use the tokenizer's default chat template
+        # the list of commonly-used chat template names for HF named templates
+        hf_chat_templates: List[str] = ['default', 'tool_use']
+        self.chat_template = chat_template \
+            if chat_template in hf_chat_templates \
+            else load_chat_template(chat_template)
+
+    async def create_tokenize(
+        self,
+        request: TokenizeRequest,
+    ) -> Union[TokenizeResponse, ErrorResponse]:
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            return error_check_ret
+
+        request_id = f"tokn-{random_uuid()}"
+
+        (
+            lora_request,
+            prompt_adapter_request,
+        ) = self._maybe_get_adapters(request)
+
+        tokenizer = await self.engine_client.get_tokenizer(lora_request)
+
+        prompt: Union[str, List[int]]
+        if isinstance(request, TokenizeChatRequest):
+            model_config = self.model_config
+
+            conversation, mm_data_future = parse_chat_messages_futures(
+                request.messages, model_config, tokenizer)
+
+            mm_data = await mm_data_future
+            if mm_data:
+                logger.warning(
+                    "Multi-modal inputs are ignored during tokenization")
+
+            if isinstance(tokenizer, MistralTokenizer):
+                prompt = apply_mistral_chat_template(
+                    tokenizer,
+                    messages=request.messages,
+                    chat_template=self.chat_template,
+                    add_generation_prompt=request.add_generation_prompt,
+                    continue_final_message=request.continue_final_message,
+                    **(request.chat_template_kwargs or {}),
+                )
+            else:
+                prompt = apply_hf_chat_template(
+                    tokenizer,
+                    conversation=conversation,
+                    chat_template=self.chat_template,
+                    add_generation_prompt=request.add_generation_prompt,
+                    continue_final_message=request.continue_final_message,
+                    **(request.chat_template_kwargs or {}),
+                )
+        else:
+            prompt = request.prompt
+
+        self._log_inputs(request_id,
+                         prompt,
+                         params=None,
+                         lora_request=lora_request,
+                         prompt_adapter_request=prompt_adapter_request)
+
+        # Silently ignore prompt adapter since it does not affect tokenization
+
+        prompt_input = self._tokenize_prompt_input(
+            request,
+            tokenizer,
+            prompt,
+            add_special_tokens=request.add_special_tokens,
+        )
+        input_ids = prompt_input["prompt_token_ids"]
+
+        return TokenizeResponse(tokens=input_ids,
+                                count=len(input_ids),
+                                max_model_len=self.max_model_len)
+
+    async def create_detokenize(
+        self,
+        request: DetokenizeRequest,
+    ) -> Union[DetokenizeResponse, ErrorResponse]:
+        error_check_ret = await self._check_model(request)
+        if error_check_ret is not None:
+            return error_check_ret
+
+        request_id = f"tokn-{random_uuid()}"
+
+        (
+            lora_request,
+            prompt_adapter_request,
+        ) = self._maybe_get_adapters(request)
+
+        tokenizer = await self.engine_client.get_tokenizer(lora_request)
+
+        self._log_inputs(request_id,
+                         request.tokens,
+                         params=None,
+                         lora_request=lora_request,
+                         prompt_adapter_request=prompt_adapter_request)
+
+        if prompt_adapter_request is not None:
+            raise NotImplementedError("Prompt adapter is not supported "
+                                      "for tokenization")
+
+        prompt_input = self._tokenize_prompt_input(
+            request,
+            tokenizer,
+            request.tokens,
+        )
+        input_text = prompt_input["prompt"]
+
+        return DetokenizeResponse(prompt=input_text)

qwen3_6_scripts/tool_parsers_init.py

@@ -1,12 +0,0 @@
-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"
-]

qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py

@@ -0,0 +1,456 @@
+import math
+from typing import List, Optional
+
+from vllm.core.block.common import BlockList
+from vllm.core.block.interfaces import Block, DeviceAwareBlockAllocator
+from vllm.utils import Device, cdiv, chunk_list
+
+
+class BlockTable:
+    """A class to manage blocks for a specific sequence.
+
+    The BlockTable maps a sequence of tokens to a list of blocks, where each
+    block represents a contiguous memory allocation for a portion of the 
+    sequence. The blocks are managed by a DeviceAwareBlockAllocator, which is
+    responsible for allocating and freeing memory for the blocks.
+
+    Args:
+        block_size (int): The maximum number of tokens that can be stored in a
+            single block.
+        block_allocator (DeviceAwareBlockAllocator): The block allocator used to
+            manage memory for the blocks.
+        _blocks (Optional[List[Block]], optional): An optional list of existing
+            blocks to initialize the BlockTable with. If not provided, an empty
+            BlockTable is created.
+        max_block_sliding_window (Optional[int], optional): The number of
+            blocks to keep around for each sequance. If None, all blocks
+            are kept (eg., when sliding window is not used).
+            It should at least fit the sliding window size of the model.
+
+    Attributes:
+        _block_size (int): The maximum number of tokens that can be stored in a
+            single block.
+        _allocator (DeviceAwareBlockAllocator): The block allocator used to
+            manage memory for the blocks.
+        _blocks (Optional[List[Block]]): The list of blocks managed by this
+            BlockTable.
+        _num_full_slots (int): The number of tokens currently stored in the
+            blocks.
+    """
+
+    def __init__(
+        self,
+        block_size: int,
+        block_allocator: DeviceAwareBlockAllocator,
+        _blocks: Optional[List[Block]] = None,
+        max_block_sliding_window: Optional[int] = None,
+        cache_namespace: Optional[bytes] = None,
+    ):
+        self._block_size = block_size
+        self._allocator = block_allocator
+        self._cache_namespace = cache_namespace
+        if _blocks is None:
+            _blocks = []
+        self._blocks: BlockList = BlockList(_blocks)
+
+        self._max_block_sliding_window = max_block_sliding_window
+        self._num_full_slots = self._get_num_token_ids()
+
+    @staticmethod
+    def get_num_required_blocks(token_ids: List[int],
+                                block_size: int,
+                                num_lookahead_slots: int = 0) -> int:
+        """Calculates the minimum number of blocks required to store a given
+        sequence of token IDs along with any look-ahead slots that may be
+        required (like in multi-step + chunked-prefill).
+
+        This assumes worst-case scenario, where every block requires a new
+        allocation (e.g. ignoring prefix caching).
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be stored.
+            block_size (int): The maximum number of tokens that can be stored in
+                a single block.
+            num_lookahead_slots (int): look-ahead slots that the sequence may
+                require.
+
+        Returns:
+            int: The minimum number of blocks required to store the given
+                sequence of token IDs along with any required look-ahead slots.
+        """
+        return cdiv(len(token_ids) + num_lookahead_slots, block_size)
+
+    def allocate(self,
+                 token_ids: List[int],
+                 device: Device = Device.GPU) -> None:
+        """Allocates memory blocks for storing the given sequence of token IDs.
+
+        This method allocates the required number of blocks to store the given
+        sequence of token IDs.
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be stored.
+            device (Device, optional): The device on which the blocks should be
+                allocated. Defaults to Device.GPU.
+        """
+        assert not self._is_allocated
+        assert token_ids
+        blocks = self._allocate_blocks_for_token_ids(prev_block=None,
+                                                     token_ids=token_ids,
+                                                     device=device)
+        self.update(blocks)
+        self._num_full_slots = len(token_ids)
+
+    def update(self, blocks: List[Block]) -> None:
+        """Resets the table to the newly provided blocks
+        (with their corresponding block ids)
+        """
+        self._blocks.update(blocks)
+
+    def get_content_hashes(self) -> List[bytes]:
+        """Returns block-level content hashes for full blocks in order."""
+        content_hashes: List[bytes] = []
+        for block in self._blocks:
+            block_hash = block.content_hash
+            if block_hash is not None:
+                content_hashes.append(block_hash)
+        return content_hashes
+
+    def append_token_ids(self,
+                         token_ids: List[int],
+                         num_lookahead_slots: int = 0,
+                         num_computed_slots: Optional[int] = None) -> None:
+        """Appends a sequence of token IDs to the existing blocks in the
+        BlockTable.
+
+        This method appends the given sequence of token IDs to the existing
+        blocks in the BlockTable. If there is not enough space in the existing
+        blocks, new blocks are allocated using the `ensure_num_empty_slots`
+        method to accommodate the additional tokens.
+
+        The token IDs are divided into chunks of size `block_size` (except for
+        the first chunk, which may be smaller), and each chunk is appended to a
+        separate block.
+
+        Args:
+            token_ids (List[int]): The sequence of token IDs to be appended.
+            num_computed_slots (Optional[int]): The number of KV cache slots
+                that are already filled (computed).
+                When sliding window is enabled, this is used to compute how many
+                blocks to drop at the front of the sequence.
+                Without sliding window, None can be passed.
+                Without chunked prefill, it should be the same as
+                _num_full_slots.
+        """
+        assert self._is_allocated, "no blocks have been allocated"
+        assert len(self._blocks) > 0
+
+        # Drop blocks that are no longer needed due to sliding window
+        if self._max_block_sliding_window is not None:
+            null_block = self._allocator.allocate_or_get_null_block()
+            assert num_computed_slots is not None
+            end_block_idx = (num_computed_slots //
+                             self._block_size) - self._max_block_sliding_window
+            for idx in range(0, end_block_idx):
+                b = self._blocks[idx]
+                if b is not null_block:
+                    self._allocator.free(b)
+                    self._blocks[idx] = null_block
+
+        # Ensure there are enough empty slots for the new tokens plus
+        # lookahead slots
+        self.ensure_num_empty_slots(num_empty_slots=len(token_ids) +
+                                    num_lookahead_slots)
+
+        # Update the blocks with the new tokens
+        first_block_idx = self._num_full_slots // self._block_size
+        token_blocks = self._chunk_token_blocks_for_append(token_ids)
+
+        for i, token_block in enumerate(token_blocks):
+            self._blocks.append_token_ids(first_block_idx + i, token_block)
+
+        self._num_full_slots += len(token_ids)
+
+    def ensure_num_empty_slots(self, num_empty_slots: int) -> None:
+        """Ensures that the BlockTable has at least the specified number of
+        empty slots available.
+
+        This method checks if the BlockTable has enough empty slots (i.e.,
+        available space) to accommodate the requested number of tokens. If not,
+        it allocates additional blocks on the GPU to ensure that the required
+        number of empty slots is available.
+
+        Args:
+            num_empty_slots (int): The minimum number of empty slots required.
+        """
+        # Currently the block table only supports
+        # appending tokens to GPU blocks.
+        device = Device.GPU
+        assert self._is_allocated
+
+        if self._num_empty_slots >= num_empty_slots:
+            return
+
+        slots_to_allocate = num_empty_slots - self._num_empty_slots
+        blocks_to_allocate = cdiv(slots_to_allocate, self._block_size)
+
+        for _ in range(blocks_to_allocate):
+            assert len(self._blocks) > 0
+            self._blocks.append(
+                self._allocator.allocate_mutable_block(
+                    prev_block=self._blocks[-1], device=device))
+
+    def fork(self) -> "BlockTable":
+        """Creates a new BlockTable instance with a copy of the blocks from the
+        current instance.
+
+        This method creates a new BlockTable instance with the same block size,
+        block allocator, and a copy of the blocks from the current instance. The
+        new BlockTable has its own independent set of blocks, but shares the
+        same underlying memory allocation with the original BlockTable.
+
+        Returns:
+            BlockTable: A new BlockTable instance with a copy of the blocks from
+                the current instance.
+        """
+        assert self._is_allocated
+        assert len(self._blocks) > 0
+        forked_blocks = self._allocator.fork(self._blocks[-1])
+        return BlockTable(
+            block_size=self._block_size,
+            block_allocator=self._allocator,
+            _blocks=forked_blocks,
+            max_block_sliding_window=self._max_block_sliding_window,
+            cache_namespace=self._cache_namespace,
+        )
+
+    def free(self) -> None:
+        """Frees the memory occupied by the blocks in the BlockTable.
+
+        This method iterates over all the blocks in the `_blocks` list and calls
+        the `free` method of the `_allocator` object to release the memory
+        occupied by each block. After freeing all the blocks, the `_blocks` list
+        is set to `None`.
+        """
+        for block in self.blocks:
+            self._allocator.free(block)
+        self._blocks.reset()
+
+    @property
+    def physical_block_ids(self) -> List[int]:
+        """Returns a list of physical block indices for the blocks in the
+        BlockTable.
+
+        This property returns a list of integers, where each integer represents
+        the physical block index of a corresponding block in the `_blocks` list.
+        The physical block index is a unique identifier for the memory location
+        occupied by the block.
+
+        Returns:
+            List[int]: A list of physical block indices for the blocks in the
+                BlockTable.
+        """
+        return self._blocks.ids()
+
+    def get_unseen_token_ids(self, sequence_token_ids: List[int]) -> List[int]:
+        """Get the number of "unseen" tokens in the sequence.
+
+        Unseen tokens are tokens in the sequence corresponding to this block
+        table, but are not yet appended to this block table.
+
+        Args:
+            sequence_token_ids (List[int]): The list of token ids in the
+                sequence.
+
+        Returns:
+            List[int]: The postfix of sequence_token_ids that has not yet been
+                appended to the block table.
+        """
+
+        # Since the block table is append-only, the unseen token ids are the
+        # ones after the appended ones.
+        return sequence_token_ids[self.num_full_slots:]
+
+    def _allocate_blocks_for_token_ids(self, prev_block: Optional[Block],
+                                       token_ids: List[int],
+                                       device: Device) -> List[Block]:
+        blocks: List[Block] = []
+
+        block_token_ids = []
+        tail_token_ids = []
+        for cur_token_ids in chunk_list(token_ids, self._block_size):
+            if len(cur_token_ids) == self._block_size:
+                block_token_ids.append(cur_token_ids)
+            else:
+                tail_token_ids.append(cur_token_ids)
+
+        if block_token_ids:
+            blocks.extend(self._allocate_immutable_blocks(
+                prev_block=prev_block,
+                block_token_ids=block_token_ids,
+                device=device))
+            prev_block = blocks[-1]
+
+        if tail_token_ids:
+            assert len(tail_token_ids) == 1
+            cur_token_ids = tail_token_ids[0]
+
+            block = self._allocate_mutable_block(prev_block=prev_block,
+                                                 device=device)
+            block.append_token_ids(cur_token_ids)
+
+            blocks.append(block)
+
+        return blocks
+
+    def _allocate_mutable_block(self, prev_block: Optional[Block],
+                                device: Device) -> Block:
+        if self._cache_namespace is None:
+            return self._allocator.allocate_mutable_block(
+                prev_block=prev_block, device=device)
+
+        with_cache_namespace = getattr(
+            self._allocator, "allocate_mutable_block_with_cache_namespace",
+            None)
+        if callable(with_cache_namespace):
+            return with_cache_namespace(
+                prev_block=prev_block,
+                cache_namespace=self._cache_namespace,
+                device=device)
+
+        backend_allocators = getattr(self._allocator, "_allocators", None)
+        if isinstance(backend_allocators, dict):
+            device_allocator = backend_allocators.get(device)
+            if device_allocator is not None:
+                with_cache_namespace = getattr(
+                    device_allocator,
+                    "allocate_mutable_block_with_cache_namespace", None)
+                if callable(with_cache_namespace):
+                    return with_cache_namespace(
+                        prev_block=prev_block,
+                        cache_namespace=self._cache_namespace)
+
+        return self._allocator.allocate_mutable_block(
+            prev_block=prev_block, device=device)
+
+    def _allocate_immutable_blocks(self,
+                                  prev_block: Optional[Block],
+                                  block_token_ids: List[List[int]],
+                                  device: Device) -> List[Block]:
+        if self._cache_namespace is None:
+            return self._allocator.allocate_immutable_blocks(
+                prev_block,
+                block_token_ids=block_token_ids,
+                device=device)
+
+        with_cache_namespace = getattr(
+            self._allocator, "allocate_immutable_blocks_with_cache_namespace", None)
+        if callable(with_cache_namespace):
+            return with_cache_namespace(
+                prev_block=prev_block,
+                block_token_ids=block_token_ids,
+                cache_namespace=self._cache_namespace,
+                device=device)
+
+        backend_allocator = getattr(self._allocator, "_allocators", None)
+        if isinstance(backend_allocator, dict):
+            device_allocator = backend_allocator.get(device)
+            if device_allocator is not None:
+                with_cache_namespace = getattr(
+                    device_allocator,
+                    "allocate_immutable_blocks_with_cache_namespace",
+                    None)
+                if callable(with_cache_namespace):
+                    return with_cache_namespace(
+                        prev_block=prev_block,
+                        block_token_ids=block_token_ids,
+                        cache_namespace=self._cache_namespace)
+
+        # Fallback: keep behavior identical when no namespace-aware allocator
+        # is available.
+        return self._allocator.allocate_immutable_blocks(
+            prev_block,
+            block_token_ids=block_token_ids,
+            device=device)
+
+    def _get_all_token_ids(self) -> List[int]:
+        # NOTE: This function is O(seq_len); use sparingly.
+        token_ids: List[int] = []
+
+        if not self._is_allocated:
+            return token_ids
+
+        for block in self.blocks:
+            token_ids.extend(block.token_ids)
+
+        return token_ids
+
+    def _get_num_token_ids(self) -> int:
+        res = 0
+        for block in self.blocks:
+            res += len(block.token_ids)
+
+        return res
+
+    @property
+    def _is_allocated(self) -> bool:
+        return len(self._blocks) > 0
+
+    @property
+    def blocks(self) -> List[Block]:
+        return self._blocks.list()
+
+    @property
+    def _num_empty_slots(self) -> int:
+        assert self._is_allocated
+        return len(self._blocks) * self._block_size - self._num_full_slots
+
+    @property
+    def num_full_slots(self) -> int:
+        """Returns the total number of tokens currently stored in the
+        BlockTable.
+
+        Returns:
+            int: The total number of tokens currently stored in the BlockTable.
+        """
+        return self._num_full_slots
+
+    def get_num_blocks_touched_by_append_slots(
+            self, token_ids: List[int], num_lookahead_slots: int) -> int:
+        """Determine how many blocks will be "touched" by appending the token
+        ids.
+
+        This is required for the scheduler to determine whether a sequence can
+        continue generation, or if it must be preempted.
+        """
+        # Math below is equivalent to:
+        # all_token_ids = token_ids + [-1] * num_lookahead_slots
+        # token_blocks = self._chunk_token_blocks_for_append(all_token_ids)
+        # return len(token_blocks)
+
+        num_token_ids = len(token_ids) + num_lookahead_slots
+        first_chunk_size = self._block_size - (self._num_full_slots %
+                                               self._block_size)
+        num_token_blocks = (1 + math.ceil(
+            (num_token_ids - first_chunk_size) / self._block_size))
+        return num_token_blocks
+
+    def _chunk_token_blocks_for_append(
+            self, token_ids: List[int]) -> List[List[int]]:
+        """Split the token ids into block-sized chunks so they can be easily
+        appended to blocks. The first such "token block" may have less token ids
+        than the block size, since the last allocated block may be partially
+        full.
+
+        If no token ids are provided, then no chunks are returned.
+        """
+
+        if not token_ids:
+            return []
+
+        first_chunk_size = self._block_size - (self._num_full_slots %
+                                               self._block_size)
+        token_blocks = [token_ids[:first_chunk_size]]
+        token_blocks.extend(
+            chunk_list(token_ids[first_chunk_size:], self._block_size))
+        return token_blocks

qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py

@@ -0,0 +1,475 @@
+from typing import Dict, FrozenSet, List, Optional, Tuple
+
+from vllm.core.block.cpu_kv_content_cache import (CpuKvContentCache,
+                                                   cpu_kv_offload_enabled)
+from vllm.core.block.interfaces import (Block, BlockAllocator, BlockId,
+                                        DeviceAwareBlockAllocator)
+from vllm.core.block.naive_block import NaiveBlock, NaiveBlockAllocator
+from vllm.core.block.prefix_caching_block import PrefixCachingBlockAllocator
+from vllm.utils import Device
+
+
+class CpuGpuBlockAllocator(DeviceAwareBlockAllocator):
+    """A block allocator that can allocate blocks on both CPU and GPU memory.
+
+    This class implements the `DeviceAwareBlockAllocator` interface and provides
+    functionality for allocating and managing blocks of memory on both CPU and
+    GPU devices.
+
+    The `CpuGpuBlockAllocator` maintains separate memory pools for CPU and GPU
+    blocks, and allows for allocation, deallocation, forking, and swapping of
+    blocks across these memory pools.
+    """
+
+    @staticmethod
+    def create(
+        allocator_type: str,
+        num_gpu_blocks: int,
+        num_cpu_blocks: int,
+        block_size: int,
+    ) -> DeviceAwareBlockAllocator:
+        """Creates a CpuGpuBlockAllocator instance with the specified
+        configuration.
+
+        This static method creates and returns a CpuGpuBlockAllocator instance
+        based on the provided parameters. It initializes the CPU and GPU block
+        allocators with the specified number of blocks, block size, and
+        allocator type.
+
+        Args:
+            allocator_type (str): The type of block allocator to use for CPU
+                and GPU blocks. Currently supported values are "naive" and
+                "prefix_caching".
+            num_gpu_blocks (int): The number of blocks to allocate for GPU
+                memory.
+            num_cpu_blocks (int): The number of blocks to allocate for CPU
+                memory.
+            block_size (int): The size of each block in number of tokens.
+
+        Returns:
+            DeviceAwareBlockAllocator: A CpuGpuBlockAllocator instance with the
+                specified configuration.
+
+        Notes:
+            - The block IDs are assigned contiguously, with GPU block IDs coming
+                before CPU block IDs.
+        """
+        content_offload = cpu_kv_offload_enabled()
+        if content_offload and allocator_type != "prefix_caching":
+            raise RuntimeError(
+                "BI100_CPU_KV_OFFLOAD=1 requires prefix caching")
+        if content_offload and num_cpu_blocks <= 0:
+            raise RuntimeError(
+                "BI100_CPU_KV_OFFLOAD=1 requires at least one CPU KV block")
+
+        block_ids = list(range(num_gpu_blocks + num_cpu_blocks))
+        gpu_block_ids = block_ids[:num_gpu_blocks]
+        cpu_block_ids = block_ids[num_gpu_blocks:]
+
+        if allocator_type == "naive":
+            gpu_allocator: BlockAllocator = NaiveBlockAllocator(
+                create_block=NaiveBlock,  # type: ignore
+                num_blocks=num_gpu_blocks,
+                block_size=block_size,
+                block_ids=gpu_block_ids,
+            )
+
+            cpu_allocator: BlockAllocator = NaiveBlockAllocator(
+                create_block=NaiveBlock,  # type: ignore
+                num_blocks=num_cpu_blocks,
+                block_size=block_size,
+                block_ids=cpu_block_ids,
+            )
+        elif allocator_type == "prefix_caching":
+            gpu_allocator = PrefixCachingBlockAllocator(
+                num_blocks=num_gpu_blocks,
+                block_size=block_size,
+                block_ids=gpu_block_ids,
+            )
+
+            cpu_allocator = PrefixCachingBlockAllocator(
+                num_blocks=num_cpu_blocks,
+                block_size=block_size,
+                block_ids=cpu_block_ids,
+            )
+        else:
+            raise ValueError(f"Unknown allocator type {allocator_type=}")
+
+        return CpuGpuBlockAllocator(
+            cpu_block_allocator=cpu_allocator,
+            gpu_block_allocator=gpu_allocator,
+            cpu_content_cache=(CpuKvContentCache(num_cpu_blocks)
+                               if content_offload else None),
+        )
+
+    def __init__(self, cpu_block_allocator: BlockAllocator,
+                 gpu_block_allocator: BlockAllocator,
+                 cpu_content_cache: Optional[CpuKvContentCache] = None):
+        assert not (
+            cpu_block_allocator.all_block_ids
+            & gpu_block_allocator.all_block_ids
+        ), "cpu and gpu block allocators can't have intersection of block ids"
+
+        self._allocators = {
+            Device.CPU: cpu_block_allocator,
+            Device.GPU: gpu_block_allocator,
+        }
+
+        self._swap_mapping: Dict[int, int] = {}
+        self._null_block: Optional[Block] = None
+        self._cpu_content_cache = cpu_content_cache
+
+        self._block_ids_to_allocator: Dict[int, BlockAllocator] = {}
+        for _, allocator in self._allocators.items():
+            for block_id in allocator.all_block_ids:
+                self._block_ids_to_allocator[block_id] = allocator
+
+        if self._cpu_content_cache is not None:
+            if not isinstance(gpu_block_allocator,
+                              PrefixCachingBlockAllocator):
+                raise RuntimeError(
+                    "CPU KV content tier requires PrefixCachingBlockAllocator")
+            if (self._cpu_content_cache.capacity !=
+                    cpu_block_allocator.get_num_total_blocks()):
+                raise RuntimeError(
+                    "CPU KV content capacity must cover the complete CPU cache")
+            gpu_block_allocator.set_external_cache_callbacks(
+                claim=self._claim_cpu_content,
+                load=self._stage_cpu_to_gpu,
+                cancel=self._cancel_cpu_claim,
+                store=self._stage_gpu_to_cpu,
+            )
+
+    @property
+    def content_offload_enabled(self) -> bool:
+        return self._cpu_content_cache is not None
+
+    def _claim_cpu_content(self, content_hash: bytes) -> Optional[int]:
+        assert self._cpu_content_cache is not None
+        return self._cpu_content_cache.claim_load(content_hash)
+
+    def _cancel_cpu_claim(self, content_hash: bytes, cpu_slot: int) -> None:
+        assert self._cpu_content_cache is not None
+        self._cpu_content_cache.cancel_load(content_hash, cpu_slot)
+
+    def _stage_cpu_to_gpu(self, content_hash: bytes, cpu_slot: int,
+                          gpu_block_id: BlockId) -> None:
+        assert self._cpu_content_cache is not None
+        gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
+        self._cpu_content_cache.stage_load(
+            content_hash, cpu_slot, gpu_slot)
+
+    def _stage_gpu_to_cpu(self, content_hash: bytes,
+                          gpu_block_id: BlockId) -> bool:
+        assert self._cpu_content_cache is not None
+        gpu_slot = self.get_physical_block_id(Device.GPU, gpu_block_id)
+        return self._cpu_content_cache.stage_store(content_hash, gpu_slot)
+
+    def allocate_or_get_null_block(self) -> Block:
+        if self._null_block is None:
+            self._null_block = NullBlock(
+                self.allocate_mutable_block(None, Device.GPU))
+        return self._null_block
+
+    def allocate_mutable_block(self, prev_block: Optional[Block],
+                               device: Device) -> Block:
+        """Allocates a new mutable block on the specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block to in the sequence.
+                Used for prefix hashing.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            Block: The newly allocated mutable block.
+        """
+        return self._allocators[device].allocate_mutable_block(prev_block)
+
+    def allocate_immutable_blocks(self, prev_block: Optional[Block],
+                                  block_token_ids: List[List[int]],
+                                  device: Device) -> List[Block]:
+        """Allocates a new group of immutable blocks with the provided block 
+        token IDs on the specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block in the sequence.
+                Used for prefix hashing.
+            block_token_ids (List[int]): The list of block token IDs to be 
+                stored in the new blocks.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            List[Block]: The newly allocated list of immutable blocks 
+                containing the provided block token IDs.
+        """
+        return self._allocators[device].allocate_immutable_blocks(
+            prev_block, block_token_ids)
+
+    def allocate_immutable_block(self, prev_block: Optional[Block],
+                                 token_ids: List[int],
+                                 device: Device) -> Block:
+        """Allocates a new immutable block with the provided token IDs on the
+        specified device.
+
+        Args:
+            prev_block (Optional[Block]): The previous block in the sequence.
+                Used for prefix hashing.
+            token_ids (List[int]): The list of token IDs to be stored in the new
+                block.
+            device (Device): The device on which to allocate the new block.
+
+        Returns:
+            Block: The newly allocated immutable block containing the provided
+                token IDs.
+        """
+        return self._allocators[device].allocate_immutable_block(
+            prev_block, token_ids)
+
+    def free(self, block: Block) -> None:
+        """Frees the memory occupied by the given block.
+
+        Args:
+            block (Block): The block to be freed.
+        """
+        # Null block should never be freed
+        if isinstance(block, NullBlock):
+            return
+        block_id = block.block_id
+        assert block_id is not None
+        allocator = self._block_ids_to_allocator[block_id]
+        allocator.free(block)
+
+    def fork(self, last_block: Block) -> List[Block]:
+        """Creates a new sequence of blocks that shares the same underlying
+            memory as the original sequence.
+
+        Args:
+            last_block (Block): The last block in the original sequence.
+
+        Returns:
+            List[Block]: A new list of blocks that shares the same memory as the
+                original sequence.
+        """
+        # do not attempt to fork the null block
+        assert not isinstance(last_block, NullBlock)
+        block_id = last_block.block_id
+        assert block_id is not None
+        allocator = self._block_ids_to_allocator[block_id]
+        return allocator.fork(last_block)
+
+    def get_num_free_blocks(self, device: Device) -> int:
+        """Returns the number of free blocks available on the specified device.
+
+        Args:
+            device (Device): The device for which to query the number of free
+                blocks. AssertionError is raised if None is passed.
+
+        Returns:
+            int: The number of free blocks available on the specified device.
+        """
+        return self._allocators[device].get_num_free_blocks()
+
+    def get_num_total_blocks(self, device: Device) -> int:
+        return self._allocators[device].get_num_total_blocks()
+
+    def get_physical_block_id(self, device: Device, absolute_id: int) -> int:
+        """Returns the zero-offset block id on certain device given the 
+        absolute block id.
+
+        Args:
+            device (Device): The device for which to query relative block id.
+                absolute_id (int): The absolute block id for the block in 
+                whole allocator.
+
+        Returns:
+            int: The zero-offset block id on certain device.
+        """
+        return self._allocators[device].get_physical_block_id(absolute_id)
+
+    def swap(self, blocks: List[Block], src_device: Device,
+             dst_device: Device) -> Dict[int, int]:
+        """Execute the swap for the given blocks from source_device
+        on to dest_device, save the current swap mapping and append 
+        them to the accumulated `self._swap_mapping` for each 
+        scheduling move.
+
+        Args:
+            blocks: List of blocks to be swapped.
+            src_device (Device): Device to swap the 'blocks' from.
+            dst_device (Device): Device to swap the 'blocks' to.
+        
+        Returns:
+            Dict[int, int]: Swap mapping from source_device
+                on to dest_device.
+        """
+        if self.content_offload_enabled:
+            raise RuntimeError(
+                "request-level preemption swap cannot share CPU slots with "
+                "BI100_CPU_KV_OFFLOAD")
+
+        src_block_ids = [block.block_id for block in blocks]
+        self._allocators[src_device].swap_out(blocks)
+        self._allocators[dst_device].swap_in(blocks)
+        dst_block_ids = [block.block_id for block in blocks]
+
+        current_swap_mapping: Dict[int, int] = {}
+        for src_block_id, dst_block_id in zip(src_block_ids, dst_block_ids):
+            if src_block_id is not None and dst_block_id is not None:
+                self._swap_mapping[src_block_id] = dst_block_id
+                current_swap_mapping[src_block_id] = dst_block_id
+        return current_swap_mapping
+
+    def get_num_full_blocks_touched(self, blocks: List[Block],
+                                    device: Device) -> int:
+        """Returns the number of full blocks that will be touched by
+        swapping in/out the given blocks on to the 'device'.
+
+        Args:
+            blocks: List of blocks to be swapped.
+            device (Device): Device to swap the 'blocks' on.
+
+        Returns:
+            int: the number of full blocks that will be touched by
+                swapping in/out the given blocks on to the 'device'.
+                Non full blocks are ignored when deciding the number
+                of blocks to touch.
+        """
+        return self._allocators[device].get_num_full_blocks_touched(blocks)
+
+    def clear_copy_on_writes(self) -> List[Tuple[int, int]]:
+        """Clears the copy-on-write (CoW) state and returns the mapping of
+            source to destination block IDs.
+
+        Returns:
+            List[Tuple[int, int]]: A list mapping source block IDs to 
+                destination block IDs.
+        """
+        # CoW only supported on GPU
+        device = Device.GPU
+        return self._allocators[device].clear_copy_on_writes()
+
+    def mark_blocks_as_accessed(self, block_ids: List[int],
+                                now: float) -> None:
+        """Mark blocks as accessed, only use for prefix caching."""
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].mark_blocks_as_accessed(block_ids, now)
+
+    def mark_blocks_as_computed(self, block_ids: List[int]) -> None:
+        """Mark blocks as accessed, only use for prefix caching."""
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].mark_blocks_as_computed(block_ids)
+
+    def get_computed_block_ids(self, prev_computed_block_ids: List[int],
+                               block_ids: List[int],
+                               skip_last_block_id: bool) -> List[int]:
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].get_computed_block_ids(
+            prev_computed_block_ids, block_ids, skip_last_block_id)
+
+    def get_common_computed_block_ids(
+            self, computed_seq_block_ids: List[List[int]]) -> List[int]:
+        # Prefix caching only supported on GPU.
+        device = Device.GPU
+        return self._allocators[device].get_common_computed_block_ids(
+            computed_seq_block_ids)
+
+    @property
+    def all_block_ids(self) -> FrozenSet[int]:
+        return frozenset(self._block_ids_to_allocator.keys())
+
+    def get_prefix_cache_hit_rate(self, device: Device) -> float:
+        """Prefix cache hit rate. -1 means not supported or disabled."""
+        assert device in self._allocators
+        return self._allocators[device].get_prefix_cache_hit_rate()
+
+    def get_and_reset_swaps(self) -> List[Tuple[int, int]]:
+        """Returns and clears the mapping of source to destination block IDs.
+        Will be called after every swapping operations for now, and after every
+        schedule when BlockManagerV2 become default. Currently not useful.
+
+        Returns:
+            List[Tuple[int, int]]: A mapping of source to destination block IDs.
+        """
+        mapping = self._swap_mapping.copy()
+        self._swap_mapping.clear()
+        return list(mapping.items())
+
+    def get_and_reset_prefix_swaps(
+            self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
+        """Return scheduler-owned (CPU->GPU, GPU->CPU) content maps."""
+        if self._cpu_content_cache is None:
+            return [], []
+        return self._cpu_content_cache.drain_step()
+
+    def begin_prefix_cache_step(self) -> None:
+        if self._cpu_content_cache is not None:
+            self._cpu_content_cache.begin_step()
+
+
+class NullBlock(Block):
+    """
+    Null blocks are used as a placeholders for KV cache blocks that have
+    been dropped due to sliding window.
+    This implementation just wraps an ordinary block and prevents it from
+    being modified. It also allows for testing if a block is NullBlock
+    via isinstance().
+    """
+
+    def __init__(self, proxy: Block):
+        super().__init__()
+        self._proxy = proxy
+
+    def append_token_ids(self, token_ids: List[BlockId]):
+        raise ValueError("null block should not be modified")
+
+    @property
+    def block_id(self):
+        return self._proxy.block_id
+
+    @block_id.setter
+    def block_id(self, value: Optional[BlockId]):
+        raise ValueError("null block should not be modified")
+
+    @property
+    def token_ids(self) -> List[BlockId]:
+        return self._proxy.token_ids
+
+    @property
+    def num_tokens_total(self) -> int:
+        raise NotImplementedError(
+            "num_tokens_total is not used for null block")
+
+    @property
+    def num_empty_slots(self) -> BlockId:
+        return self._proxy.num_empty_slots
+
+    @property
+    def is_full(self):
+        return self._proxy.is_full
+
+    @property
+    def prev_block(self):
+        return self._proxy.prev_block
+
+    @property
+    def computed(self):
+        return self._proxy.computed
+
+    @computed.setter
+    def computed(self, value):
+        self._proxy.computed = value
+
+    @property
+    def last_accessed(self) -> float:
+        return self._proxy.last_accessed
+
+    @last_accessed.setter
+    def last_accessed(self, last_accessed_ts: float):
+        self._proxy.last_accessed = last_accessed_ts
+
+    @property
+    def content_hash(self):
+        return self._proxy.content_hash

qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py

@@ -0,0 +1,255 @@
+"""Scheduler-owned content index for an inclusive CPU KV cache tier."""
+
+from __future__ import annotations
+
+import heapq
+import os
+from collections import OrderedDict
+from typing import Dict, List, Mapping, Optional, Set, Tuple
+
+
+ContentHash = bytes
+SwapMapping = List[Tuple[int, int]]
+
+
+def cpu_kv_offload_enabled(
+    environ: Optional[Mapping[str, str]] = None,
+) -> bool:
+    """Read the experimental selector without accepting ambiguous values."""
+    source = os.environ if environ is None else environ
+    value = source.get("BI100_CPU_KV_OFFLOAD", "0")
+    if value == "0":
+        return False
+    if value == "1":
+        return True
+    raise RuntimeError(
+        "BI100_CPU_KV_OFFLOAD must be exactly '0' or '1', "
+        f"got {value!r}")
+
+
+class CpuKvContentCache:
+    """Track immutable KV blocks held in the worker's pinned CPU cache.
+
+    The scheduler owns this metadata and sends identical physical block maps
+    to every tensor-parallel worker. CPU copies are inclusive: loading a block
+    back to GPU does not remove its CPU entry. Slots touched by either transfer
+    direction are pinned for the whole scheduling step so a D2H destination
+    can never overwrite an H2D source before workers execute the maps.
+    """
+
+    def __init__(self, capacity: int) -> None:
+        if capacity <= 0:
+            raise ValueError("CPU KV content cache capacity must be positive")
+        self.capacity = capacity
+        self._hash_to_slot: Dict[ContentHash, int] = {}
+        self._slot_to_hash: Dict[int, ContentHash] = {}
+        self._ready_slots: Set[int] = set()
+        self._lru: OrderedDict[int, None] = OrderedDict()
+        self._free_slots = list(range(capacity))
+        heapq.heapify(self._free_slots)
+
+        self._step_slots_in_use: Set[int] = set()
+        self._step_load_slots: Set[int] = set()
+        self._step_h2d: Dict[int, int] = {}
+        self._step_d2h: Dict[int, int] = {}
+        self._deferred_d2h: Dict[int, ContentHash] = {}
+        self._deferred_hashes: Set[ContentHash] = set()
+        self._pending_ready_slots: Set[int] = set()
+
+        self.hits = 0
+        self.misses = 0
+        self.stores = 0
+        self.deduplicated_stores = 0
+        self.evictions = 0
+        self.skipped_stores = 0
+
+    @staticmethod
+    def _validate_hash(content_hash: ContentHash) -> None:
+        if not isinstance(content_hash, bytes) or len(content_hash) != 32:
+            raise ValueError("CPU KV cache key must be a 32-byte content hash")
+
+    @staticmethod
+    def _validate_block_id(name: str, block_id: int) -> None:
+        if not isinstance(block_id, int) or isinstance(block_id, bool):
+            raise TypeError(f"{name} must be an integer")
+        if block_id < 0:
+            raise ValueError(f"{name} must be non-negative")
+
+    def _touch(self, slot: int) -> None:
+        self._lru.pop(slot, None)
+        self._lru[slot] = None
+
+    def _select_store_slot(self) -> Optional[int]:
+        if self._free_slots:
+            return heapq.heappop(self._free_slots)
+        for slot in self._lru:
+            if slot not in self._step_slots_in_use:
+                return slot
+        return None
+
+    def _commit_store(self, content_hash: ContentHash,
+                      gpu_block: int, slot: int) -> None:
+        old_hash = self._slot_to_hash.get(slot)
+        if old_hash is not None:
+            if slot in self._step_slots_in_use:
+                raise RuntimeError("selected an in-use CPU KV slot for eviction")
+            del self._hash_to_slot[old_hash]
+            self._ready_slots.discard(slot)
+            self.evictions += 1
+
+        if slot in self._step_h2d:
+            raise RuntimeError(
+                "a CPU KV slot cannot be an H2D source and D2H destination "
+                "in one scheduler step")
+        if slot in self._step_d2h.values():
+            raise RuntimeError(f"duplicate D2H destination CPU slot {slot}")
+
+        self._hash_to_slot[content_hash] = slot
+        self._slot_to_hash[slot] = content_hash
+        self._ready_slots.discard(slot)
+        self._step_slots_in_use.add(slot)
+        self._step_d2h[gpu_block] = slot
+        self._touch(slot)
+        self.stores += 1
+
+    def begin_step(self) -> None:
+        """Publish D2H stores returned by the preceding synchronous step."""
+        if (self._step_slots_in_use or self._step_h2d or self._step_d2h
+                or self._deferred_d2h or self._deferred_hashes):
+            raise RuntimeError("cannot begin a CPU KV step before draining it")
+        self._ready_slots.update(self._pending_ready_slots)
+        self._pending_ready_slots.clear()
+
+    def _require_step_started(self) -> None:
+        if self._pending_ready_slots:
+            raise RuntimeError(
+                "CPU KV step must begin before content lookup or eviction")
+
+    def claim_load(self, content_hash: ContentHash) -> Optional[int]:
+        """Pin and return a ready CPU source for this scheduling step."""
+        self._validate_hash(content_hash)
+        self._require_step_started()
+        slot = self._hash_to_slot.get(content_hash)
+        if slot is None or slot not in self._ready_slots:
+            self.misses += 1
+            return None
+        if slot in self._step_slots_in_use:
+            raise RuntimeError(
+                f"CPU KV slot {slot} was claimed twice in one scheduler step")
+        self._step_slots_in_use.add(slot)
+        self._step_load_slots.add(slot)
+        self._touch(slot)
+        self.hits += 1
+        return slot
+
+    def cancel_load(self, content_hash: ContentHash, cpu_slot: int) -> None:
+        """Release a claim when GPU allocation fails before H2D is staged."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("cpu_slot", cpu_slot)
+        if self._hash_to_slot.get(content_hash) != cpu_slot:
+            raise RuntimeError("CPU KV load cancellation key/slot mismatch")
+        if cpu_slot in self._step_h2d:
+            raise RuntimeError("cannot cancel a CPU KV load after H2D staging")
+        if cpu_slot not in self._step_slots_in_use:
+            raise RuntimeError("cannot cancel an unclaimed CPU KV load")
+        self._step_slots_in_use.remove(cpu_slot)
+        self._step_load_slots.remove(cpu_slot)
+
+    def stage_load(self, content_hash: ContentHash, cpu_slot: int,
+                   gpu_block: int) -> None:
+        """Stage one CPU-to-GPU promotion after the GPU slot is reserved."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("cpu_slot", cpu_slot)
+        self._validate_block_id("gpu_block", gpu_block)
+        if self._hash_to_slot.get(content_hash) != cpu_slot:
+            raise RuntimeError("CPU KV load key/slot mismatch")
+        if cpu_slot not in self._ready_slots:
+            raise RuntimeError("CPU KV load source is not ready")
+        if cpu_slot not in self._step_slots_in_use:
+            raise RuntimeError("CPU KV load source was not claimed")
+        if cpu_slot in self._step_h2d:
+            raise RuntimeError(f"duplicate H2D source CPU slot {cpu_slot}")
+        if gpu_block in self._step_h2d.values():
+            raise RuntimeError(f"duplicate H2D destination GPU block {gpu_block}")
+        if cpu_slot in self._step_d2h.values():
+            raise RuntimeError(
+                "a CPU KV slot cannot be an H2D source and D2H destination "
+                "in one scheduler step")
+        self._step_h2d[cpu_slot] = gpu_block
+
+    def stage_store(self, content_hash: ContentHash,
+                    gpu_block: int) -> bool:
+        """Stage a lazy GPU-to-CPU copy for an evicted immutable block."""
+        self._validate_hash(content_hash)
+        self._validate_block_id("gpu_block", gpu_block)
+        self._require_step_started()
+        if (content_hash in self._hash_to_slot
+                or content_hash in self._deferred_hashes):
+            slot = self._hash_to_slot.get(content_hash)
+            if slot is not None:
+                self._touch(slot)
+            self.deduplicated_stores += 1
+            return False
+        if gpu_block in self._step_d2h or gpu_block in self._deferred_d2h:
+            raise RuntimeError(f"duplicate D2H source GPU block {gpu_block}")
+
+        if self._free_slots:
+            self._commit_store(
+                content_hash, gpu_block, heapq.heappop(self._free_slots))
+            return True
+
+        # Do not replace resident content until every lookup in this scheduler
+        # step is known. A later H2D claim can refer to any current LRU entry.
+        self._deferred_d2h[gpu_block] = content_hash
+        self._deferred_hashes.add(content_hash)
+        return True
+
+    def _resolve_deferred_stores(self) -> None:
+        if self._step_load_slots:
+            self.skipped_stores += len(self._deferred_d2h)
+        else:
+            for gpu_block, content_hash in self._deferred_d2h.items():
+                slot = self._select_store_slot()
+                if slot is None:
+                    self.skipped_stores += 1
+                    continue
+                self._commit_store(content_hash, gpu_block, slot)
+        self._deferred_d2h.clear()
+        self._deferred_hashes.clear()
+
+    def drain_step(self) -> Tuple[SwapMapping, SwapMapping]:
+        """Finalize this synchronous step and return (H2D, D2H) maps."""
+        self._resolve_deferred_stores()
+        transfer_slots = (
+            set(self._step_h2d) | set(self._step_d2h.values()))
+        if transfer_slots != self._step_slots_in_use:
+            raise RuntimeError(
+                "CPU KV scheduler step contains an uncommitted slot claim")
+        if set(self._step_h2d) & set(self._step_d2h.values()):
+            raise RuntimeError(
+                "CPU KV scheduler step reuses a CPU slot across directions")
+        if set(self._step_h2d) != self._step_load_slots:
+            raise RuntimeError(
+                "CPU KV scheduler step contains an unstaged load claim")
+
+        swap_in = sorted(self._step_h2d.items())
+        swap_out = sorted(self._step_d2h.items())
+        self._pending_ready_slots.update(self._step_d2h.values())
+        self._step_h2d.clear()
+        self._step_d2h.clear()
+        self._step_slots_in_use.clear()
+        self._step_load_slots.clear()
+        return swap_in, swap_out
+
+    def resident_slot(self, content_hash: ContentHash) -> Optional[int]:
+        self._validate_hash(content_hash)
+        return self._hash_to_slot.get(content_hash)
+
+    def is_ready(self, content_hash: ContentHash) -> bool:
+        self._validate_hash(content_hash)
+        slot = self._hash_to_slot.get(content_hash)
+        return slot is not None and slot in self._ready_slots
+
+    @property
+    def resident_count(self) -> int:
+        return len(self._hash_to_slot)

qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py

@@ -0,0 +1,769 @@
+"""A block manager that manages token blocks."""
+import hashlib
+import os
+import struct
+from collections.abc import Mapping
+from typing import Any, Dict, List, Optional, Sequence as GenericSequence, Tuple
+
+try:
+    from PIL import Image
+except Exception:  # pragma: no cover - optional dependency in some envs
+    Image = None  # type: ignore
+try:
+    import torch
+except Exception:  # pragma: no cover - optional dependency in some envs
+    torch = None  # type: ignore
+
+from vllm.core.block.block_table import BlockTable
+from vllm.core.block.cpu_gpu_block_allocator import CpuGpuBlockAllocator
+from vllm.core.block.interfaces import Block
+from vllm.core.block.prefix_caching_block import (ComputedBlocksTracker,
+                                                  LastAccessBlocksTracker)
+from vllm.core.block.utils import check_no_caching_or_swa_for_blockmgr_encdec
+from vllm.core.interfaces import AllocStatus, BlockSpaceManager
+from vllm.logger import init_logger
+from vllm.sequence import Sequence, SequenceGroup, SequenceStatus
+from vllm.utils import Device
+
+SeqId = int
+EncoderSeqId = str
+
+logger = init_logger(__name__)
+
+
+class BlockSpaceManagerV2(BlockSpaceManager):
+    """BlockSpaceManager which manages the allocation of KV cache.
+
+    It owns responsibility for allocation, swapping, allocating memory for
+    autoregressively-generated tokens, and other advanced features such as
+    prefix caching, forking/copy-on-write, and sliding-window memory allocation.
+
+    This class implements the design described in
+    https://github.com/vllm-project/vllm/pull/3492.
+
+    Lookahead slots
+        The block manager has the notion of a "lookahead slot". These are slots
+        in the KV cache that are allocated for a sequence. Unlike the other
+        allocated slots, the content of these slots is undefined -- the worker
+        may use the memory allocations in any way.
+
+        In practice, a worker could use these lookahead slots to run multiple
+        forward passes for a single scheduler invocation. Each successive
+        forward pass would write KV activations to the corresponding lookahead
+        slot. This allows low inter-token latency use-cases, where the overhead
+        of continuous batching scheduling is amortized over >1 generated tokens.
+
+        Speculative decoding uses lookahead slots to store KV activations of
+        proposal tokens.
+
+        See https://github.com/vllm-project/vllm/pull/3250 for more information
+        on lookahead scheduling.
+
+    Args:
+        block_size (int): The size of each memory block.
+        num_gpu_blocks (int): The number of memory blocks allocated on GPU.
+        num_cpu_blocks (int): The number of memory blocks allocated on CPU.
+        watermark (float, optional): The threshold used for memory swapping.
+            Defaults to 0.01.
+        sliding_window (Optional[int], optional): The size of the sliding
+            window. Defaults to None.
+        enable_caching (bool, optional): Flag indicating whether caching is
+            enabled. Defaults to False.
+    """
+
+    def __init__(
+        self,
+        block_size: int,
+        num_gpu_blocks: int,
+        num_cpu_blocks: int,
+        watermark: float = 0.01,
+        sliding_window: Optional[int] = None,
+        enable_caching: bool = False,
+    ) -> None:
+        self.block_size = block_size
+        self.num_total_gpu_blocks = num_gpu_blocks
+        self.num_total_cpu_blocks = num_cpu_blocks
+
+        self.sliding_window = sliding_window
+        # max_block_sliding_window is the max number of blocks that need to be
+        # allocated
+        self.max_block_sliding_window = None
+        if sliding_window is not None:
+            # +1 here because // rounds down
+            num_blocks = sliding_window // block_size + 1
+            # +1 here because the last block may not be full,
+            # and so the sequence stretches one more block at the beginning
+            # For example, if sliding_window is 3 and block_size is 4,
+            # we may need 2 blocks when the second block only holds 1 token.
+            self.max_block_sliding_window = num_blocks + 1
+
+        self.watermark = watermark
+        assert watermark >= 0.0
+
+        self.enable_caching = enable_caching
+
+        self.watermark_blocks = int(watermark * num_gpu_blocks)
+
+        self.block_allocator = CpuGpuBlockAllocator.create(
+            allocator_type="prefix_caching" if enable_caching else "naive",
+            num_gpu_blocks=num_gpu_blocks,
+            num_cpu_blocks=num_cpu_blocks,
+            block_size=block_size,
+        )
+
+        self.block_tables: Dict[SeqId, BlockTable] = {}
+        self.cross_block_tables: Dict[EncoderSeqId, BlockTable] = {}
+        self._warned_mm_namespace_requests = set[str]()
+        self._request_local_namespace: Dict[str, bytes] = {}
+        self._runtime_cache_namespace = self._build_runtime_cache_namespace()
+
+        self._computed_blocks_tracker = ComputedBlocksTracker(
+            self.block_allocator)
+        self._last_access_blocks_tracker = LastAccessBlocksTracker(
+            self.block_allocator)
+
+    def can_allocate(self,
+                     seq_group: SequenceGroup,
+                     num_lookahead_slots: int = 0) -> AllocStatus:
+        # FIXME(woosuk): Here we assume that all sequences in the group share
+        # the same prompt. This may not be true for preempted sequences.
+
+        check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
+
+        seq = seq_group.get_seqs(status=SequenceStatus.WAITING)[0]
+        num_required_blocks = BlockTable.get_num_required_blocks(
+            seq.get_token_ids(),
+            block_size=self.block_size,
+            num_lookahead_slots=num_lookahead_slots,
+        )
+
+        if seq_group.is_encoder_decoder():
+            encoder_seq = seq_group.get_encoder_seq()
+            assert encoder_seq is not None
+            num_required_blocks += BlockTable.get_num_required_blocks(
+                encoder_seq.get_token_ids(),
+                block_size=self.block_size,
+            )
+
+        if self.max_block_sliding_window is not None:
+            num_required_blocks = min(num_required_blocks,
+                                      self.max_block_sliding_window)
+
+        num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
+            device=Device.GPU)
+
+        # Use watermark to avoid frequent cache eviction.
+        if (self.num_total_gpu_blocks - num_required_blocks <
+                self.watermark_blocks):
+            return AllocStatus.NEVER
+        if num_free_gpu_blocks - num_required_blocks >= self.watermark_blocks:
+            return AllocStatus.OK
+        else:
+            return AllocStatus.LATER
+
+    def _allocate_sequence(
+        self,
+        seq: Sequence,
+        cache_namespace: Optional[bytes] = None,
+    ) -> BlockTable:
+        block_table = BlockTable(
+            block_size=self.block_size,
+            block_allocator=self.block_allocator,
+            max_block_sliding_window=self.max_block_sliding_window,
+            cache_namespace=cache_namespace,
+        )
+        if seq.get_token_ids():
+            # Add blocks to the block table only if the sequence is non empty.
+            block_table.allocate(seq.get_token_ids())
+
+        return block_table
+
+    def allocate(self, seq_group: SequenceGroup) -> None:
+
+        # Allocate self-attention block tables for decoder sequences
+        waiting_seqs = seq_group.get_seqs(status=SequenceStatus.WAITING)
+        assert not (set(seq.seq_id for seq in waiting_seqs)
+                    & self.block_tables.keys()), "block table already exists"
+
+        # NOTE: Here we assume that all sequences in the group have the same
+        # prompt.
+        seq = waiting_seqs[0]
+        request_id = seq_group.request_id
+        cache_namespace = self._get_cache_namespace(
+            seq,
+            request_id=request_id,
+            seq_group=seq_group,
+        )
+        block_table: BlockTable = self._allocate_sequence(
+            seq,
+            cache_namespace=cache_namespace,
+        )
+        self.block_tables[seq.seq_id] = block_table
+
+        # Track seq
+        self._computed_blocks_tracker.add_seq(seq.seq_id)
+        self._last_access_blocks_tracker.add_seq(seq.seq_id)
+
+        # Assign the block table for each sequence.
+        for seq in waiting_seqs[1:]:
+            self.block_tables[seq.seq_id] = block_table.fork()
+
+            # Track seq
+            self._computed_blocks_tracker.add_seq(seq.seq_id)
+            self._last_access_blocks_tracker.add_seq(seq.seq_id)
+
+        # Allocate cross-attention block table for encoder sequence
+        #
+        # NOTE: Here we assume that all sequences in the group have the same
+        # encoder prompt.
+        request_id = seq_group.request_id
+
+        assert (request_id
+                not in self.cross_block_tables), \
+            "block table already exists"
+
+        check_no_caching_or_swa_for_blockmgr_encdec(self, seq_group)
+
+        if seq_group.is_encoder_decoder():
+            encoder_seq = seq_group.get_encoder_seq()
+            assert encoder_seq is not None
+            encoder_cache_namespace = self._get_cache_namespace(
+                encoder_seq,
+                request_id=request_id,
+                seq_group=seq_group)
+            block_table = self._allocate_sequence(
+                encoder_seq, cache_namespace=encoder_cache_namespace)
+            self.cross_block_tables[request_id] = block_table
+
+    @staticmethod
+    def _has_multi_modal_payload(multi_modal_data: Any) -> bool:
+        if multi_modal_data is None:
+            return False
+        if isinstance(multi_modal_data, Mapping):
+            try:
+                return len(multi_modal_data) > 0
+            except (TypeError, ValueError, RuntimeError, OSError,
+                    OverflowError, AttributeError, LookupError, struct.error):
+                # Treat an unusual mapping as payload and let normalization
+                # either identify it or select request-local isolation.
+                return True
+        return True
+
+    def _get_cache_namespace(self, seq: Sequence, request_id: str,
+                             seq_group: SequenceGroup) -> bytes:
+        digest = hashlib.sha256()
+        digest.update(b"bi100-request-prefix-namespace-v1|")
+        digest.update(self._runtime_cache_namespace)
+        digest.update(self._adapter_cache_namespace(seq_group))
+
+        multi_modal_data = seq.multi_modal_data
+        if self._has_multi_modal_payload(multi_modal_data):
+            try:
+                mm_namespace = self._hash_multi_modal_namespace(
+                    multi_modal_data)
+            except (TypeError, ValueError, RuntimeError, OSError,
+                    OverflowError, AttributeError, LookupError, struct.error):
+                if request_id not in self._warned_mm_namespace_requests:
+                    logger.warning(
+                        "Request %s has multimodal input that cannot be "
+                        "normalized for cache namespace hashing. Falling "
+                        "back to "
+                        "request-local namespace isolation.",
+                        request_id,
+                    )
+                    self._warned_mm_namespace_requests.add(request_id)
+                mm_namespace = self._request_local_fallback_cache_namespace(
+                    request_id=request_id)
+            digest.update(b"mm|")
+            digest.update(mm_namespace)
+        else:
+            digest.update(b"text|")
+        return digest.digest()
+
+    def can_append_slots(self, seq_group: SequenceGroup,
+                         num_lookahead_slots: int) -> bool:
+        """Determine if there is enough space in the GPU KV cache to continue
+        generation of the specified sequence group.
+
+        We use a worst-case heuristic: assume each touched block will require a
+        new allocation (either via CoW or new block). We can append slots if the
+        number of touched blocks is less than the number of free blocks.
+
+        "Lookahead slots" are slots that are allocated in addition to the slots
+        for known tokens. The contents of the lookahead slots are not defined.
+        This is used by speculative decoding when speculating future tokens.
+        """
+
+        num_touched_blocks = 0
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            block_table = self.block_tables[seq.seq_id]
+
+            num_touched_blocks += (
+                block_table.get_num_blocks_touched_by_append_slots(
+                    token_ids=block_table.get_unseen_token_ids(
+                        seq.get_token_ids()),
+                    num_lookahead_slots=num_lookahead_slots,
+                ))
+
+        num_free_gpu_blocks = self.block_allocator.get_num_free_blocks(
+            Device.GPU)
+        return num_touched_blocks <= num_free_gpu_blocks
+
+    def append_slots(
+        self,
+        seq: Sequence,
+        num_lookahead_slots: int,
+    ) -> List[Tuple[int, int]]:
+
+        block_table = self.block_tables[seq.seq_id]
+
+        block_table.append_token_ids(
+            token_ids=block_table.get_unseen_token_ids(seq.get_token_ids()),
+            num_lookahead_slots=num_lookahead_slots,
+            num_computed_slots=seq.data.get_num_computed_tokens(),
+        )
+        # Return any new copy-on-writes.
+        new_cows = self.block_allocator.clear_copy_on_writes()
+        return new_cows
+
+    def free(self, seq: Sequence) -> None:
+        seq_id = seq.seq_id
+
+        if seq_id not in self.block_tables:
+            # Already freed or haven't been scheduled yet.
+            return
+
+        # Update seq block ids with the latest access time
+        self._last_access_blocks_tracker.update_seq_blocks_last_access(
+            seq_id, self.block_tables[seq.seq_id].physical_block_ids)
+
+        # Untrack seq
+        self._last_access_blocks_tracker.remove_seq(seq_id)
+        self._computed_blocks_tracker.remove_seq(seq_id)
+
+        # Free table/blocks
+        self.block_tables[seq_id].free()
+        del self.block_tables[seq_id]
+
+    def free_cross(self, seq_group: SequenceGroup) -> None:
+        request_id = seq_group.request_id
+        if request_id not in self.cross_block_tables:
+            # Already freed or hasn't been scheduled yet.
+            return
+        self.cross_block_tables[request_id].free()
+        del self.cross_block_tables[request_id]
+
+    def get_block_table(self, seq: Sequence) -> List[int]:
+        block_ids = self.block_tables[seq.seq_id].physical_block_ids
+        return block_ids  # type: ignore
+
+    def get_cross_block_table(self, seq_group: SequenceGroup) -> List[int]:
+        request_id = seq_group.request_id
+        assert request_id in self.cross_block_tables
+        block_ids = self.cross_block_tables[request_id].physical_block_ids
+        assert all(b is not None for b in block_ids)
+        return block_ids  # type: ignore
+
+    def access_all_blocks_in_seq(self, seq: Sequence, now: float):
+        if self.enable_caching:
+            # Record the latest access time for the sequence. The actual update
+            # of the block ids is deferred to the sequence free(..) call, since
+            # only during freeing of block ids, the blocks are actually added to
+            # the evictor (which is when the most updated time is required)
+            # (This avoids expensive calls to mark_blocks_as_accessed(..))
+            self._last_access_blocks_tracker.update_last_access(
+                seq.seq_id, now)
+
+    def mark_blocks_as_computed(self, seq_group: SequenceGroup,
+                                token_chunk_size: int):
+        # If prefix caching is enabled, mark immutable blocks as computed
+        # right after they have been scheduled (for prefill). This assumes
+        # the scheduler is synchronous so blocks are actually computed when
+        # scheduling the next batch.
+        self.block_allocator.mark_blocks_as_computed([])
+
+    def get_common_computed_block_ids(
+            self, seqs: List[Sequence]) -> GenericSequence[int]:
+        """Determine which blocks for which we skip prefill.
+
+        With prefix caching we can skip prefill for previously-generated blocks.
+        Currently, the attention implementation only supports skipping cached
+        blocks if they are a contiguous prefix of cached blocks.
+
+        This method determines which blocks can be safely skipped for all
+        sequences in the sequence group.
+        """
+        computed_seq_block_ids = []
+        for seq in seqs:
+            computed_seq_block_ids.append(
+                self._computed_blocks_tracker.
+                get_cached_computed_blocks_and_update(
+                    seq.seq_id,
+                    self.block_tables[seq.seq_id].physical_block_ids))
+
+        # NOTE(sang): This assumes seq_block_ids doesn't contain any None.
+        return self.block_allocator.get_common_computed_block_ids(
+            computed_seq_block_ids)  # type: ignore
+
+    def get_content_hashes(self, seq: Sequence) -> List[bytes]:
+        return self.block_tables[seq.seq_id].get_content_hashes()
+
+    def get_and_reset_prefix_swaps(
+            self) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]]]:
+        """Return scheduler-owned (CPU->GPU, GPU->CPU) content transfers."""
+        return self.block_allocator.get_and_reset_prefix_swaps()
+
+    def begin_prefix_cache_step(self) -> None:
+        self.block_allocator.begin_prefix_cache_step()
+
+    def _build_runtime_cache_namespace(self) -> bytes:
+        """Bind first-block hashes to the fixed model runtime identity."""
+        model = os.getenv("BI100_PREFIX_MODEL_FINGERPRINT",
+                          "Qwen3.6-35B-A3B")
+        dtype = os.getenv("BI100_PREFIX_DTYPE", "float16")
+        tp_raw = os.getenv("BI100_PREFIX_TP_SIZE", "4")
+        try:
+            tp_size = int(tp_raw)
+        except ValueError as exc:
+            raise RuntimeError(
+                "BI100_PREFIX_TP_SIZE must be a positive integer") from exc
+        if tp_size <= 0:
+            raise RuntimeError(
+                "BI100_PREFIX_TP_SIZE must be a positive integer")
+
+        digest = hashlib.sha256()
+        digest.update(b"bi100-runtime-prefix-identity-v1|")
+        for label, value in (
+                (b"model", model),
+                (b"dtype", dtype),
+                (b"tp", str(tp_size)),
+                (b"block_size", str(self.block_size))):
+            encoded = value.encode("utf-8")
+            digest.update(label)
+            digest.update(struct.pack("!Q", len(encoded)))
+            digest.update(encoded)
+        return digest.digest()
+
+    @staticmethod
+    def _adapter_cache_namespace(seq_group: SequenceGroup) -> bytes:
+        digest = hashlib.sha256()
+        digest.update(b"bi100-adapter-prefix-identity-v1|")
+        lora = getattr(seq_group, "lora_request", None)
+        prompt_adapter = getattr(seq_group, "prompt_adapter_request", None)
+        identities = (
+            ("lora", lora, ("lora_name", "lora_int_id", "lora_path",
+                             "base_model_name")),
+            ("prompt", prompt_adapter,
+             ("prompt_adapter_name", "prompt_adapter_id",
+              "prompt_adapter_local_path",
+              "prompt_adapter_num_virtual_tokens")),
+        )
+        for kind, adapter, fields in identities:
+            digest.update(kind.encode("ascii"))
+            if adapter is None:
+                digest.update(b"none|")
+                continue
+            for field in fields:
+                value = str(getattr(adapter, field, ""))
+                encoded = value.encode("utf-8")
+                digest.update(field.encode("ascii"))
+                digest.update(struct.pack("!Q", len(encoded)))
+                digest.update(encoded)
+        return digest.digest()
+
+    def _request_local_fallback_cache_namespace(self,
+                                               request_id: str) -> bytes:
+        namespace = self._request_local_namespace.get(request_id)
+        if namespace is None:
+            digest = hashlib.sha256()
+            digest.update(b"multimodal-unsupported-request-local-v1|")
+            digest.update(self._runtime_cache_namespace)
+            digest.update(os.urandom(32))
+            digest.update(request_id.encode("utf-8"))
+            namespace = digest.digest()
+            self._request_local_namespace[request_id] = namespace
+        return namespace
+
+    def release_request_cache_namespace(self, request_id: str) -> None:
+        """Release request-local isolation state after request completion."""
+        self._request_local_namespace.pop(request_id, None)
+        self._warned_mm_namespace_requests.discard(request_id)
+
+    def _hash_multi_modal_namespace(self, mm_data: Any) -> bytes:
+        digest = hashlib.sha256()
+        self._hash_multi_modal_obj(digest, mm_data)
+        return digest.digest()
+
+    @staticmethod
+    def _sort_map_keys(mm_map: Mapping[Any, Any]) -> List[Any]:
+        return sorted(mm_map.keys(), key=lambda key: repr(key))
+
+    @classmethod
+    def _hash_multi_modal_obj(cls, digest: Any, value: Any) -> None:
+        if value is None:
+            digest.update(b"none|")
+            return
+        if isinstance(value, Mapping):
+            digest.update(b"map|")
+            digest.update(struct.pack("!Q", len(value)))
+            for key in cls._sort_map_keys(value):
+                digest.update(b"k|")
+                cls._hash_multi_modal_obj(digest, key)
+                digest.update(b"v|")
+                cls._hash_multi_modal_obj(digest, value[key])
+            return
+        if isinstance(value, list):
+            digest.update(b"list|")
+            digest.update(struct.pack("!Q", len(value)))
+            for item in value:
+                cls._hash_multi_modal_obj(digest, item)
+            return
+        if isinstance(value, tuple):
+            digest.update(b"tuple|")
+            digest.update(struct.pack("!Q", len(value)))
+            for item in value:
+                cls._hash_multi_modal_obj(digest, item)
+            return
+        if isinstance(value, str):
+            encoded = value.encode()
+            digest.update(b"str|")
+            digest.update(struct.pack("!Q", len(encoded)))
+            digest.update(encoded)
+            return
+        if isinstance(value, bytes):
+            digest.update(b"bytes|")
+            digest.update(struct.pack("!Q", len(value)))
+            digest.update(value)
+            return
+        if isinstance(value, bytearray):
+            cls._hash_multi_modal_obj(digest, bytes(value))
+            return
+        if isinstance(value, bool):
+            digest.update(b"bool|")
+            digest.update(b"1" if value else b"0")
+            return
+        if isinstance(value, int):
+            digest.update(b"int|")
+            digest.update(str(value).encode())
+            return
+        if isinstance(value, float):
+            digest.update(b"float|")
+            digest.update(struct.pack("!d", value))
+            return
+        if torch is not None and isinstance(value, torch.Tensor):
+            digest.update(b"tensor|")
+            tensor = value.detach().cpu().contiguous()
+            digest.update(struct.pack("!Q", len(tensor.shape)))
+            for dim in tensor.shape:
+                digest.update(struct.pack("!Q", int(dim)))
+            digest.update(str(tensor.dtype).encode())
+            # Byte views work for bfloat16 and other dtypes that NumPy cannot
+            # materialize directly.
+            tensor_bytes = tensor.view(torch.uint8).numpy().tobytes()
+            digest.update(struct.pack("!Q", len(tensor_bytes)))
+            digest.update(tensor_bytes)
+            return
+        if Image is not None and isinstance(value, Image.Image):
+            digest.update(b"image|")
+            digest.update(value.mode.encode())
+            digest.update(struct.pack("!II", value.width, value.height))
+            image_bytes = value.tobytes()
+            digest.update(struct.pack("!Q", len(image_bytes)))
+            digest.update(image_bytes)
+            palette = value.getpalette()
+            digest.update(b"palette-mode|")
+            cls._hash_multi_modal_obj(
+                digest, getattr(getattr(value, "palette", None), "mode", None))
+            digest.update(b"palette|")
+            cls._hash_multi_modal_obj(digest, palette)
+            digest.update(b"transparency|")
+            cls._hash_multi_modal_obj(
+                digest, value.info.get("transparency"))
+            return
+
+        raise TypeError(f"Unsupported multimodal namespace value type {type(value)}")
+
+
+    def fork(self, parent_seq: Sequence, child_seq: Sequence) -> None:
+        if parent_seq.seq_id not in self.block_tables:
+            # Parent sequence has either been freed or never existed.
+            return
+        src_block_table = self.block_tables[parent_seq.seq_id]
+        self.block_tables[child_seq.seq_id] = src_block_table.fork()
+
+        # Track child seq
+        self._computed_blocks_tracker.add_seq(child_seq.seq_id)
+        self._last_access_blocks_tracker.add_seq(child_seq.seq_id)
+
+    def can_swap_in(self, seq_group: SequenceGroup,
+                    num_lookahead_slots: int) -> AllocStatus:
+        """Returns the AllocStatus for the given sequence_group 
+        with num_lookahead_slots.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            AllocStatus: The AllocStatus for the given sequence group.
+        """
+        if self.block_allocator.content_offload_enabled:
+            return AllocStatus.NEVER
+        return self._can_swap(seq_group, Device.GPU, SequenceStatus.SWAPPED,
+                              num_lookahead_slots)
+
+    def swap_in(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
+        """Returns the block id mapping (from CPU to GPU) generated by
+        swapping in the given seq_group with num_lookahead_slots.
+
+        Args:
+            seq_group (SequenceGroup): The sequence group to swap in.
+
+        Returns:
+            List[Tuple[int, int]]: The mapping of swapping block from CPU 
+                to GPU.
+        """
+        physical_block_id_mapping = []
+        for seq in seq_group.get_seqs(status=SequenceStatus.SWAPPED):
+            blocks = self.block_tables[seq.seq_id].blocks
+            if len(blocks) == 0:
+                continue
+
+            seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
+                                                         src_device=Device.CPU,
+                                                         dst_device=Device.GPU)
+
+            # Refresh the block ids of the table (post-swap)
+            self.block_tables[seq.seq_id].update(blocks)
+
+            seq_physical_block_id_mapping = {
+                self.block_allocator.get_physical_block_id(
+                    Device.CPU, cpu_block_id):
+                self.block_allocator.get_physical_block_id(
+                    Device.GPU, gpu_block_id)
+                for cpu_block_id, gpu_block_id in seq_swap_mapping.items()
+            }
+
+            physical_block_id_mapping.extend(
+                list(seq_physical_block_id_mapping.items()))
+
+        return physical_block_id_mapping
+
+    def can_swap_out(self, seq_group: SequenceGroup) -> bool:
+        """Returns whether we can swap out the given sequence_group 
+        with num_lookahead_slots.
+
+        Args:
+            seq_group (SequenceGroup): The sequence group to swap in.
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            bool: Whether it's possible to swap out current sequence group.
+        """
+        if self.block_allocator.content_offload_enabled:
+            return False
+        alloc_status = self._can_swap(seq_group, Device.CPU,
+                                      SequenceStatus.RUNNING)
+        return alloc_status == AllocStatus.OK
+
+    def swap_out(self, seq_group: SequenceGroup) -> List[Tuple[int, int]]:
+        """Returns the block id mapping (from GPU to CPU) generated by
+        swapping out the given sequence_group with num_lookahead_slots.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+
+        Returns:
+            List[Tuple[int, int]]: The mapping of swapping block from 
+                GPU to CPU.
+        """
+        physical_block_id_mapping = []
+        for seq in seq_group.get_seqs(status=SequenceStatus.RUNNING):
+            blocks = self.block_tables[seq.seq_id].blocks
+            if len(blocks) == 0:
+                continue
+
+            seq_swap_mapping = self.block_allocator.swap(blocks=blocks,
+                                                         src_device=Device.GPU,
+                                                         dst_device=Device.CPU)
+
+            # Refresh the block ids of the table (post-swap)
+            self.block_tables[seq.seq_id].update(blocks)
+
+            seq_physical_block_id_mapping = {
+                self.block_allocator.get_physical_block_id(
+                    Device.GPU, gpu_block_id):
+                self.block_allocator.get_physical_block_id(
+                    Device.CPU, cpu_block_id)
+                for gpu_block_id, cpu_block_id in seq_swap_mapping.items()
+            }
+
+            physical_block_id_mapping.extend(
+                list(seq_physical_block_id_mapping.items()))
+
+        return physical_block_id_mapping
+
+    def get_num_free_gpu_blocks(self) -> int:
+        return self.block_allocator.get_num_free_blocks(Device.GPU)
+
+    def get_num_free_cpu_blocks(self) -> int:
+        return self.block_allocator.get_num_free_blocks(Device.CPU)
+
+    def get_prefix_cache_hit_rate(self, device: Device) -> float:
+        return self.block_allocator.get_prefix_cache_hit_rate(device)
+
+    def _can_swap(self,
+                  seq_group: SequenceGroup,
+                  device: Device,
+                  status: SequenceStatus,
+                  num_lookahead_slots: int = 0) -> AllocStatus:
+        """Returns the AllocStatus for swapping in/out the given sequence_group 
+        on to the 'device'.
+
+        Args:
+            sequence_group (SequenceGroup): The sequence group to swap in.
+            device (Device): device to swap the 'seq_group' on.
+            status (SequenceStatus): The status of sequence which is needed
+                for action. RUNNING for swap out and SWAPPED for swap in
+            num_lookahead_slots (int): Number of lookahead slots used in 
+                speculative decoding, default to 0.
+
+        Returns:
+            AllocStatus: The AllocStatus for swapping in/out the given 
+                sequence_group on to the 'device'.
+        """
+        # First determine the number of blocks that will be touched by this
+        # swap. Then verify if there are available blocks in the device
+        # to perform the swap.
+        num_blocks_touched = 0
+        blocks: List[Block] = []
+        for seq in seq_group.get_seqs(status=status):
+            block_table = self.block_tables[seq.seq_id]
+            if block_table.blocks is not None:
+                # Compute the number blocks to touch for the tokens to be
+                # appended. This does NOT include the full blocks that need
+                # to be touched for the swap.
+                num_blocks_touched += \
+                    block_table.get_num_blocks_touched_by_append_slots(
+                        block_table.get_unseen_token_ids(seq.get_token_ids()),
+                        num_lookahead_slots=num_lookahead_slots)
+                blocks.extend(block_table.blocks)
+        # Compute the number of full blocks to touch and add it to the
+        # existing count of blocks to touch.
+        num_blocks_touched += self.block_allocator.get_num_full_blocks_touched(
+            blocks, device=device)
+
+        watermark_blocks = 0
+        if device == Device.GPU:
+            watermark_blocks = self.watermark_blocks
+
+        if self.block_allocator.get_num_total_blocks(
+                device) < num_blocks_touched:
+            return AllocStatus.NEVER
+        elif self.block_allocator.get_num_free_blocks(
+                device) - num_blocks_touched >= watermark_blocks:
+            return AllocStatus.OK
+        else:
+            return AllocStatus.LATER

qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py

@@ -0,0 +1,272 @@
+import enum
+import heapq
+import os
+from abc import ABC, abstractmethod
+from collections.abc import Mapping
+from typing import Dict, List, OrderedDict, Tuple
+
+
+ContentHash = bytes
+
+
+class EvictionPolicy(enum.Enum):
+    """Enum for eviction policy used by make_evictor to instantiate the correct
+       Evictor subclass.
+    """
+    LRU = enum.auto()
+    FREQUENCY_AWARE = enum.auto()
+
+
+class Evictor(ABC):
+    """The Evictor subclasses should be used by the BlockAllocator class to
+    handle eviction of freed PhysicalTokenBlocks.
+    """
+
+    @abstractmethod
+    def __init__(self):
+        pass
+
+    @abstractmethod
+    def __contains__(self, block_id: int) -> bool:
+        pass
+
+    @abstractmethod
+    def evict(self) -> Tuple[int, ContentHash]:
+        """Runs the eviction algorithm and returns the evicted block's
+        content hash along with physical block id along with physical block id
+        """
+        pass
+
+    @abstractmethod
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int,
+            last_accessed: float):
+        """Adds block to the evictor, making it a candidate for eviction"""
+        pass
+
+    @abstractmethod
+    def update(self, block_id: int, last_accessed: float):
+        """Update corresponding block's access time in metadata"""
+        pass
+
+    @abstractmethod
+    def remove(self, block_id: int):
+        """Remove a given block id from the cache."""
+        pass
+
+    @property
+    @abstractmethod
+    def num_blocks(self) -> int:
+        pass
+
+
+class BlockMetaData():
+    """Data structure for storing key data describe cached block, so that
+    evitor could use to make its decision which one to choose for eviction
+
+    Here we use physical block id as the dict key, as there maybe several
+    blocks with the same content hash, but their physical id is unique.
+    """
+
+    def __init__(self, content_hash: ContentHash, num_hashed_tokens: int,
+                 last_accessed: float):
+        self.content_hash = content_hash
+        self.num_hashed_tokens = num_hashed_tokens
+        self.last_accessed = last_accessed
+
+
+class LRUEvictor(Evictor):
+    """Evicts in a least-recently-used order using the last_accessed timestamp
+    that's recorded in the PhysicalTokenBlock. If there are multiple blocks with
+    the same last_accessed time, then the one with the largest num_hashed_tokens
+    will be evicted. If two blocks each have the lowest last_accessed time and
+    highest num_hashed_tokens value, then one will be chose arbitrarily
+    """
+
+    def __init__(self):
+        self.free_table: OrderedDict[int, BlockMetaData] = OrderedDict()
+
+    def __contains__(self, block_id: int) -> bool:
+        return block_id in self.free_table
+
+    def evict(self) -> Tuple[int, ContentHash]:
+        if len(self.free_table) == 0:
+            raise ValueError("No usable cache memory left")
+
+        evicted_block, evicted_block_id = None, None
+        # The blocks with the lowest timestamps should be placed consecutively
+        # at the start of OrderedDict. Loop through all these blocks to
+        # find the one with maximum number of hashed tokens.
+        for _id, block in self.free_table.items():
+            if evicted_block is None:
+                evicted_block, evicted_block_id = block, _id
+                continue
+            if evicted_block.last_accessed < block.last_accessed:
+                break
+            if evicted_block.num_hashed_tokens < block.num_hashed_tokens:
+                evicted_block, evicted_block_id = block, _id
+
+        assert evicted_block is not None
+        assert evicted_block_id is not None
+        self.free_table.pop(evicted_block_id)
+
+        return evicted_block_id, evicted_block.content_hash
+
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int,
+            last_accessed: float):
+        self.free_table[block_id] = BlockMetaData(content_hash,
+                                                  num_hashed_tokens,
+                                                  last_accessed)
+
+    def update(self, block_id: int, last_accessed: float):
+        self.free_table[block_id].last_accessed = last_accessed
+
+    def remove(self, block_id: int):
+        if block_id not in self.free_table:
+            raise ValueError(
+                "Attempting to remove block that's not in the evictor")
+        self.free_table.pop(block_id)
+
+    @property
+    def num_blocks(self) -> int:
+        return len(self.free_table)
+
+
+class FrequencyAwareEvictor(Evictor):
+    """Evict the least frequently reused logical prefix content first.
+
+    Content frequency survives physical block reuse. Heap entries carry a
+    generation and are lazily invalidated so eviction remains O(log N) without
+    allowing stale entries to grow without bound.
+    """
+
+    _COMPACTION_FACTOR = 2
+    _COMPACTION_SLACK = 1
+
+    def __init__(self):
+        self.free_table: Dict[int, BlockMetaData] = {}
+        self.frequency_by_hash: Dict[ContentHash, int] = {}
+        self._heap: List[Tuple[int, float, int, int, int]] = []
+        self._generations: Dict[int, int] = {}
+        self._next_generation = 0
+
+    @staticmethod
+    def _validate_content_hash(content_hash: ContentHash) -> None:
+        if not isinstance(content_hash, bytes) or len(content_hash) != 32:
+            raise ValueError(
+                "frequency-aware eviction requires a 32-byte content hash")
+
+    def __contains__(self, block_id: int) -> bool:
+        return block_id in self.free_table
+
+    def _heap_key(self, block_id: int, block: BlockMetaData,
+                  generation: int) -> Tuple[int, float, int, int, int]:
+        return (
+            self.frequency_by_hash[block.content_hash],
+            block.last_accessed,
+            -block.num_hashed_tokens,
+            block_id,
+            generation,
+        )
+
+    def _push(self, block_id: int) -> None:
+        self._next_generation += 1
+        generation = self._next_generation
+        self._generations[block_id] = generation
+        heapq.heappush(
+            self._heap,
+            self._heap_key(
+                block_id, self.free_table[block_id], generation),
+        )
+
+    def _compact_if_needed(self) -> None:
+        limit = (
+            self._COMPACTION_FACTOR * len(self.free_table)
+            + self._COMPACTION_SLACK
+        )
+        if len(self._heap) <= limit:
+            return
+        self._heap = [
+            self._heap_key(block_id, block, self._generations[block_id])
+            for block_id, block in self.free_table.items()
+        ]
+        heapq.heapify(self._heap)
+
+    def evict(self) -> Tuple[int, ContentHash]:
+        if not self.free_table:
+            raise ValueError("No usable cache memory left")
+
+        while self._heap:
+            entry = heapq.heappop(self._heap)
+            frequency, _, _, block_id, generation = entry
+            block = self.free_table.get(block_id)
+            if (
+                block is None
+                or self._generations.get(block_id) != generation
+            ):
+                continue
+            if self.frequency_by_hash[block.content_hash] != frequency:
+                heapq.heappush(
+                    self._heap,
+                    self._heap_key(block_id, block, generation),
+                )
+                continue
+
+            block = self.free_table.pop(block_id)
+            self._generations.pop(block_id)
+            self._compact_if_needed()
+            return block_id, block.content_hash
+
+        raise RuntimeError("Evictor heap has no usable entry")
+
+    def add(self, block_id: int, content_hash: ContentHash,
+            num_hashed_tokens: int, last_accessed: float):
+        self._validate_content_hash(content_hash)
+        self.frequency_by_hash[content_hash] = (
+            self.frequency_by_hash.get(content_hash, 0) + 1)
+        self.free_table[block_id] = BlockMetaData(
+            content_hash, num_hashed_tokens, last_accessed)
+        self._push(block_id)
+        self._compact_if_needed()
+
+    def update(self, block_id: int, last_accessed: float):
+        self.free_table[block_id].last_accessed = last_accessed
+        self._push(block_id)
+        self._compact_if_needed()
+
+    def remove(self, block_id: int):
+        if block_id not in self.free_table:
+            raise ValueError(
+                "Attempting to remove block that's not in the evictor")
+        self.free_table.pop(block_id)
+        self._generations.pop(block_id)
+        self._compact_if_needed()
+
+    @property
+    def num_blocks(self) -> int:
+        return len(self.free_table)
+
+
+def eviction_policy_from_env(
+    environ: Mapping[str, str] | None = None,
+) -> EvictionPolicy:
+    source = os.environ if environ is None else environ
+    value = source.get("BI100_KV_EVICTION_POLICY", "lru").strip().lower()
+    policies = {
+        "lru": EvictionPolicy.LRU,
+        "frequency": EvictionPolicy.FREQUENCY_AWARE,
+    }
+    if value not in policies:
+        raise ValueError(
+            "BI100_KV_EVICTION_POLICY must be one of: frequency, lru")
+    return policies[value]
+
+
+def make_evictor(eviction_policy: EvictionPolicy) -> Evictor:
+    if eviction_policy == EvictionPolicy.LRU:
+        return LRUEvictor()
+    elif eviction_policy == EvictionPolicy.FREQUENCY_AWARE:
+        return FrequencyAwareEvictor()
+    else:
+        raise ValueError(f"Unknown cache eviction policy: {eviction_policy}")

qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py

@@ -331,33 +331,9 @@ def _get_bin_counts_and_mask(
 ) -> Tuple[torch.Tensor, torch.Tensor]:
     # Compute the bin counts for the tokens.
     # vocab_size + 1 for padding.
-    #
-    # CCCL bit_packed_counter pattern (catch2_test_memcpy_bitpacked_counter.cu):
-    # Pack counters using minimum bits needed. Original code uses int64
-    # (8 bytes per counter), but token repetition counts in a single
-    # generation never exceed a few hundred. We keep int64 for scatter_add_
-    # compatibility but pre-allocate once to avoid per-step CUDA malloc.
-    #
-    # CCCL dispatch_reduce.cuh alias_temporaries: pre-allocate, reuse.
-    # For Qwen3.6 (vocab=152064, batch=8 decode):
-    #   bin_counts = 8 × 152065 × 8 = 9.7 MB, allocated ONCE, reused.
-    #   scatter_add_ requires int64 on CUDA, so dtype cannot change.
-    #
-    # Future: if scatter_add_ supports int16/int32, switch to reduce 4x.
-    _cache_key = ("bin_counts", vocab_size, num_seqs, tokens.device)
-    global _sampler_cache
-    if '_sampler_cache' not in dir():
-        _sampler_cache = {}
-    cached = _sampler_cache.get(_cache_key)
-    if cached is not None and cached.shape == (num_seqs, vocab_size + 1):
-        bin_counts = cached
-        bin_counts.zero_()
-    else:
-        bin_counts = torch.zeros((num_seqs, vocab_size + 1),
-                                 dtype=torch.long,
-                                 device=tokens.device)
-        _sampler_cache[_cache_key] = bin_counts
-
+    bin_counts = torch.zeros((num_seqs, vocab_size + 1),
+                             dtype=torch.long,
+                             device=tokens.device)
     bin_counts.scatter_add_(1, tokens, torch.ones_like(tokens))
     bin_counts = bin_counts[:, :vocab_size]
     mask = bin_counts > 0
@@ -440,61 +416,7 @@ def _apply_top_k_top_p(
     p: torch.Tensor,
     k: torch.Tensor,
 ) -> torch.Tensor:
-    # CCCL dispatch_topk.cuh architecture (480 lines, full read):
-    #
-    # 1. Multi-pass radix selection: O(N × bits_per_pass) not O(N log N)
-    #    pass 0: DeviceTopKHistogramKernel (histogram only, no filter)
-    #    pass 1+: DeviceTopKKernel (fused filter + histogram)
-    #    last: DeviceTopKLastFilterKernel (filter only)
-    #
-    # 2. DoubleBuffer<key_in_t> pattern (dispatch_topk.cuh line ~430):
-    #    key_bufs = DoubleBuffer(alloc[3], alloc[2])  // ping-pong
-    #    for pass: use Current() as input, Alternate() as output, then swap
-    #    → zero allocation in the hot loop
-    #
-    # 3. candidate_buffer_length = num_items / 128
-    #    Only 1/128 of input needs buffer space for candidates
-    #    vocab=152064 → 1188 candidates max
-    #
-    # PyTorch translation below uses pre-allocated buffers where possible
-    # to avoid per-step allocation overhead (BI-V100 has no async allocator).
-
-    # Fast path: when ALL sequences use top_p=1.0 (no nucleus sampling),
-    # we only need top-k selection, not full sort + cumsum.
-    all_top_p_disabled = (p >= 1.0 - 1e-6).all()
-    if all_top_p_disabled:
-        max_k = k.max().item()
-        if max_k > 0 and max_k < logits.size(1):
-            # CCCL DeviceTopK env API (catch2_test_device_topk_env_api.cu):
-            # output_ordering::unsorted — top-k results don't need sorting.
-            # torch.topk(sorted=False) skips the final sort step, saving
-            # ~10% of the radix select time. We only need the threshold
-            # value (min of top-k), not their ordering.
-            topk_vals, topk_idx = torch.topk(logits, int(max_k), dim=-1, sorted=False)
-            actual_k_mask = torch.arange(int(max_k), device=k.device).unsqueeze(0) < k.unsqueeze(1)
-            topk_vals.masked_fill_(~actual_k_mask, -float("inf"))
-            threshold = topk_vals.min(dim=-1, keepdim=True).values
-            logits = logits.masked_fill(logits < threshold, -float("inf"))
-            return logits
-
-    # Full path: sort + top-k + top-p (cumsum)
-    # CCCL DoubleBuffer insight: reuse sort output tensors across calls
-    # by caching them keyed on (batch_size, vocab_size, device).
-    # This avoids torch.sort allocating 2 new tensors (152064×4B each)
-    # on every single decode step.
-    _buf_key = (logits.shape[0], logits.shape[1], str(logits.device))
-    _bufs = getattr(_apply_top_k_top_p, '_sort_bufs', {}).get(_buf_key)
-    if _bufs is not None:
-        logits_sort, logits_idx = _bufs
-        # In-place sort into pre-allocated buffers
-        torch.sort(logits, dim=-1, descending=False, out=(logits_sort, logits_idx))
-    else:
-        logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
-        # Cache for next call (CCCL DoubleBuffer pattern)
-        if not hasattr(_apply_top_k_top_p, '_sort_bufs'):
-            _apply_top_k_top_p._sort_bufs = {}
-        _apply_top_k_top_p._sort_bufs[_buf_key] = (
-            logits_sort.clone(), logits_idx.clone())  # pre-alloc buffers
+    logits_sort, logits_idx = logits.sort(dim=-1, descending=False)
 
     # Apply top-k.
     top_k_mask = logits_sort.size(1) - k.to(torch.long)
@@ -1058,17 +980,12 @@ def get_logprobs(
 
         assert len(next_token_ids) == len(query_indices)
 
-    if len(query_indices) == 0:
-        empty_sampled_logprob: SampleLogprobs = []
-        empty_prompt_logprob: Optional[PromptLogprobs] = None
-        return [empty_prompt_logprob], [empty_sampled_logprob]
-
     selected_logprobs, ranks = None, None
     top_logprobs, top_token_ids = None, None
 
     # If largest_num_logprobs == -1, i.e. no logprobs are requested, we can
     # skip the whole logprob calculation.
-    if largest_num_logprobs >= 0:
+    if query_indices and largest_num_logprobs >= 0:
         query_indices_gpu = torch.tensor(query_indices, device=logprobs.device)
         next_token_ids_gpu = torch.tensor(next_token_ids,
                                           device=logprobs.device)
@@ -1137,18 +1054,45 @@ def _get_prompt_logprob_if_needed(
     # Find prompt logprobs
     prompt_logprobs: Optional[PromptLogprobs] = None
     if is_prompt and sampling_params.prompt_logprobs is not None:
-        prompt_logprobs = []
+        query_len = seq_group.query_len
+        assert query_len is not None
+        requested_prompt_logprob_len = (
+            query_len - len(seq_group.seq_ids)
+            if seq_group.do_sample else query_len)
+        seq_data = seq_group.seq_data[seq_group.seq_ids[0]]
+        available_next_tokens = max(
+            0,
+            len(seq_data.prompt_token_ids)
+            - seq_data.get_num_computed_tokens()
+            - 1,
+        )
+        full_prompt_logprob_len = min(
+            requested_prompt_logprob_len,
+            available_next_tokens,
+        )
+        prompt_logprobs = [None] * full_prompt_logprob_len
         num_logprobs = sampling_params.prompt_logprobs
         next_prompt_tokens = _get_next_prompt_tokens(seq_group)
+        assert (len(next_prompt_tokens)
+                == len(seq_group.prompt_logprob_indices)
+                == len(seq_group.prompt_logprob_output_indices))
         # Pre-select indexes and create a list. It is faster than calling .item
         # repetitively.
-        selected_logprob_items = selected_logprobs[
-            selected_logprobs_idx:selected_logprobs_idx +
-            len(next_prompt_tokens)].tolist()
-        rank_items = ranks[selected_logprobs_idx:selected_logprobs_idx +
-                           len(next_prompt_tokens)].tolist()
-
-        for idx, token_id in enumerate(next_prompt_tokens):
+        if next_prompt_tokens:
+            assert selected_logprobs is not None
+            assert ranks is not None
+            selected_logprob_items = selected_logprobs[
+                selected_logprobs_idx:selected_logprobs_idx +
+                len(next_prompt_tokens)].tolist()
+            rank_items = ranks[
+                selected_logprobs_idx:selected_logprobs_idx +
+                len(next_prompt_tokens)].tolist()
+        else:
+            selected_logprob_items = []
+            rank_items = []
+        for idx, (token_id, output_index) in enumerate(zip(
+                next_prompt_tokens,
+                seq_group.prompt_logprob_output_indices)):
             # Calculate the prompt logprob of the real prompt tokens.
             # {token_id: (logprob, rank_from_vocab)}
             prompt_logprobs_dict: Dict[int, Tuple[float, int]] = {
@@ -1157,6 +1101,8 @@ def _get_prompt_logprob_if_needed(
 
             # Add top K prompt logprobs along with its rank.
             if num_logprobs > 0:
+                assert top_token_ids is not None
+                assert top_logprobs is not None
                 top_ids = top_token_ids[
                     top_logprob_idx, :num_logprobs].tolist()
                 top_probs = top_logprobs[
@@ -1169,10 +1115,10 @@ def _get_prompt_logprob_if_needed(
                     for top_id, top_prob, rank in zip(top_ids, top_probs,
                                                       top_ranks)
                 })
-            prompt_logprobs.append({
+            prompt_logprobs[output_index] = {
                 token_id: Logprob(*logprob_and_rank)
                 for token_id, logprob_and_rank in prompt_logprobs_dict.items()
-            })
+            }
             # + 1 to go to the next prompt token.
             top_logprob_idx += 1
 
@@ -1386,10 +1332,9 @@ def _get_next_prompt_tokens(seq_group: SequenceGroupToSample) -> List[int]:
     seq_data = seq_group.seq_data[seq_ids[0]]
     computed_len = seq_data.get_num_computed_tokens()
     prompt_tokens = seq_data.prompt_token_ids
-    # +1 because we are looking for a next prompt token.
-    next_token_index_start = computed_len + 1
-    next_token_index_end = min(computed_len + query_len + 1,
-                               len(prompt_tokens))
-    next_prompt_tokens = prompt_tokens[
-        next_token_index_start:next_token_index_end]
+    next_prompt_tokens = []
+    for output_index in seq_group.prompt_logprob_output_indices:
+        token_index = computed_len + output_index + 1
+        assert token_index < len(prompt_tokens)
+        next_prompt_tokens.append(prompt_tokens[token_index])
     return next_prompt_tokens

qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py

@@ -0,0 +1,644 @@
+from array import array
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from vllm.sampling_params import SamplingParams, SamplingType
+from vllm.sequence import (VLLM_TOKEN_ID_ARRAY_TYPE, SequenceData,
+                           SequenceGroupMetadata)
+from vllm.utils import (PyObjectCache, async_tensor_h2d,
+                        is_pin_memory_available, make_tensor_with_pad)
+
+_SAMPLING_EPS = 1e-5
+
+
+@dataclass
+class SequenceGroupToSample:
+    # |---------- N-1 iteration --------|
+    # |---------------- N iteration ---------------------|
+    # |- tokenA -|......................|-- newTokens ---|
+    # |---------- context_len ----------|
+    # |-------------------- seq_len ----------------------|
+    #                                   |-- query_len ---|
+
+    # Sequence ids for the sequence group in a previous step.
+    seq_ids: List[int]
+    sampling_params: SamplingParams
+    # seq_id -> sequence data.
+    seq_data: Dict[int, SequenceData]
+    # The length of the sequence (all tokens seen in the past + new token to
+    # compute attention) of the sequence group. None if it is in a decode
+    # stage.
+    seq_len: Optional[int]
+    # The length of new query tokens to compute in the current step. None if it
+    # is in a decode stage. The length of query_len <= seq_len if chunked
+    # prefill is enabled.
+    query_len: Optional[int]
+    # A random number generator for sampling.
+    generator: Optional[torch.Generator]
+    # True if the sequence group is in prefill stage. False if it is in a
+    # decode stage.
+    is_prompt: bool
+    # Query token indices from logits. to compute prompt logprob. Empty if
+    # prompt logprob is not required.
+    prompt_logprob_indices: List[int]
+    # Output offsets within this prefill chunk. Sparse diagnostic requests use
+    # this to retain the standard full-length prompt-logprob response shape.
+    prompt_logprob_output_indices: List[int]
+    # Sample token indices from logits. Empty if sampling is not required.
+    sample_indices: List[int]
+
+    @property
+    def do_sample(self):
+        return len(self.sample_indices) > 0
+
+    def __post_init__(self):
+        if len(self.prompt_logprob_indices) > 0:
+            assert self.sampling_params.prompt_logprobs is not None
+        assert (len(self.prompt_logprob_indices)
+                == len(self.prompt_logprob_output_indices))
+        assert self.prompt_logprob_output_indices == sorted(
+            set(self.prompt_logprob_output_indices))
+        if self.is_prompt:
+            assert self.seq_len is not None
+            assert self.query_len is not None
+            assert all(
+                0 <= index < self.query_len
+                for index in self.prompt_logprob_output_indices)
+
+
+def gen_seq_group_to_sample_builder(num_seqs: int):
+    return lambda: SequenceGroupToSample(
+        seq_ids=[0] * num_seqs,
+        sampling_params=None,
+        seq_data=None,  # type: ignore
+        seq_len=0,
+        query_len=0,
+        generator=None,
+        is_prompt=True,
+        prompt_logprob_indices=[],
+        prompt_logprob_output_indices=[],
+        sample_indices=[],
+    )
+
+
+class SamplingMetadataCache:
+    """Used to cache SamplingMetadata objects between scheduler iterations"""
+
+    def __init__(self):
+        self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {}
+
+    def get_cached_seq_group_to_sample(self, num_seqs):
+        if num_seqs not in self._seq_group_to_sample_cache:
+            self._seq_group_to_sample_cache[num_seqs] = PyObjectCache(
+                gen_seq_group_to_sample_builder(num_seqs))
+
+        obj = self._seq_group_to_sample_cache[num_seqs].get_object()
+        return obj
+
+    def reset(self):
+        for cache in self._seq_group_to_sample_cache.values():
+            cache.reset()
+
+
+class SamplingMetadata:
+    """Metadata for input sequences. Used in sampler.
+
+    The usage is as follow;
+    ```
+    hidden_states = execute_model(...)
+    logits = hidden_states[sampling_metadata.selected_token_indices]
+    sample(logits)
+
+    def sample(logits):
+        # Use categorized_sample_indices for sampling....
+    ```
+
+    Args:
+        seq_groups: List of batched sequence groups.
+        selected_token_indices: (num_query_tokens_to_logprob). Indices to find
+            logits from the initial model output hidden states.
+        categorized_sample_indices: SamplingType -> token indices to sample.
+            Each token indices is 2D tensor of (num_indices, num_indices) where
+            the first item means the sample index within the returned logit
+            (before pruning padding), and the second item means the sample
+            index after pruning using selected_token_indices.
+            For example, if the returned logit is [1, 2, 3], and we select
+            [1, 2] for sampling, the pruned logit will be [2, 3]. In this case,
+            The first tuple is [1, 2] (sampled index within original logit),
+            and the second tuple is [0, 1] (sampled index within pruned logit).
+        num_prompts: Number of prompt sequence groups in seq_groups.
+        skip_sampler_cpu_output: Indicates if we want to skip the GPU=>CPU
+            serialization of token outputs.
+        reuse_sampling_tensors: Indicates if we want to reuse sampling
+            tensors that are part of the sampler forward pass. Currently,
+            it is mainly used for multi-step decode.
+
+    """
+
+    def __init__(
+        self,
+        seq_groups: List[SequenceGroupToSample],
+        selected_token_indices: torch.Tensor,
+        categorized_sample_indices: Dict[SamplingType, torch.Tensor],
+        num_prompts: int,
+        skip_sampler_cpu_output: bool = False,
+        reuse_sampling_tensors: bool = False,
+    ) -> None:
+        self.seq_groups = seq_groups
+        self.selected_token_indices = selected_token_indices
+        self.categorized_sample_indices = categorized_sample_indices
+        self.num_prompts = num_prompts
+        self.skip_sampler_cpu_output = skip_sampler_cpu_output
+        self.reuse_sampling_tensors = reuse_sampling_tensors
+
+    @staticmethod
+    def prepare(
+        seq_group_metadata_list: List[SequenceGroupMetadata],
+        seq_lens: List[int],
+        query_lens: List[int],
+        device: str,
+        pin_memory: bool,
+        generators: Optional[Dict[str, torch.Generator]] = None,
+        cache: Optional[SamplingMetadataCache] = None,
+    ) -> "SamplingMetadata":
+        (
+            seq_groups,
+            selected_token_indices,
+            categorized_sample_indices,
+            num_prompts,
+        ) = _prepare_seq_groups(seq_group_metadata_list, seq_lens, query_lens,
+                                device, generators, cache)
+        selected_token_indices = async_tensor_h2d(
+            selected_token_indices,
+            dtype=torch.long,
+            target_device=device,
+            pin_memory=pin_memory,
+        )
+        categorized_sample_indices = {
+            t: async_tensor_h2d(
+                seq_ids,
+                dtype=torch.int,
+                target_device=device,
+                pin_memory=pin_memory,
+            )
+            for t, seq_ids in categorized_sample_indices.items()
+        }
+
+        sampling_metadata = SamplingMetadata(
+            seq_groups=seq_groups,
+            selected_token_indices=selected_token_indices,
+            categorized_sample_indices=categorized_sample_indices,
+            num_prompts=num_prompts,
+        )
+        return sampling_metadata
+
+    def __repr__(self) -> str:
+        return (
+            "SamplingMetadata("
+            f"seq_groups={self.seq_groups}, "
+            f"selected_token_indices={self.selected_token_indices}, "
+            f"categorized_sample_indices={self.categorized_sample_indices}), ")
+
+
+def _get_prompt_logprob_output_indices(
+    sampling_params: SamplingParams,
+    seq_data: SequenceData,
+    prompt_logprob_len: int,
+) -> List[int]:
+    if sampling_params.prompt_logprobs is None or prompt_logprob_len <= 0:
+        return []
+    positions = sampling_params.prompt_logprob_positions
+    computed_len = seq_data.get_num_computed_tokens()
+    available_next_tokens = max(
+        0,
+        len(seq_data.prompt_token_ids) - computed_len - 1,
+    )
+    materialized_len = min(prompt_logprob_len, available_next_tokens)
+    if positions is None:
+        return list(range(materialized_len))
+
+    output_indices = [
+        position - computed_len - 1
+        for position in positions
+        if computed_len < position
+        <= computed_len + materialized_len
+    ]
+    assert output_indices == sorted(set(output_indices))
+    assert all(0 <= index < materialized_len for index in output_indices)
+    return output_indices
+
+
+def _prepare_seq_groups(
+    seq_group_metadata_list: List[SequenceGroupMetadata],
+    seq_lens: List[int],
+    query_lens: List[int],
+    device: str,
+    generators: Optional[Dict[str, torch.Generator]] = None,
+    cache: Optional[SamplingMetadataCache] = None,
+) -> Tuple[List[SequenceGroupToSample], List[int], Dict[SamplingType,
+                                                        List[int]], int, ]:
+    """Prepare sequence groups and indices for sampling.
+
+    Args:
+        seq_group_metadata_list: A list of sequence group to batch.
+        seq_lens: A list of sequence lens per sequence group.
+            Index of prompt len should match with seq_group_metadata_list.
+        query_lens: A list of query lengths. Prompt lens include the length
+            of entire prompt tokens, and it could be shorter.
+        device: A device to use for random number generators,
+            `SequenceGroupToSample.generator`.
+        generators: A store of per-request random number generators used
+            for seeded requests.
+
+    Returns:
+        seq_groups: A list of sequence group to sample.
+        selected_token_indices: See the definition from `SamplingMetadata`.
+        categorized_sample_indices: See the definition from `SamplingMetadata`.
+        num_prompts: Total number of prompts from `seq_group_metadata_list`.
+    """
+    # Batched sequence groups for the current model forward stsep.
+    seq_groups: List[SequenceGroupToSample] = []
+    # A list of token indices to sample/compute logprob. It is used to
+    # prune the outcome logits from the model for the performance.
+    selected_token_indices: List[int] = []
+    # Used for selected_token_indices.
+    model_output_idx = 0
+
+    # Sampling type -> (
+    # indices to sample/prompt logprob within pruned output logits,
+    # indices to sample within pruned logits)
+    categorized_sample_indices: Dict[SamplingType, List[int]] = {
+        t: []
+        for t in SamplingType
+    }
+    # Index of logits to compute logprob. Logits include both prompt logprob
+    # and sample logprob indices.
+    logit_idx = 0
+    # Total number of prompts from given sequence groups.
+    num_prompts = 0
+
+    for i, seq_group_metadata in enumerate(seq_group_metadata_list):
+        seq_ids = seq_group_metadata.seq_data.keys()
+
+        if cache is not None:
+            sample_obj = cache.get_cached_seq_group_to_sample(len(seq_ids))
+
+            for j, seq_id in enumerate(seq_ids):
+                sample_obj.seq_ids[j] = seq_id
+
+            sample_obj.prompt_logprob_indices.clear()
+            sample_obj.prompt_logprob_output_indices.clear()
+            sample_obj.sample_indices.clear()
+
+        sampling_params = seq_group_metadata.sampling_params
+        is_prompt = seq_group_metadata.is_prompt
+        generator: Optional[torch.Generator] = None
+        # If the current seq group is in decode stage, it is None.
+        seq_len: Optional[int] = None
+        query_len: Optional[int] = None
+        prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices
+                                             if cache is not None else [])
+        prompt_logprob_output_indices: List[int] = (
+            sample_obj.prompt_logprob_output_indices
+            if cache is not None else [])
+        sample_indices: List[int] = (sample_obj.sample_indices
+                                     if cache is not None else [])
+        do_sample = seq_group_metadata.do_sample
+
+        if seq_group_metadata.is_prompt:
+            if sampling_params.seed is not None:
+                generator = torch.Generator(device=device).manual_seed(
+                    sampling_params.seed)
+                if generators is not None:
+                    generators[seq_group_metadata.request_id] = generator
+
+            num_prompts += 1
+            num_prefill_sample = len(seq_ids)
+            assert num_prefill_sample == 1
+            assert query_lens is not None and seq_lens is not None
+            query_len, seq_len = query_lens[i], seq_lens[i]
+            # If we need sampling, exclude num_prefill_sample tokens from
+            # prompt logprob.
+            prompt_logprob_len = (query_len - num_prefill_sample
+                                  if do_sample else query_len)
+            sample_len = num_prefill_sample if do_sample else 0
+        else:
+            # Decode
+            prompt_logprob_len = 0
+            query_len = query_lens[i] if query_lens is not None else 1
+            sample_len = len(seq_ids) * query_len if do_sample else 0
+
+            if sampling_params.seed is not None and generators is not None:
+                generator = generators.get(seq_group_metadata.request_id)
+
+        seq_data = next(iter(seq_group_metadata.seq_data.values()))
+        prompt_logprob_output_indices.extend(
+            _get_prompt_logprob_output_indices(
+                sampling_params,
+                seq_data,
+                prompt_logprob_len,
+            ))
+
+        # Update indices to select from the model output.
+        """
+        This blocks computes selected_token_indices which is used in the
+        following way.
+
+        hidden_states = model(...)
+        logits = hidden_states[selected_token_indices]
+        """
+
+        if sampling_params.prompt_logprobs is not None:
+            selected_token_indices.extend(
+                model_output_idx + output_index
+                for output_index in prompt_logprob_output_indices)
+        model_output_idx += prompt_logprob_len
+        if do_sample:
+            selected_token_indices.extend(
+                range(model_output_idx, model_output_idx + sample_len))
+        model_output_idx += sample_len
+
+        # We now find indices for logprob computation and sampling.
+        """
+        This block computes categorized_sample_indices which is used in the
+        following way.
+
+        hidden_states = model(...)
+        logits = hidden_states[selected_token_indices]
+        def sample(logits):
+           # Use categorized_sample_indices for sampling.
+           # prompt_logprob_indices to find prompt logprob indices.
+           # sample_indices to find sample indices.
+        """
+
+        if sampling_params.prompt_logprobs is not None:
+            prompt_logprob_indices.extend(
+                range(logit_idx,
+                      logit_idx + len(prompt_logprob_output_indices)))
+            logit_idx += len(prompt_logprob_output_indices)
+        if do_sample:
+            sample_indices.extend(range(logit_idx, logit_idx + sample_len))
+            categorized_sample_indices[sampling_params.sampling_type].extend(
+                list(range(logit_idx, logit_idx + sample_len)))
+            logit_idx += sample_len
+
+        if cache is not None:
+            sample_obj.sampling_params = sampling_params
+            sample_obj.seq_data = seq_group_metadata.seq_data
+            sample_obj.seq_len = seq_len
+            sample_obj.query_len = query_len
+            sample_obj.generator = generator
+            sample_obj.is_prompt = is_prompt
+        else:
+            sample_obj = SequenceGroupToSample(
+                seq_ids=list(seq_ids),
+                sampling_params=sampling_params,
+                seq_data=seq_group_metadata.seq_data,
+                seq_len=seq_len,
+                query_len=query_len,
+                generator=generator,
+                is_prompt=is_prompt,
+                prompt_logprob_indices=list(prompt_logprob_indices),
+                prompt_logprob_output_indices=list(
+                    prompt_logprob_output_indices),
+                sample_indices=list(sample_indices),
+            )
+
+        assert (len(sample_obj.prompt_logprob_indices)
+                == len(sample_obj.prompt_logprob_output_indices))
+        seq_groups.append(sample_obj)
+
+    if cache is not None:
+        cache.reset()
+
+    return (seq_groups, selected_token_indices, categorized_sample_indices,
+            num_prompts)
+
+
+@dataclass
+class SamplingTensors:
+    """Tensors for sampling."""
+
+    temperatures: torch.Tensor
+    top_ps: torch.Tensor
+    top_ks: torch.Tensor
+    min_ps: torch.Tensor
+    presence_penalties: torch.Tensor
+    frequency_penalties: torch.Tensor
+    repetition_penalties: torch.Tensor
+    prompt_tokens: torch.Tensor
+    output_tokens: torch.Tensor
+
+    @classmethod
+    def from_sampling_metadata(
+        cls,
+        sampling_metadata: "SamplingMetadata",
+        vocab_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> Tuple["SamplingTensors", bool, bool, bool]:
+        prompt_tokens: List[array] = []
+        output_tokens: List[array] = []
+        top_ks: List[int] = []
+        temperatures: List[float] = []
+        top_ps: List[float] = []
+        min_ps: List[float] = []
+        presence_penalties: List[float] = []
+        frequency_penalties: List[float] = []
+        repetition_penalties: List[float] = []
+        do_penalties = False
+        do_top_p_top_k = False
+        do_min_p = False
+
+        assert sampling_metadata.seq_groups is not None
+        for seq_group in sampling_metadata.seq_groups:
+            seq_ids = seq_group.seq_ids
+            sampling_params = seq_group.sampling_params
+            temperature = sampling_params.temperature
+            p = sampling_params.presence_penalty
+            f = sampling_params.frequency_penalty
+            r = sampling_params.repetition_penalty
+            top_p = sampling_params.top_p
+            min_p = sampling_params.min_p
+
+            # k should not be greater than the vocab size.
+            top_k = min(sampling_params.top_k, vocab_size)
+            top_k = vocab_size if top_k == -1 else top_k
+            if temperature < _SAMPLING_EPS:
+                # NOTE: Zero temperature means deterministic sampling
+                # (i.e., greedy sampling or beam search).
+                # Set the temperature to 1 to avoid division by zero.
+                temperature = 1.0
+            if not do_top_p_top_k and (top_p < 1.0 - _SAMPLING_EPS
+                                       or top_k != vocab_size):
+                do_top_p_top_k = True
+            if not do_min_p and min_p > _SAMPLING_EPS:
+                do_min_p = True
+            if not do_penalties and (abs(p) >= _SAMPLING_EPS
+                                     or abs(f) >= _SAMPLING_EPS
+                                     or abs(r - 1.0) >= _SAMPLING_EPS):
+                do_penalties = True
+
+            is_prompt = seq_group.is_prompt
+            if is_prompt and sampling_params.prompt_logprobs is not None:
+                # For tokens in the prompt that we only need to get
+                # their logprobs
+                query_len = seq_group.query_len
+                assert query_len is not None
+                prefill_len = len(seq_group.prompt_logprob_indices)
+                temperatures += [temperature] * prefill_len
+                top_ps += [top_p] * prefill_len
+                top_ks += [top_k] * prefill_len
+                min_ps += [min_p] * prefill_len
+                presence_penalties += [0] * prefill_len
+                frequency_penalties += [0] * prefill_len
+                repetition_penalties += [1] * prefill_len
+
+            if seq_group.do_sample:
+                sample_lens = len(seq_group.sample_indices)
+                assert sample_lens >= len(seq_ids)
+                temperatures += [temperature] * sample_lens
+                top_ps += [top_p] * sample_lens
+                top_ks += [top_k] * sample_lens
+                min_ps += [min_p] * sample_lens
+                presence_penalties += [p] * sample_lens
+                frequency_penalties += [f] * sample_lens
+                repetition_penalties += [r] * sample_lens
+
+        if do_penalties:
+            for seq_group in sampling_metadata.seq_groups:
+                seq_ids = seq_group.seq_ids
+                if (seq_group.is_prompt
+                        and sampling_params.prompt_logprobs is not None):
+                    prefill_len = len(seq_group.prompt_logprob_indices)
+                    prompt_tokens.extend(
+                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
+                        for _ in range(prefill_len))
+                    output_tokens.extend(
+                        array(VLLM_TOKEN_ID_ARRAY_TYPE)
+                        for _ in range(prefill_len))
+                if seq_group.do_sample:
+                    for seq_id in seq_ids:
+                        seq_data = seq_group.seq_data[seq_id]
+                        prompt_tokens.append(seq_data.prompt_token_ids_array)
+                        output_tokens.append(seq_data.output_token_ids_array)
+
+        sampling_tensors = SamplingTensors.from_lists(
+            temperatures,
+            top_ps,
+            top_ks,
+            min_ps,
+            presence_penalties,
+            frequency_penalties,
+            repetition_penalties,
+            prompt_tokens,
+            output_tokens,
+            vocab_size,
+            device,
+            dtype,
+        )
+        return (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p)
+
+    @classmethod
+    def from_lists(
+        cls,
+        temperatures: List[float],
+        top_ps: List[float],
+        top_ks: List[int],
+        min_ps: List[float],
+        presence_penalties: List[float],
+        frequency_penalties: List[float],
+        repetition_penalties: List[float],
+        prompt_tokens: List[array],
+        output_tokens: List[array],
+        vocab_size: int,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> "SamplingTensors":
+        # Note that the performance will be very bad without
+        # pinned memory.
+        pin_memory = is_pin_memory_available()
+
+        do_penalties = prompt_tokens or output_tokens
+
+        if do_penalties:
+            prompt_t = make_tensor_with_pad(
+                prompt_tokens,
+                vocab_size,
+                device="cpu",
+                dtype=torch.int64,
+                pin_memory=pin_memory,
+            )
+            output_t = make_tensor_with_pad(
+                output_tokens,
+                vocab_size,
+                device="cpu",
+                dtype=torch.int64,
+                pin_memory=pin_memory,
+            )
+        else:
+            empty_tensor = torch.empty(0, device=device, dtype=torch.long)
+            prompt_t = empty_tensor
+            output_t = empty_tensor
+
+        temperatures_t = torch.tensor(
+            temperatures,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        top_ps_t = torch.tensor(
+            top_ps,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        min_ps_t = torch.tensor(
+            min_ps,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        presence_penalties_t = torch.tensor(
+            presence_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        frequency_penalties_t = torch.tensor(
+            frequency_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        repetition_penalties_t = torch.tensor(
+            repetition_penalties,
+            device="cpu",
+            dtype=dtype,
+            pin_memory=pin_memory,
+        )
+        top_ks_t = torch.tensor(
+            top_ks,
+            device="cpu",
+            dtype=torch.int,
+            pin_memory=pin_memory,
+        )
+        # Because the memory is pinned, we can do non-blocking
+        # transfer to device.
+
+        return cls(
+            temperatures=temperatures_t.to(device=device, non_blocking=True),
+            top_ps=top_ps_t.to(device=device, non_blocking=True),
+            top_ks=top_ks_t.to(device=device, non_blocking=True),
+            min_ps=min_ps_t.to(device=device, non_blocking=True),
+            presence_penalties=presence_penalties_t.to(device=device,
+                                                       non_blocking=True),
+            frequency_penalties=frequency_penalties_t.to(device=device,
+                                                         non_blocking=True),
+            repetition_penalties=repetition_penalties_t.to(device=device,
+                                                           non_blocking=True),
+            prompt_tokens=prompt_t.to(device=device, non_blocking=True),
+            output_tokens=output_t.to(device=device, non_blocking=True),
+        )

qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py

@@ -0,0 +1,520 @@
+"""Sampling parameters for text generation."""
+import copy
+from dataclasses import dataclass
+from enum import Enum, IntEnum
+from functools import cached_property
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import msgspec
+import torch
+from pydantic import BaseModel
+from typing_extensions import Annotated
+
+from vllm.logger import init_logger
+
+logger = init_logger(__name__)
+
+_SAMPLING_EPS = 1e-5
+_MAX_TEMP = 1e-2
+
+
+class SamplingType(IntEnum):
+    GREEDY = 0
+    RANDOM = 1
+    RANDOM_SEED = 2
+
+
+LogitsProcessor = Union[Callable[[List[int], torch.Tensor], torch.Tensor],
+                        Callable[[List[int], List[int], torch.Tensor],
+                                 torch.Tensor]]
+"""LogitsProcessor is a function that takes a list
+of previously generated tokens, the logits tensor
+for the next token and, optionally, prompt tokens as a
+first argument, and returns a modified tensor of logits
+to sample from."""
+
+
+# maybe make msgspec?
+@dataclass
+class GuidedDecodingParams:
+    """One of these fields will be used to build a logit processor."""
+    json: Optional[Union[str, Dict]] = None
+    regex: Optional[str] = None
+    choice: Optional[List[str]] = None
+    grammar: Optional[str] = None
+    json_object: Optional[bool] = None
+    """These are other options that can be set"""
+    backend: Optional[str] = None
+    whitespace_pattern: Optional[str] = None
+
+    @staticmethod
+    def from_optional(
+        json: Optional[Union[Dict, BaseModel, str]],
+        regex: Optional[str] = None,
+        choice: Optional[List[str]] = None,
+        grammar: Optional[str] = None,
+        json_object: Optional[bool] = None,
+        backend: Optional[str] = None,
+        whitespace_pattern: Optional[str] = None,
+    ) -> "GuidedDecodingParams":
+        # Extract json schemas from pydantic models
+        if isinstance(json, (BaseModel, type(BaseModel))):
+            json = json.model_json_schema()
+        return GuidedDecodingParams(
+            json=json,
+            regex=regex,
+            choice=choice,
+            grammar=grammar,
+            json_object=json_object,
+            backend=backend,
+            whitespace_pattern=whitespace_pattern,
+        )
+
+    def __post_init__(self):
+        """Validate that some fields are mutually exclusive."""
+        guide_count = sum([
+            self.json is not None, self.regex is not None, self.choice
+            is not None, self.grammar is not None, self.json_object is not None
+        ])
+        if guide_count > 1:
+            raise ValueError(
+                "You can only use one kind of guided decoding but multiple are "
+                f"specified: {self.__dict__}")
+
+
+class RequestOutputKind(Enum):
+    # Return entire output so far in every RequestOutput
+    CUMULATIVE = 0
+    # Return only deltas in each RequestOutput
+    DELTA = 1
+    # Do not return intermediate RequestOuputs
+    FINAL_ONLY = 2
+
+
+class SamplingParams(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        # required for @cached_property.
+        dict=True):  # type: ignore[call-arg]
+    """Sampling parameters for text generation.
+
+    Overall, we follow the sampling parameters from the OpenAI text completion
+    API (https://platform.openai.com/docs/api-reference/completions/create).
+    In addition, we support beam search, which is not supported by OpenAI.
+
+    Args:
+        n: Number of output sequences to return for the given prompt.
+        best_of: Number of output sequences that are generated from the prompt.
+            From these `best_of` sequences, the top `n` sequences are returned.
+            `best_of` must be greater than or equal to `n`. By default,
+            `best_of` is set to `n`.
+        presence_penalty: Float that penalizes new tokens based on whether they
+            appear in the generated text so far. Values > 0 encourage the model
+            to use new tokens, while values < 0 encourage the model to repeat
+            tokens.
+        frequency_penalty: Float that penalizes new tokens based on their
+            frequency in the generated text so far. Values > 0 encourage the
+            model to use new tokens, while values < 0 encourage the model to
+            repeat tokens.
+        repetition_penalty: Float that penalizes new tokens based on whether
+            they appear in the prompt and the generated text so far. Values > 1
+            encourage the model to use new tokens, while values < 1 encourage
+            the model to repeat tokens.
+        temperature: Float that controls the randomness of the sampling. Lower
+            values make the model more deterministic, while higher values make
+            the model more random. Zero means greedy sampling.
+        top_p: Float that controls the cumulative probability of the top tokens
+            to consider. Must be in (0, 1]. Set to 1 to consider all tokens.
+        top_k: Integer that controls the number of top tokens to consider. Set
+            to -1 to consider all tokens.
+        min_p: Float that represents the minimum probability for a token to be
+            considered, relative to the probability of the most likely token.
+            Must be in [0, 1]. Set to 0 to disable this.
+        seed: Random seed to use for the generation.
+        stop: List of strings that stop the generation when they are generated.
+            The returned output will not contain the stop strings.
+        stop_token_ids: List of tokens that stop the generation when they are
+            generated. The returned output will contain the stop tokens unless
+            the stop tokens are special tokens.
+        include_stop_str_in_output: Whether to include the stop strings in
+            output text. Defaults to False.
+        ignore_eos: Whether to ignore the EOS token and continue generating
+            tokens after the EOS token is generated.
+        max_tokens: Maximum number of tokens to generate per output sequence.
+        min_tokens: Minimum number of tokens to generate per output sequence
+            before EOS or stop_token_ids can be generated
+        logprobs: Number of log probabilities to return per output token.
+            When set to None, no probability is returned. If set to a non-None
+            value, the result includes the log probabilities of the specified
+            number of most likely tokens, as well as the chosen tokens.
+            Note that the implementation follows the OpenAI API: The API will
+            always return the log probability of the sampled token, so there
+            may be up to `logprobs+1` elements in the response.
+        prompt_logprobs: Number of log probabilities to return per prompt token.
+        detokenize: Whether to detokenize the output. Defaults to True.
+        skip_special_tokens: Whether to skip special tokens in the output.
+        spaces_between_special_tokens: Whether to add spaces between special
+            tokens in the output.  Defaults to True.
+        logits_processors: List of functions that modify logits based on
+            previously generated tokens, and optionally prompt tokens as
+            a first argument.
+        truncate_prompt_tokens: If set to an integer k, will use only the last k
+            tokens from the prompt (i.e., left truncation). Defaults to None
+            (i.e., no truncation).
+        guided_decoding: If provided, the engine will construct a guided
+            decoding logits processor from these parameters. Defaults to None.
+        logit_bias: If provided, the engine will construct a logits processor
+            that applies these logit biases. Defaults to None.
+        allowed_token_ids: If provided, the engine will construct a logits
+            processor which only retains scores for the given token ids.
+            Defaults to None.
+        prompt_logprob_positions: Optional prompt-token positions whose logits
+            should be materialized. None preserves the standard all-position
+            prompt-logprob behavior.
+    """
+
+    n: int = 1
+    best_of: Optional[int] = None
+    _real_n: Optional[int] = None
+    presence_penalty: float = 0.0
+    frequency_penalty: float = 0.0
+    repetition_penalty: float = 1.0
+    temperature: float = 1.0
+    top_p: float = 1.0
+    top_k: int = -1
+    min_p: float = 0.0
+    seed: Optional[int] = None
+    stop: Optional[Union[str, List[str]]] = None
+    stop_token_ids: Optional[List[int]] = None
+    ignore_eos: bool = False
+    max_tokens: Optional[int] = 16
+    min_tokens: int = 0
+    logprobs: Optional[int] = None
+    prompt_logprobs: Optional[int] = None
+    # NOTE: This parameter is only exposed at the engine level for now.
+    # It is not exposed in the OpenAI API server, as the OpenAI API does
+    # not support returning only a list of token IDs.
+    detokenize: bool = True
+    skip_special_tokens: bool = True
+    spaces_between_special_tokens: bool = True
+    # Optional[List[LogitsProcessor]] type. We use Any here because
+    # Optional[List[LogitsProcessor]] type is not supported by msgspec.
+    logits_processors: Optional[Any] = None
+    include_stop_str_in_output: bool = False
+    truncate_prompt_tokens: Optional[Annotated[int, msgspec.Meta(ge=1)]] = None
+    output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE
+
+    # The below fields are not supposed to be used as an input.
+    # They are set in post_init.
+    output_text_buffer_length: int = 0
+    _all_stop_token_ids: Set[int] = msgspec.field(default_factory=set)
+
+    # Fields used to construct logits processors
+    guided_decoding: Optional[GuidedDecodingParams] = None
+    logit_bias: Optional[Dict[int, float]] = None
+    allowed_token_ids: Optional[List[int]] = None
+    prompt_logprob_positions: Optional[List[int]] = None
+
+    @staticmethod
+    def from_optional(
+        n: Optional[int] = 1,
+        best_of: Optional[int] = None,
+        presence_penalty: Optional[float] = 0.0,
+        frequency_penalty: Optional[float] = 0.0,
+        repetition_penalty: Optional[float] = 1.0,
+        temperature: Optional[float] = 1.0,
+        top_p: Optional[float] = 1.0,
+        top_k: int = -1,
+        min_p: float = 0.0,
+        seed: Optional[int] = None,
+        stop: Optional[Union[str, List[str]]] = None,
+        stop_token_ids: Optional[List[int]] = None,
+        include_stop_str_in_output: bool = False,
+        ignore_eos: bool = False,
+        max_tokens: Optional[int] = 16,
+        min_tokens: int = 0,
+        logprobs: Optional[int] = None,
+        prompt_logprobs: Optional[int] = None,
+        detokenize: bool = True,
+        skip_special_tokens: bool = True,
+        spaces_between_special_tokens: bool = True,
+        logits_processors: Optional[List[LogitsProcessor]] = None,
+        truncate_prompt_tokens: Optional[Annotated[int,
+                                                   msgspec.Meta(ge=1)]] = None,
+        output_kind: RequestOutputKind = RequestOutputKind.CUMULATIVE,
+        guided_decoding: Optional[GuidedDecodingParams] = None,
+        logit_bias: Optional[Union[Dict[int, float], Dict[str, float]]] = None,
+        allowed_token_ids: Optional[List[int]] = None,
+        prompt_logprob_positions: Optional[List[int]] = None,
+    ) -> "SamplingParams":
+        if logit_bias is not None:
+            logit_bias = {
+                int(token): bias
+                for token, bias in logit_bias.items()
+            }
+
+        return SamplingParams(
+            n=1 if n is None else n,
+            best_of=best_of,
+            presence_penalty=0.0
+            if presence_penalty is None else presence_penalty,
+            frequency_penalty=0.0
+            if frequency_penalty is None else frequency_penalty,
+            repetition_penalty=1.0
+            if repetition_penalty is None else repetition_penalty,
+            temperature=1.0 if temperature is None else temperature,
+            top_p=1.0 if top_p is None else top_p,
+            top_k=top_k,
+            min_p=min_p,
+            seed=seed,
+            stop=stop,
+            stop_token_ids=stop_token_ids,
+            include_stop_str_in_output=include_stop_str_in_output,
+            ignore_eos=ignore_eos,
+            max_tokens=max_tokens,
+            min_tokens=min_tokens,
+            logprobs=logprobs,
+            prompt_logprobs=prompt_logprobs,
+            detokenize=detokenize,
+            skip_special_tokens=skip_special_tokens,
+            spaces_between_special_tokens=spaces_between_special_tokens,
+            logits_processors=logits_processors,
+            truncate_prompt_tokens=truncate_prompt_tokens,
+            output_kind=output_kind,
+            guided_decoding=guided_decoding,
+            logit_bias=logit_bias,
+            allowed_token_ids=allowed_token_ids,
+            prompt_logprob_positions=prompt_logprob_positions,
+        )
+
+    def __post_init__(self) -> None:
+        # how we deal with `best_of``:
+        # if `best_of`` is not set, we default to `n`;
+        # if `best_of`` is set, we set `n`` to `best_of`,
+        # and set `_real_n`` to the original `n`.
+        # when we return the result, we will check
+        # if we need to return `n` or `_real_n` results
+        if self.best_of:
+            if self.best_of < self.n:
+                raise ValueError(
+                    f"best_of must be greater than or equal to n, "
+                    f"got n={self.n} and best_of={self.best_of}.")
+            self._real_n = self.n
+            self.n = self.best_of
+        if 0 < self.temperature < _MAX_TEMP:
+            logger.warning(
+                "temperature %s is less than %s, which may cause numerical "
+                "errors nan or inf in tensors. We have maxed it out to %s.",
+                self.temperature, _MAX_TEMP, _MAX_TEMP)
+            self.temperature = max(self.temperature, _MAX_TEMP)
+        if self.seed == -1:
+            self.seed = None
+        else:
+            self.seed = self.seed
+        if self.stop is None:
+            self.stop = []
+        elif isinstance(self.stop, str):
+            self.stop = [self.stop]
+        else:
+            self.stop = list(self.stop)
+        if self.stop_token_ids is None:
+            self.stop_token_ids = []
+        else:
+            self.stop_token_ids = list(self.stop_token_ids)
+        self.logprobs = 1 if self.logprobs is True else self.logprobs
+        self.prompt_logprobs = (1 if self.prompt_logprobs is True else
+                                self.prompt_logprobs)
+        if self.prompt_logprob_positions is not None:
+            self.prompt_logprob_positions = list(
+                self.prompt_logprob_positions)
+
+        # Number of characters to hold back for stop string evaluation
+        # until sequence is finished.
+        if self.stop and not self.include_stop_str_in_output:
+            self.output_text_buffer_length = max(len(s) for s in self.stop) - 1
+
+        self._verify_args()
+
+        if self.temperature < _SAMPLING_EPS:
+            # Zero temperature means greedy sampling.
+            self.top_p = 1.0
+            self.top_k = -1
+            self.min_p = 0.0
+            self._verify_greedy_sampling()
+        # eos_token_id is added to this by the engine
+        self._all_stop_token_ids = set(self.stop_token_ids)
+
+    def _verify_args(self) -> None:
+        if not isinstance(self.n, int):
+            raise ValueError(f"n must be an int, but is of "
+                             f"type {type(self.n)}")
+        if self.n < 1:
+            raise ValueError(f"n must be at least 1, got {self.n}.")
+        if not -2.0 <= self.presence_penalty <= 2.0:
+            raise ValueError("presence_penalty must be in [-2, 2], got "
+                             f"{self.presence_penalty}.")
+        if not -2.0 <= self.frequency_penalty <= 2.0:
+            raise ValueError("frequency_penalty must be in [-2, 2], got "
+                             f"{self.frequency_penalty}.")
+        if not 0.0 < self.repetition_penalty <= 2.0:
+            raise ValueError("repetition_penalty must be in (0, 2], got "
+                             f"{self.repetition_penalty}.")
+        if self.temperature < 0.0:
+            raise ValueError(
+                f"temperature must be non-negative, got {self.temperature}.")
+        if not 0.0 < self.top_p <= 1.0:
+            raise ValueError(f"top_p must be in (0, 1], got {self.top_p}.")
+        if self.top_k < -1 or self.top_k == 0:
+            raise ValueError(f"top_k must be -1 (disable), or at least 1, "
+                             f"got {self.top_k}.")
+        if not isinstance(self.top_k, int):
+            raise TypeError(
+                f"top_k must be an integer, got {type(self.top_k).__name__}")
+        if not 0.0 <= self.min_p <= 1.0:
+            raise ValueError("min_p must be in [0, 1], got "
+                             f"{self.min_p}.")
+        if self.max_tokens is not None and self.max_tokens < 1:
+            raise ValueError(
+                f"max_tokens must be at least 1, got {self.max_tokens}.")
+        if self.min_tokens < 0:
+            raise ValueError(f"min_tokens must be greater than or equal to 0, "
+                             f"got {self.min_tokens}.")
+        if self.max_tokens is not None and self.min_tokens > self.max_tokens:
+            raise ValueError(
+                f"min_tokens must be less than or equal to "
+                f"max_tokens={self.max_tokens}, got {self.min_tokens}.")
+        if self.logprobs is not None and self.logprobs < 0:
+            raise ValueError(
+                f"logprobs must be non-negative, got {self.logprobs}.")
+        if self.prompt_logprobs is not None and self.prompt_logprobs < 0:
+            raise ValueError(f"prompt_logprobs must be non-negative, got "
+                             f"{self.prompt_logprobs}.")
+        if self.prompt_logprob_positions is not None:
+            if self.prompt_logprobs is None:
+                raise ValueError(
+                    "prompt_logprob_positions requires prompt_logprobs.")
+            if (
+                not self.prompt_logprob_positions
+                or any(
+                    not isinstance(position, int)
+                    or isinstance(position, bool)
+                    or position <= 0
+                    for position in self.prompt_logprob_positions
+                )
+                or self.prompt_logprob_positions
+                != sorted(set(self.prompt_logprob_positions))
+            ):
+                raise ValueError(
+                    "prompt_logprob_positions must be a sorted unique list "
+                    "of positive integers.")
+        if (self.truncate_prompt_tokens is not None
+                and self.truncate_prompt_tokens < 1):
+            raise ValueError(f"truncate_prompt_tokens must be >= 1, "
+                             f"got {self.truncate_prompt_tokens}")
+        assert isinstance(self.stop, list)
+        if any(not stop_str for stop_str in self.stop):
+            raise ValueError("stop cannot contain an empty string.")
+        if self.stop and not self.detokenize:
+            raise ValueError(
+                "stop strings are only supported when detokenize is True. "
+                "Set detokenize=True to use stop.")
+        if self.best_of != self._real_n and self.output_kind == (
+                RequestOutputKind.DELTA):
+            raise ValueError("best_of must equal n to use output_kind=DELTA")
+
+    def _verify_greedy_sampling(self) -> None:
+        if self.n > 1:
+            raise ValueError("n must be 1 when using greedy sampling, "
+                             f"got {self.n}.")
+
+    def update_from_generation_config(
+            self,
+            generation_config: Dict[str, Any],
+            model_eos_token_id: Optional[int] = None) -> None:
+        """Update if there are non-default values from generation_config"""
+
+        if model_eos_token_id is not None:
+            # Add the eos token id into the sampling_params to support
+            # min_tokens processing.
+            self._all_stop_token_ids.add(model_eos_token_id)
+
+        # Update eos_token_id for generation
+        if (eos_ids := generation_config.get("eos_token_id")) is not None:
+            # it can be either int or list of int
+            eos_ids = {eos_ids} if isinstance(eos_ids, int) else set(eos_ids)
+            if model_eos_token_id is not None:
+                # We don't need to include the primary eos_token_id in
+                # stop_token_ids since it's handled separately for stopping
+                # purposes.
+                eos_ids.discard(model_eos_token_id)
+            if eos_ids:
+                self._all_stop_token_ids.update(eos_ids)
+                if not self.ignore_eos:
+                    eos_ids.update(self.stop_token_ids)
+                    self.stop_token_ids = list(eos_ids)
+
+    @cached_property
+    def sampling_type(self) -> SamplingType:
+        if self.temperature < _SAMPLING_EPS:
+            return SamplingType.GREEDY
+        if self.seed is not None:
+            return SamplingType.RANDOM_SEED
+        return SamplingType.RANDOM
+
+    @property
+    def all_stop_token_ids(self) -> Set[int]:
+        return self._all_stop_token_ids
+
+    def clone(self) -> "SamplingParams":
+        """Deep copy excluding LogitsProcessor objects.
+
+        LogitsProcessor objects are excluded because they may contain an
+        arbitrary, nontrivial amount of data.
+        See https://github.com/vllm-project/vllm/issues/3087
+        """
+
+        logit_processor_refs = None if self.logits_processors is None else {
+            id(lp): lp
+            for lp in self.logits_processors
+        }
+        return copy.deepcopy(self, memo=logit_processor_refs)
+
+    def __repr__(self) -> str:
+        return (
+            f"SamplingParams(n={self.n}, "
+            f"presence_penalty={self.presence_penalty}, "
+            f"frequency_penalty={self.frequency_penalty}, "
+            f"repetition_penalty={self.repetition_penalty}, "
+            f"temperature={self.temperature}, "
+            f"top_p={self.top_p}, "
+            f"top_k={self.top_k}, "
+            f"min_p={self.min_p}, "
+            f"seed={self.seed}, "
+            f"stop={self.stop}, "
+            f"stop_token_ids={self.stop_token_ids}, "
+            f"include_stop_str_in_output={self.include_stop_str_in_output}, "
+            f"ignore_eos={self.ignore_eos}, "
+            f"max_tokens={self.max_tokens}, "
+            f"min_tokens={self.min_tokens}, "
+            f"logprobs={self.logprobs}, "
+            f"prompt_logprobs={self.prompt_logprobs}, "
+            "prompt_logprob_positions="
+            f"{self.prompt_logprob_positions}, "
+            f"skip_special_tokens={self.skip_special_tokens}, "
+            "spaces_between_special_tokens="
+            f"{self.spaces_between_special_tokens}, "
+            f"truncate_prompt_tokens={self.truncate_prompt_tokens}), "
+            f"guided_decoding={self.guided_decoding}")
+
+
+class BeamSearchParams(
+        msgspec.Struct,
+        omit_defaults=True,  # type: ignore[call-arg]
+        # required for @cached_property.
+        dict=True):  # type: ignore[call-arg]
+    """Beam search parameters for text generation."""
+    beam_width: int
+    max_tokens: int
+    ignore_eos: bool = False
+    temperature: float = 0.0
+    length_penalty: float = 1.0