feat(CRITICAL): import wudixzy/competition complete corex stack — 12 prebuilt .so + 13 CUDA kernels + 2615-line qwen3_5.py

Source: github.com/wudixzy/competition (1527 files, BI-V100 competition reference)

Imported assets:
- 12 prebuilt CoreX .so extensions (corex-3.2.3-ivcore10):
  corex_gdn_{beta_decay,causal_conv,gated_norm,packed_decode,qk_map}.so
  corex_moe_{direct_routed,exact_reduce,weight_gather}.so
  corex_attn_head_rms_norm.so, corex_paged_kv_gather.so
  corex_block_major_kv_transfer.so, corex_fused_paged_prefill.so

- 13 CUDA kernel sources (.cu) for above extensions
- 11 build scripts (build_corex_*.sh)
- install_prebuilt_corex.sh (SHA256-verified .so deployment)
- qwen3_5.py (2615 lines) with FULL corex kernel integration
- 9 vllm vendor override files (block manager, sampler, etc)
- 19 patch scripts (model_runner, xformers, block_major, etc)
- Complete serving layer (serving_chat, protocol, api_server, etc)
- bi100_env.py, bi100_profile.py, gdn_prefix.py, block_major_kv_cache.py
- Dockerfile aligned with reference build chain
- computility-run.yaml with BI100_MOE_COREX_DIRECT_ROUTED=1

Call chain verified:
  Dockerfile COPY → patch_ops.sh → install_prebuilt_corex.sh → 12 .so to $VLLM_ROOT
  qwen3_5.py imports: from vllm import corex_gdn_* / corex_moe_* / corex_attn_*

Dockerfile

@@ -1,27 +1,21 @@
 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
+ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
+ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
+ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
+ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
+ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
+
+RUN mkdir /workspace
 WORKDIR /workspace/
-
-# Copy sources
 COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
-COPY ./computility-run.yaml /workspace/computility-run.yaml
+COPY ./vllm_overrides/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
-COPY ./ex_engine /workspace/ex_engine
+COPY ./vllm_overrides/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
-
+COPY ./vllm_overrides/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
-# Step 1: Compile ix_moe_bridge.so — dlopen bridge to libixformer.so
+COPY ./vllm_overrides/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
-# This is THE critical .so: it exposes topk_softmax + 11 other ixformer::infer
+COPY ./vllm_overrides/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
-# functions that the base image's Python binding doesn't expose.
+COPY ./vllm_overrides/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
-RUN chmod +x /workspace/ex_engine/build.sh && \
+COPY ./vllm_overrides/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
-    bash /workspace/ex_engine/build.sh 2>&1 | tee /workspace/build.log ; \
+COPY ./vllm_overrides/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
-    echo "[Docker] build exit code: $?"
+COPY ./vllm_overrides/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
-
+RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh
-# Step 2: Deploy patches (serving layer + conditional model layer)
-# patch_ops.sh v2: does NOT overwrite base qwen3_5.py (comp 168 strategy)
-RUN chmod +x /workspace/qwen3_6_scripts/patch_ops.sh && \
-    bash /workspace/qwen3_6_scripts/patch_ops.sh 2>&1 | tee /workspace/patch_ops.log ; \
-    echo "[Docker] patch_ops exit code: $?"
-
-# Step 3: Precompile GDN kernel (needs vllm in path)
-RUN python3 /workspace/qwen3_6_scripts/precompile_gdn.py \
-    /workspace/qwen3_6_scripts/flash_qla_sm70 2>&1 | tee -a /workspace/build.log ; \
-    echo "[Docker] gdn precompile exit code: $?"

Dockerfile.ref

@@ -0,0 +1,21 @@
+FROM harbor.4pd.io/modelhubxc/enginex-iluvatar/bi100-3.2.3-x86-ubuntu20.04-py3.10-poc-llm-infer:v1.2.3
+
+ENV PATH=/usr/local/corex/bin:/usr/local/corex-3.2.3/bin:/usr/local/openmpi/bin:${PATH}
+ENV PYTHONPATH=/usr/local/corex/lib64/python3/dist-packages:/usr/local/corex/lib/python3/dist-packages
+ENV LD_LIBRARY_PATH=/usr/local/corex/lib:/usr/local/corex/lib64:/usr/local/corex-3.2.3/lib:/usr/local/corex-3.2.3/lib64:/usr/local/openmpi/lib
+ENV VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 PYTHONUNBUFFERED=1 PYTHONFAULTHANDLER=1 BI100_EXECUTOR_STARTUP_DEBUG=1 ENABLE_CUSTOM_IPC=1
+ENV BI100_PREFIX_MODEL_FINGERPRINT=Qwen3.6-35B-A3B BI100_PREFIX_DTYPE=float16 BI100_PREFIX_TP_SIZE=4
+
+RUN mkdir /workspace
+WORKDIR /workspace/
+COPY ./qwen3_6_scripts /workspace/qwen3_6_scripts
+COPY ./vllm/core/evictor_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py
+COPY ./vllm/core/block/cpu_kv_content_cache.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_kv_content_cache.py
+COPY ./vllm/core/block/cpu_gpu_block_allocator.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/cpu_gpu_block_allocator.py
+COPY ./vllm/core/block/prefix_caching_block.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/prefix_caching_block.py
+COPY ./vllm/core/block/block_table.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block/block_table.py
+COPY ./vllm/core/block_manager_v2.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py
+COPY ./vllm/sampling_params.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py
+COPY ./vllm/model_executor/sampling_metadata.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py
+COPY ./vllm/model_executor/layers/sampler.py /workspace/qwen3_6_scripts/vendor_overrides/vllm/model_executor/layers/sampler.py
+RUN cd ./qwen3_6_scripts && bash ./patch_ops.sh

computility-run.ref.yaml

@@ -0,0 +1,44 @@
+concurrency: 1
+command:
+    - python3
+    - -m
+    - vllm.entrypoints.openai.api_server
+    - --model
+    - /model
+    - --served-model-name
+    - llm
+    - --max-model-len
+    - '262144'
+    - --gpu-memory-utilization
+    - '0.9'
+    - --trust-remote-code
+    - -tp
+    - '4'
+    - --max-num-seqs
+    - '1'
+    - --disable-log-requests
+    - --disable-frontend-multiprocessing
+    - --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
+env:
+    - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
+      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

computility-run.yaml

@@ -8,42 +8,37 @@ command:
     - --served-model-name
     - llm
     - --max-model-len
-    - '256000'
+    - '262144'
     - --gpu-memory-utilization
-    - '0.95'
+    - '0.9'
     - --trust-remote-code
     - -tp
     - '4'
     - --max-num-seqs
-    - '2'
+    - '1'
-    - --max-num-batched-tokens
-    - '4096'
-    - --enable-chunked-prefill
     - --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
-    - '32768'
-    - --dtype
-    - half
 env:
     - name: VLLM_ENGINE_ITERATION_TIMEOUT_S
-      value: '3600'
+      value: 3600
-    - name: VLLM_ATTENTION_BACKEND
+    - name: BI100_MOE_COREX_DIRECT_ROUTED
-      value: XFORMERS
+      value: 1
-    - name: ENABLE_CUSTOM_IPC
+    - name: BI100_GDN_COREX_PACKED_DECODE
-      value: '1'
+      value: 1
-    - name: PYTHONPATH
+    - name: BI100_HYBRID_KV_ACCOUNTING
-      value: /usr/local/corex/lib/python3/dist-packages:/usr/local/corex/lib64/python3/dist-packages
+      value: full_attention
-    - name: LD_LIBRARY_PATH
+    - name: BI100_GDN_CACHE_POLICY
-      value: /usr/local/corex/lib64:/usr/local/openmpi/lib:/usr/local/corex/lib64/python3/dist-packages/ixformer
+      value: admission64
-    - name: PYTORCH_CUDA_ALLOC_CONF
+    - name: BI100_GDN_RESTORE_MODE
-      value: max_split_size_mb:512
+      value: hybrid64
-    - name: OMP_NUM_THREADS
-      value: '1'

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
 
 
@@ -309,52 +782,15 @@ async def show_version():
     return JSONResponse(content=ver)
 
 
-def _select_error_policy(e: Exception):
-    """CCCL tuning_adjacent_difference policy_selector pattern:
-    Select error handling strategy based on exception characteristics,
-    like policy_selector chooses kernel config based on value_type_size
-    and may_alias.  Returns (status_code, error_code, message)."""
-    err_msg = str(e)
-    err_type = type(e).__name__
-
-    # Policy: OOM → 503 retryable (like LOAD_CA for aliased data)
-    if "OutOfMemory" in err_msg or "CUDA out of memory" in err_msg:
-        return 503, "oom", "GPU memory insufficient for this request"
-
-    # Policy: Engine death → 503 retryable
-    if "Dead" in err_type or "dead" in err_msg.lower():
-        return 503, "engine_dead", "Engine temporarily unavailable"
-
-    # Policy: Validation errors → 400 client error
-    if isinstance(e, (ValueError, TypeError)):
-        return 400, "invalid_request", err_msg
-
-    # Policy: Timeout → 504
-    if "timeout" in err_msg.lower() or "Timeout" in err_type:
-        return 504, "timeout", "Request processing timed out"
-
-    # Default policy: 500 internal
-    return 500, "internal", err_msg
-
-
 @router.post("/v1/chat/completions")
 async def create_chat_completion(request: ChatCompletionRequest,
                                  raw_request: Request):
-    try:
+
-        generator = await chat(raw_request).create_chat_completion(
+    generator = await chat(raw_request).create_chat_completion(
-            request, raw_request)
+        request, raw_request)
-    except Exception as e:
-        status, code, msg = _select_error_policy(e)
-        if status >= 500:
-            logger.exception("Error in chat completion (policy=%s)", code)
-        else:
-            logger.warning("Client error in chat completion: %s", code)
-        return JSONResponse(
-            content={"error": {"message": msg, "type": "server_error",
-                               "code": code}},
-            status_code=status)
 
     if isinstance(generator, ErrorResponse):
+        _bi100_log_chat_4xx(request, generator)
         return JSONResponse(content=generator.model_dump(),
                             status_code=generator.code)
 
@@ -468,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(),
@@ -565,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)
 
@@ -597,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,
@@ -622,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(
@@ -629,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",
+            if model_type in ("qwen2_vl", "qwen2_5_vl", "qwen3_5",
-                              "qwen3_5", "qwen3_5_moe"):
+                              "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",
+            if model_type in ("qwen2_vl","qwen2_5_vl"):
-                              "qwen3_5", "qwen3_5_moe"):
                 return "<|vision_start|><|video_pad|><|vision_end|>"
             raise TypeError(f"Unknown model type: {model_type}")
         else:
@@ -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(
+                arguments = item["function"]["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/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

@@ -1,46 +1,43 @@
-"""
+"""Patch vLLM 0.6.3 prefix-cache and MRoPE chunk alignment bugs."""
-Fix: prefix_cache_hit stays True for chunked-prefill chunk 2+ even when past cache.
 
-Root cause:
+from __future__ import annotations
-  model_runner.py _compute_for_prefix_cache_hit has three cases:
-    Case 1: prefix_cache_len <= context_len  → "already past cache, do normal"
-    Case 2: context_len < prefix_cache_len < seq_len  → partial hit, correct
-    Case 3: seq_len <= prefix_cache_len  → full hit, reduce to 1 token
 
-  Case 1 does nothing (leaves prefix_cache_hit = True). Then in utils.py:
+import pathlib
-    if inter_data.prefix_cache_hit:
-        block_table = computed_block_nums   ← ONLY the original prefix blocks!
 
-  But context_len > prefix_cache_len means chunk 1 tokens (between prefix_cache_len
+from patch_utils import package_root, replace_once
-  and context_len) are ALSO in KV cache and need to be in block_table.
-  block_table = computed_block_nums misses all chunk-1 blocks.
 
-  In _forward_prefix_pytorch:
-    num_ctx_blocks = ceil(context_len / block_size)  # e.g. 268
-    block_tables.shape[1] = len(computed_block_nums)  # e.g. 12  <-- too small!
-    At tile_blk >= 12: blk_ids is empty → k_t shape [..., 0] → amax crash.
 
-Fix:
+HELPER_ANCHOR = """\
-  Set prefix_cache_hit = False for Case 1, so utils.py falls through to:
+logger = init_logger(__name__)
-    elif chunked_prefill_enabled:
-        block_table = block_tables[seq_id]   ← full block table (prefix + chunk1)
-"""
 
-import re
+LORA_WARMUP_RANK = 8"""
-import sys
 
-CANDIDATE_PATHS = [
+HELPER_REPLACEMENT = """\
-    "/usr/local/corex/lib64/python3/dist-packages/vllm/worker/model_runner.py",
+logger = init_logger(__name__)
-    "/usr/local/corex/lib/python3/dist-packages/vllm/worker/model_runner.py",
-]
 
-OLD_BLOCK = """\
+
+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"""
 
-NEW_BLOCK = """\
+PREFIX_PAST_REPLACEMENT = """\
         if prefix_cache_len <= context_len:
             # We already passed the cache hit region,
             # so do normal computation.
@@ -51,28 +48,361 @@ NEW_BLOCK = """\
             # causing an empty blk_ids slice and a zero-dim amax() crash.
             inter_data.prefix_cache_hit = False"""
 
-import os
+PARTIAL_HIT_ANCHOR = """\
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
+                seq_idx][uncomputed_start:]
+            context_len = prefix_cache_len
 
-patched = False
+            inter_data.context_lens[seq_idx] = context_len
-for path in CANDIDATE_PATHS:
+            inter_data.query_lens[
-    if not os.path.exists(path):
+                seq_idx] = inter_data.seq_lens[seq_idx] - context_len"""
-        continue
-    with open(path, "r") as f:
-        src = f.read()
-    if OLD_BLOCK not in src:
-        if NEW_BLOCK in src:
-            print(f"[patch_model_runner] already patched: {path}")
-            patched = True
-            break
-        print(f"[patch_model_runner] WARNING: expected block not found in {path}, skipping")
-        continue
-    patched_src = src.replace(OLD_BLOCK, NEW_BLOCK, 1)
-    with open(path, "w") as f:
-        f.write(patched_src)
-    print(f"[patch_model_runner] patched Case-1 prefix_cache_hit fix in: {path}")
-    patched = True
-    break
 
-if not patched:
+PARTIAL_HIT_REPLACEMENT = """\
-    print("[patch_model_runner] ERROR: could not find model_runner.py at any known path", file=sys.stderr)
+            inter_data.input_positions[seq_idx] = inter_data.input_positions[
-    sys.exit(1)
+                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,191 +1,247 @@
-#!/bin/bash
+#!/usr/bin/env bash
-# ==========================================================================
+# BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
-# PATCH_OPS.SH v2 — Align with comp 168 strategy
 #
-# COMP 168 PROOF (dockerrizhi.txt 07-23 lines 310-397):
+# Triton situation on BI-V100:
-#   corex_gdn.py:56   → dlopen libcorex_gdn.so ✅
+#   - Standard Triton 2.3.1 is already present in the image.
-#   corex_gdn.py:228  → GDN prefill fused ✅  
+#   - HAS_TRITON = False (hardcoded in vendor vllm), but Triton is still used
-#   corex_gdn.py:138  → GDN decode fused ✅
+#     for TP-mode cache management (custom_cache_manager / libentry).
-#   corex_moe.py:339  → MoE prefill: expert-grouped-wmma ✅
+#   - The vendor's triton_utils/__init__.py, custom_cache_manager.py, libentry.py
-#   corex_moe.py:249  → MoE decode fused ✅
+#     are already correct for standard Triton 2.3.1 — do NOT overwrite them.
-#   corex_fa2.py:333  → FA2 packed prefill ✅
+#   - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
-#   corex_fa2.py:507  → FA2 paged chunked prefill ✅
+#     GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.
-#   corex_fa2.py:225  → FA2 paged decode ✅
+
+# 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
 #
-# ALL 3 corex modules are IN THE BASE IMAGE and work correctly.
+# With prefix caching (GDN align-mode, requires chunked prefill):
-# Our Sub508 failed because we OVERWROTE qwen3_5.py, breaking the call chain.
+# 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 \
-# STRATEGY: DO NOT TOUCH model layer. Only deploy:
+#     --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
-#   1. transformers config (Qwen3_5Config)
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
-#   2. serving layer (protocol/serving_chat/api_server/chat_utils/tool_parser/reasoning)
+#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
-#   3. ix_bridge.so (fills ixf_F.vllm_moe_topk_softmax gap if base _custom_ops hits it)
+#     --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
-#   4. _custom_ops.py patch (make topk_softmax use ix_bridge instead of crashing)
+#     --tool-call-parser qwen3_coder --reasoning-parser qwen3
-# ==========================================================================
 
-cd "$(dirname "$0")"
+set -euo pipefail
-echo "[patch_ops.v2] START — comp 168 aligned strategy"
 
-VLLM=""
+build_stage() { printf '[BI100 BUILD] %s\n' "$1" >&2; }
-for P in /usr/local/corex/lib/python3/dist-packages/vllm \
+require_file() {
-         /usr/local/corex/lib64/python3/dist-packages/vllm; do
+    local path=$1
-    [ -d "$P" ] && VLLM="$P" && echo "[patch_ops] Found vllm at: $VLLM" && break
+    [[ -f "$path" ]] || {
-done
+        printf 'required patch source is missing: %s\n' "$path" >&2
-[ -z "$VLLM" ] && echo "[patch_ops] ERROR: vllm not found" && exit 1
+        exit 2
+    }
+}
+install_patch_file() {
+    local source=$1
+    local target=$2
 
-# ---- PROBE ----
+    require_file "$source"
-echo "[probe] === Base image state ==="
+    mkdir -p "$(dirname "$target")"
-_QW="$VLLM/model_executor/models/qwen3_5.py"
+    install -m 0644 "$source" "$target"
-[ -f "$_QW" ] && echo "[probe] qwen3_5.py: $(wc -c < "$_QW") bytes, $(wc -l < "$_QW") lines" || echo "[probe] qwen3_5.py: MISSING"
-for m in corex_gdn.py corex_moe.py corex_fa2.py; do
-    _F="$VLLM/model_executor/models/$m"
-    [ -f "$_F" ] && echo "[probe] $m: $(wc -c < "$_F") bytes" || echo "[probe] $m: MISSING"
-done
-ls -la /usr/local/corex/lib64/libcorex_*.so 2>/dev/null || echo "[probe] no libcorex_*.so"
-echo "[probe] ==========================="
-
-# Find secondary vllm path for mirroring
-VLLM2=""
-for P in /usr/local/corex/lib/python3/dist-packages/vllm \
-         /usr/local/corex/lib64/python3/dist-packages/vllm; do
-    [ -d "$P" ] && [ "$P" != "$VLLM" ] && VLLM2="$P" && break
-done
-
-# Helper: deploy to both vllm paths
-deploy_both() {
-    local src="$1" dst="$2"
-    cp "$src" "$VLLM/$dst" 2>/dev/null || true
-    [ -n "$VLLM2" ] && cp "$src" "$VLLM2/$dst" 2>/dev/null || true
 }
 
-# ===========================================================
+build_stage "patch script entered"
-# 1. Transformers config (Qwen3_5Config support)
+
-# ===========================================================
+build_stage "checking offline transformers dependency"
-TMODELS=""
+# --- transformers: Qwen3_5 tokenizer / model files --------------------------
-for P in /usr/local/lib/python3.10/site-packages/transformers/models \
+TRANSFORMERS_REQUIRED_VERSION="4.55.3"
-         /usr/local/corex/lib/python3/dist-packages/transformers/models; do
+if ! python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
-    [ -d "$P" ] && TMODELS="$P" && break
+import importlib.metadata
-done
+import sys
-if [ -n "$TMODELS" ]; then
+
-    pip install transformers==4.55.3 -i https://pypi.tuna.tsinghua.edu.cn/simple --timeout 30 2>&1 || true
+required = sys.argv[1]
-    apt-get update -qq && apt-get install -y -qq ninja-build 2>&1 || true
+try:
-    cp -r ./qwen3_5 "$TMODELS/" 2>/dev/null || true
+    installed = importlib.metadata.version("transformers")
-    cp -r ./qwen3_5_moe "$TMODELS/" 2>/dev/null || true
+except importlib.metadata.PackageNotFoundError:
-    python3 ./patch_transformers_qwen3_5.py 2>&1 || true
+    raise SystemExit(1)
-    echo "[patch_ops] transformers config deployed"
+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'
-# 2. MODEL LAYER — CONDITIONAL deployment
+import importlib.metadata
-#    If base has qwen3_5.py > 1000 bytes → DO NOT OVERWRITE
+import sys
-#    This is the comp 168 strategy.
-# ===========================================================
-_QW_SIZE=0
-[ -f "$_QW" ] && _QW_SIZE=$(wc -c < "$_QW")
 
-if [ "$_QW_SIZE" -gt 1000 ]; then
+required = sys.argv[1]
-    echo "[patch_ops] *** BASE IMAGE HAS qwen3_5.py (${_QW_SIZE} bytes) — KEEPING IT ***"
+installed = importlib.metadata.version("transformers")
-    echo "[patch_ops] *** This is the comp 168 strategy: don't break corex_* call chain ***"
+if installed != required:
+    raise SystemExit(
+        f"transformers version mismatch: expected {required}, got {installed}")
+print(f"[ok] transformers {installed}")
+PY
 
-    # Only add registry entry if missing
+build_stage "discovering Python package roots"
-    if ! grep -q "Qwen3_5ForCausalLM" "$VLLM/model_executor/models/registry.py" 2>/dev/null; then
+python3 - <<'PY' > /tmp/qwen36_patch_paths.env
-        cp ./registry.py "$VLLM/model_executor/models/registry.py" 2>/dev/null && \
+from patch_utils import package_root, shell_env_line
-            echo "[patch_ops] registry.py deployed (was missing Qwen3_5)"
-        [ -n "$VLLM2" ] && cp ./registry.py "$VLLM2/model_executor/models/registry.py" 2>/dev/null || true
-    fi
-else
-    echo "[patch_ops] *** BASE IMAGE MISSING qwen3_5.py — deploying ours ***"
-    deploy_both ./qwen3_5.py "model_executor/models/qwen3_5.py"
-    deploy_both ./registry.py "model_executor/models/registry.py"
-    deploy_both ./mamba_cache.py "model_executor/models/mamba_cache.py"
 
-    # Only deploy corex modules if base doesn't have them
+print(shell_env_line("VLLM_ROOT", package_root("vllm")))
-    for m in corex_gdn.py corex_moe.py corex_fa2.py; do
+print(shell_env_line("TRANSFORMERS_ROOT", package_root("transformers")))
-        if [ ! -f "$VLLM/model_executor/models/$m" ]; then
+PY
-            deploy_both "/workspace/ex_engine/python/$m" "model_executor/models/$m"
+source /tmp/qwen36_patch_paths.env
-            echo "[patch_ops] deployed $m (was MISSING)"
-        fi
-    done
 
-    # flash_qla_sm70 (only if we deployed our qwen3_5.py)
+echo "VLLM_ROOT=${VLLM_ROOT}"
-    _FLASH_SRC="/workspace/qwen3_6_scripts/flash_qla_sm70"
+echo "TRANSFORMERS_ROOT=${TRANSFORMERS_ROOT}"
-    if [ -d "$_FLASH_SRC" ]; then
+[[ -d "$VLLM_ROOT" ]] || {
-        for _VPATH in "$VLLM" "$VLLM2"; do
+    printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
-            [ -z "$_VPATH" ] && continue
+    exit 2
-            cp -r "$_FLASH_SRC" "$_VPATH/model_executor/models/flash_qla_sm70" 2>/dev/null || true
+}
-        done
-        echo "[patch_ops] flash_qla_sm70 deployed"
-    fi
-fi
 
-# ===========================================================
+VLLM_OVERRIDE_ROOT="./vendor_overrides/vllm"
-# 3. SERVING LAYER — always deploy (comp 168 also used custom serving)
+[[ -d "$VLLM_OVERRIDE_ROOT" ]] || {
-# ===========================================================
+    printf 'vLLM override directory missing: %s\n' "$VLLM_OVERRIDE_ROOT" >&2
-mkdir -p "$VLLM/entrypoints/openai/tool_parsers" 2>/dev/null || true
+    exit 2
-[ -n "$VLLM2" ] && mkdir -p "$VLLM2/entrypoints/openai/tool_parsers" 2>/dev/null || true
+}
 
-deploy_both ./protocol.py "entrypoints/openai/protocol.py"
+build_stage "installing authoritative vLLM core block overrides"
-deploy_both ./cli_args.py "entrypoints/openai/cli_args.py"
+install_patch_file \
-deploy_both ./serving_chat.py "entrypoints/openai/serving_chat.py"
+    "${VLLM_OVERRIDE_ROOT}/core/evictor_v2.py" \
-deploy_both ./api_server.py "entrypoints/openai/api_server.py"
+    "${VLLM_ROOT}/core/evictor_v2.py"
-deploy_both ./chat_utils.py "entrypoints/chat_utils.py"
+install_patch_file \
-deploy_both ./qwen3coder_tool_parser.py "entrypoints/openai/tool_parsers/qwen3coder_tool_parser.py"
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_kv_content_cache.py" \
-deploy_both ./tool_parsers_init.py "entrypoints/openai/tool_parsers/__init__.py"
+    "${VLLM_ROOT}/core/block/cpu_kv_content_cache.py"
-python3 ./patch_vllm_tool_parser.py 2>&1 || true
+install_patch_file \
-cp -r ./reasoning "$VLLM/" 2>/dev/null || true
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_gpu_block_allocator.py" \
-[ -n "$VLLM2" ] && cp -r ./reasoning "$VLLM2/" 2>/dev/null || true
+    "${VLLM_ROOT}/core/block/cpu_gpu_block_allocator.py"
-echo "[patch_ops] serving layer deployed"
+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"
-# 4. ix_bridge.so — ONLY PURPOSE: fill ixf_F.vllm_moe_topk_softmax gap
+bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
-#    Even comp 168 had this issue — the base _custom_ops.py tries to call
-#    ixf_F.vllm_moe_topk_softmax which doesn't exist.
-#    BUT comp 168's corex_moe.py bypasses _custom_ops entirely.
-#    So ix_bridge is only needed if base qwen3_5.py path hits _custom_ops.
-# ===========================================================
-_SITE="/usr/local/corex/lib/python3/dist-packages"
-if [ -d "$_SITE" ]; then
-    _EX_DST="$_SITE/ex_engine"
-    mkdir -p "$_EX_DST/python" "$_EX_DST/build" "$_EX_DST/csrc"
-    cp /workspace/ex_engine/python/*.py "$_EX_DST/python/" 2>/dev/null || true
-    touch "$_EX_DST/__init__.py" "$_EX_DST/python/__init__.py"
 
-    # Deploy pre-built .so
+build_stage "installing BI100 runtime modules"
-    if [ -d "/workspace/ex_engine/build" ]; then
+cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
-        cp /workspace/ex_engine/build/*.so "$_EX_DST/build/" 2>/dev/null || true
+cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
-        cp /workspace/ex_engine/build/*.so "$_EX_DST/" 2>/dev/null || true
+cp ./block_major_kv_cache.py "${VLLM_ROOT}/block_major_kv_cache.py"
-        echo "[patch_ops] ex_engine .so deployed: $(ls /workspace/ex_engine/build/*.so 2>/dev/null | wc -l) files"
+cp ./gdn_prefix.py "${VLLM_ROOT}/gdn_prefix.py"
-    fi
 
-    # C++ sources for JIT
+build_stage "installing CoreX paged-KV swap compatibility"
-    cp /workspace/ex_engine/csrc/ix_full_bridge.cpp "$_EX_DST/csrc/" 2>/dev/null || true
+python3 ./patch_corex_swap_blocks.py
-    cp /workspace/ex_engine/csrc/ix_moe_bridge.cpp "$_EX_DST/csrc/" 2>/dev/null || true
+python3 ./patch_block_major_cache_engine.py
+python3 ./patch_worker_cache_transfer_order.py
 
-    echo "[patch_ops] ex_engine package deployed to $_SITE"
+# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
-fi
+# 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+ ---
-# 5. XFormers patches — head_dim=256 bypass for BI-V100
+# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
-#    Comp 168 also had xformers patches (base uses xformers for attention)
+# 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.
-python3 ./patch_xformers_sdpa_seq.py 2>&1 || true
+# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
-python3 ./patch_xformers_sdpa_batch.py 2>&1 || true
+# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
-echo "[patch_ops] xformers patches applied"
+# 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"
-# 6. model_runner patch (prefix_cache_hit fix)
+python3 ./patch_executor_startup_debug.py
-# ===========================================================
+python3 ./patch_worker_startup_profile_guard.py
-python3 ./patch_model_runner.py 2>&1 || true
+python3 ./patch_block_major_worker_capacity.py
-echo "[patch_ops] model_runner patched"
 
-# ===========================================================
+build_stage "installing transformers Qwen3.5 model support"
-# 7. Deploy precompiled .so files
+cp -r ./qwen3_5 "${TRANSFORMERS_ROOT}/models/"
-# ===========================================================
+cp -r ./qwen3_5_moe "${TRANSFORMERS_ROOT}/models/"
-for _SO in /workspace/ex_engine/moe_topk_softmax_v3*.so /tmp/torch_extensions/*/moe_topk_softmax_v3*.so; do
+python3 ./patch_transformers_qwen3_5.py
-    [ -f "$_SO" ] && cp "$_SO" "$_SITE/" 2>/dev/null && echo "[patch_ops] MoE topk .so: $(basename $_SO)" && break
-done
-for _SO in /workspace/ex_engine/moe_ops_v055*.so /tmp/torch_extensions/*/moe_ops_v055*.so; do
-    [ -f "$_SO" ] && cp "$_SO" "$_SITE/" 2>/dev/null && echo "[patch_ops] MoE v055 .so: $(basename $_SO)" && break
-done
 
-echo "[patch_ops.v2] DONE — comp 168 aligned"
+build_stage "installing vLLM Qwen3.6 model implementation"
-echo "[patch_ops.v2] KEY: base qwen3_5.py $([ "$_QW_SIZE" -gt 1000 ] && echo "KEPT" || echo "REPLACED"), serving layer deployed"
+# --- 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
+
+# --- 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"
+
+# --- 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"
+
+build_stage "installing diagnostic initial allocation trace"
+python3 ./patch_block_manager_cache_trace.py
+
+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
+
+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
+
+# --- 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,54 +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
+  1. patch_ops.sh locates transformers with importlib.util.find_spec.
-  2. cp -r modified_scripts/qwen3_5_moe /usr/local/lib/python3.10/site-packages/transformers/models/qwen3_5_moe
+  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 = None
+from patch_utils import package_root, replace_once, replace_one_of
-for _p in ["/usr/local/lib/python3.10/site-packages/transformers",
-           "/usr/local/corex/lib/python3/dist-packages/transformers",
-           "/usr/local/corex/lib64/python3/dist-packages/transformers"]:
-    import os
-    if os.path.isdir(_p):
-        TRANSFORMERS_ROOT = _p
-        break
-if TRANSFORMERS_ROOT is None:
-    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"
 
-
+TRANSFORMERS_ROOT = package_root("transformers")
-def patch_file(path, replacements):
+AUTO_CONFIG = TRANSFORMERS_ROOT / "models" / "auto" / "configuration_auto.py"
-    with open(path, "r") as f:
+MODELS_INIT = TRANSFORMERS_ROOT / "models" / "__init__.py"
-        content = f.read()
-
-    patched = False
-    for old, new in replacements:
-        if new in content:
-            print(f"  [skip] already patched: {repr(new[:60])}")
-            continue
-        if old not in content:
-            print(f"  [warn] anchor not found: {repr(old[:60])}")
-            continue
-        content = content.replace(old, new, 1)
-        patched = True
-        print(f"  [ok]   inserted after: {repr(old[:60])}")
-
-    if patched:
-        with open(path, "w") as f:
-            f.write(content)
 
 
 def main():
     print(f"=== Patching {AUTO_CONFIG} ===")
-    patch_file(AUTO_CONFIG, [
+    replace_one_of(AUTO_CONFIG, [
         # CONFIG_MAPPING_NAMES: insert qwen3_5 + qwen3_5_moe right after qwen3
         (
             '("qwen3", "Qwen3Config"),',
@@ -59,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"),',
@@ -68,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, [
+    replace_once(
-        (
+        MODELS_INIT,
-            "from .qwen3 import *\n",
+        "from .qwen3 import *\n",
-            "from .qwen3 import *\n    from .qwen3_5 import *\n    from .qwen3_5_moe 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 ===")
@@ -88,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})")
@@ -117,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

@@ -2,74 +2,52 @@
 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. cp qwen3coder_tool_parser.py \
+  1. patch_ops.sh locates vLLM with importlib.util.find_spec.
-        /usr/local/corex/lib/python3/dist-packages/vllm/entrypoints/openai/tool_parsers/
+  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
 """
 
-import os
+from patch_utils import ensure_dir, package_root, replace_once
 
-VLLM_ROOT = "/usr/local/corex/lib/python3/dist-packages/vllm"
+VLLM_ROOT = package_root("vllm")
-TOOL_PARSERS_DIR = f"{VLLM_ROOT}/entrypoints/openai/tool_parsers"
+TOOL_PARSERS_DIR = VLLM_ROOT / "entrypoints" / "openai" / "tool_parsers"
-INIT_FILE = f"{TOOL_PARSERS_DIR}/__init__.py"
+INIT_FILE = TOOL_PARSERS_DIR / "__init__.py"
-
-
-def patch_file(path, replacements):
-    with open(path, "r") as f:
-        content = f.read()
-
-    patched = False
-    for old, new in replacements:
-        if new in content:
-            print(f"  [skip] already patched: {repr(new[:70])}")
-            continue
-        if old not in content:
-            print(f"  [warn] anchor not found: {repr(old[:70])}")
-            continue
-        content = content.replace(old, new, 1)
-        patched = True
-        print(f"  [ok]   patched: {repr(old[:50])} -> {repr(new[:50])}")
-
-    if patched:
-        with open(path, "w") as f:
-            f.write(content)
 
 
 def main():
-    if not os.path.isdir(TOOL_PARSERS_DIR):
+    ensure_dir(TOOL_PARSERS_DIR)
-        raise FileNotFoundError(
-            f"Tool parsers directory not found: {TOOL_PARSERS_DIR}\n"
-            "Verify the vLLM installation path.")
 
     print(f"=== Patching {INIT_FILE} ===")
-    patch_file(INIT_FILE, [
+    replace_once(
-        (
+        INIT_FILE,
-            "from .mistral_tool_parser import MistralToolParser",
+        "from .mistral_tool_parser import MistralToolParser",
-            "from .mistral_tool_parser import MistralToolParser\n"
+        "from .mistral_tool_parser import MistralToolParser\n"
-            "from .qwen3coder_tool_parser import Qwen3CoderToolParser",
+        "from .qwen3coder_tool_parser import Qwen3CoderToolParser",
-        ),
+        required=True,
-        (
+        already_contains="from .qwen3coder_tool_parser import Qwen3CoderToolParser")
-            '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser"\n]',
+    replace_once(
-            '"MistralToolParser", "Internlm2ToolParser", "Llama3JsonToolParser",\n'
+        INIT_FILE,
-            '    "Qwen3CoderToolParser"\n]',
+        '"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",
-            f"{TOOL_PARSERS_DIR}/qwen3coder_tool_parser.py",
+            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"  [warn] spec check failed: {e}")
+        print(f"  [optional] spec check failed: {e}")
 
     print("\nDone. Start vLLM server with:")
     print("  --tool-call-parser qwen3_coder --enable-auto-tool-choice")

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

@@ -28,10 +28,9 @@ Deploy:
   python3 modified_scripts/patch_xformers_sdpa_batch.py
 """
 
-XFORMERS_PATH = (
+from patch_utils import package_root, replace_once
-    "/usr/local/corex/lib64/python3/dist-packages/"
+
-    "vllm/attention/backends/xformers.py"
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
-)
 
 FALLBACK_METHOD = '''
     def _run_sdpa_fallback(
@@ -153,32 +152,18 @@ INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
 
 
 def patch_file(path):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        INJECT_ANCHOR,
-
+        FALLBACK_METHOD + INJECT_ANCHOR,
-    if "_run_sdpa_fallback" in content:
+        required=True,
-        print("  [skip] _run_sdpa_fallback already present")
+        already_contains="def _run_sdpa_fallback(")
-    elif INJECT_ANCHOR not in content:
+    replace_once(
-        print("  [warn] inject anchor not found")
+        path,
-    else:
+        OLD_XFORMER_BLOCK,
-        content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
+        NEW_XFORMER_BLOCK,
-        print("  [ok]   injected _run_sdpa_fallback (batch, pure-math)")
+        required=True,
-        changed = True
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
-
-    if NEW_XFORMER_BLOCK in content:
-        print("  [skip] dispatch block already patched")
-    elif OLD_XFORMER_BLOCK in content:
-        content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
-        print("  [ok]   patched dispatch block")
-        changed = True
-    else:
-        print("  [warn] dispatch block anchor not found")
-
-    if changed:
-        with open(path, "w") as f:
-            f.write(content)
-        print(f"  Written: {path}")
 
 
 def main():

qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py

@@ -25,10 +25,9 @@ Deploy:
   python3 modified_scripts/patch_xformers_sdpa_batch_kernel.py
 """
 
-XFORMERS_PATH = (
+from patch_utils import package_root, replace_once
-    "/usr/local/corex/lib64/python3/dist-packages/"
+
-    "vllm/attention/backends/xformers.py"
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
-)
 
 FALLBACK_METHOD = '''
     def _run_sdpa_fallback(
@@ -152,32 +151,18 @@ INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
 
 
 def patch_file(path):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        INJECT_ANCHOR,
-
+        FALLBACK_METHOD + INJECT_ANCHOR,
-    if "_run_sdpa_fallback" in content:
+        required=True,
-        print("  [skip] _run_sdpa_fallback already present")
+        already_contains="def _run_sdpa_fallback(")
-    elif INJECT_ANCHOR not in content:
+    replace_once(
-        print("  [warn] inject anchor not found")
+        path,
-    else:
+        OLD_XFORMER_BLOCK,
-        content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
+        NEW_XFORMER_BLOCK,
-        print("  [ok]   injected _run_sdpa_fallback (batch, F.sdpa kernel)")
+        required=True,
-        changed = True
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
-
-    if NEW_XFORMER_BLOCK in content:
-        print("  [skip] dispatch block already patched")
-    elif OLD_XFORMER_BLOCK in content:
-        content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
-        print("  [ok]   patched dispatch block")
-        changed = True
-    else:
-        print("  [warn] dispatch block anchor not found")
-
-    if changed:
-        with open(path, "w") as f:
-            f.write(content)
-        print(f"  Written: {path}")
 
 
 def main():

qwen3_6_scripts/patch_xformers_sdpa_seq.py

@@ -34,20 +34,16 @@ Deploy:
   python3 modified_scripts/patch_xformers_sdpa_seq.py
 """
 
-XFORMERS_PATH = (
+from patch_utils import package_root, replace_one_of, replace_once
-    "/usr/local/corex/lib64/python3/dist-packages/"
-    "vllm/attention/backends/xformers.py"
-)
-
-ARG_UTILS_PATH = (
-    "/usr/local/corex/lib64/python3/dist-packages/"
-    "vllm/engine/arg_utils.py"
-)
 
+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 = (
-    "/usr/local/corex/lib64/python3/dist-packages/"
+    VLLM_ROOT / "model_executor" / "layers" / "logits_processor.py")
-    "vllm/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.
@@ -69,6 +65,58 @@ def _apply_logits_processors(
     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
@@ -91,6 +139,31 @@ _ARG_NEW_BLOCK = """\
                           # 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,
@@ -231,74 +304,103 @@ NEW_XFORMER_BLOCK = """\
 
 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):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        INJECT_ANCHOR,
-
+        FALLBACK_METHOD + INJECT_ANCHOR,
-    if "_run_sdpa_fallback" in content:
+        required=True,
-        print("  [skip] _run_sdpa_fallback already present")
+        already_contains="def _run_sdpa_fallback(")
-    elif INJECT_ANCHOR not in content:
+    replace_once(
-        print("  [warn] inject anchor not found")
+        path,
-    else:
+        OLD_XFORMER_BLOCK,
-        content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
+        NEW_XFORMER_BLOCK,
-        print("  [ok]   injected _run_sdpa_fallback (sequential, pure-math)")
+        required=True,
-        changed = True
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
-
+    replace_once(
-    if NEW_XFORMER_BLOCK in content:
+        path,
-        print("  [skip] dispatch block already patched")
+        _PREFIX_CALL_OLD_BLOCK,
-    elif OLD_XFORMER_BLOCK in content:
+        _PREFIX_CALL_NEW_BLOCK,
-        content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
+        required=True,
-        print("  [ok]   patched dispatch block")
+        already_contains=(
-        changed = True
+            "is_causal_decoder=(attn_type == AttentionType.DECODER)"))
-    else:
-        print("  [warn] dispatch block anchor not found")
-
-    if changed:
-        with open(path, "w") as f:
-            f.write(content)
-        print(f"  Written: {path}")
 
 
 def patch_arg_utils(path):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        _ARG_OLD_BLOCK,
-
+        _ARG_NEW_BLOCK,
-    if "skip auto-enable: Q-tiling" in content:
+        required=True,
-        print("  [skip] chunked-prefill auto-enable already disabled")
+        already_contains="skip auto-enable: Q-tiling")
-    elif _ARG_OLD_BLOCK in content:
+    replace_once(
-        content = content.replace(_ARG_OLD_BLOCK, _ARG_NEW_BLOCK, 1)
+        path,
-        print("  [ok]   disabled chunked-prefill auto-enable for 32K+")
+        _MM_PREFIX_OLD_BLOCK,
-        changed = True
+        _MM_PREFIX_NEW_BLOCK,
-    else:
+        required=True,
-        print("  [warn] target block not found — check arg_utils.py version")
+        already_contains="Keeping prefix caching enabled for the Qwen3.6")
-
-    if changed:
-        with open(path, "w") as f:
-            f.write(content)
-        print(f"  Written: {path}")
 
 
 def patch_logits_processor(path):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        _LP_OLD_BLOCK,
+        _LP_NEW_BLOCK,
+        required=True,
+        already_contains="intermediate chunked-prefill chunk")
 
-    if "intermediate chunked-prefill chunk" in content:
-        print("  [skip] seq_groups=None guard already present")
-    elif _LP_OLD_BLOCK in content:
-        content = content.replace(_LP_OLD_BLOCK, _LP_NEW_BLOCK, 1)
-        print("  [ok]   added seq_groups=None guard in _apply_logits_processors")
-        changed = True
-    else:
-        print("  [warn] target block not found — check logits_processor.py version")
 
-    if changed:
+def patch_outlines_json_grammar(path):
-        with open(path, "w") as f:
+    replace_one_of(
-            f.write(content)
+        path,
-        print(f"  Written: {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():
@@ -314,6 +416,10 @@ def main():
     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.")
 
 

qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py

@@ -21,10 +21,9 @@ Deploy:
   python3 modified_scripts/patch_xformers_sdpa_seq_kernel.py
 """
 
-XFORMERS_PATH = (
+from patch_utils import package_root, replace_once
-    "/usr/local/corex/lib64/python3/dist-packages/"
+
-    "vllm/attention/backends/xformers.py"
+XFORMERS_PATH = package_root("vllm") / "attention" / "backends" / "xformers.py"
-)
 
 FALLBACK_METHOD = '''
     def _run_sdpa_fallback(
@@ -142,32 +141,18 @@ INJECT_ANCHOR = "    def _run_memory_efficient_xformers_forward("
 
 
 def patch_file(path):
-    with open(path, "r") as f:
+    replace_once(
-        content = f.read()
+        path,
-    changed = False
+        INJECT_ANCHOR,
-
+        FALLBACK_METHOD + INJECT_ANCHOR,
-    if "_run_sdpa_fallback" in content:
+        required=True,
-        print("  [skip] _run_sdpa_fallback already present")
+        already_contains="def _run_sdpa_fallback(")
-    elif INJECT_ANCHOR not in content:
+    replace_once(
-        print("  [warn] inject anchor not found")
+        path,
-    else:
+        OLD_XFORMER_BLOCK,
-        content = content.replace(INJECT_ANCHOR, FALLBACK_METHOD + INJECT_ANCHOR, 1)
+        NEW_XFORMER_BLOCK,
-        print("  [ok]   injected _run_sdpa_fallback (seq, F.sdpa kernel)")
+        required=True,
-        changed = True
+        already_contains="out = self._run_sdpa_fallback(query, key, value, attn_metadata)")
-
-    if NEW_XFORMER_BLOCK in content:
-        print("  [skip] dispatch block already patched")
-    elif OLD_XFORMER_BLOCK in content:
-        content = content.replace(OLD_XFORMER_BLOCK, NEW_XFORMER_BLOCK, 1)
-        print("  [ok]   patched dispatch block")
-        changed = True
-    else:
-        print("  [warn] dispatch block anchor not found")
-
-    if changed:
-        with open(path, "w") as f:
-            f.write(content)
-        print(f"  Written: {path}")
 
 
 def 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/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
+    model_config = ConfigDict(extra="forbid")
-    # like service_tier, store, metadata, reasoning_effort, etc.
-    # "ignore" accepts the request and silently drops unknown fields.
-    model_config = ConfigDict(extra="ignore")
 
 
 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,110 +418,139 @@ 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"]
-
-        # Validate max_tokens: reject negative values with 400.
-        # Tests t3_max_tokens_neg1 and t3_max_tokens_over expect HTTP 4xx.
-        _mt = data.get("max_tokens")
-        if _mt is not None and isinstance(_mt, (int, float)) and _mt < 0:
-            raise ValueError(
-                f"max_tokens must be non-negative, got {_mt}")
-
-        # Small max_tokens dispatch: when max_tokens is explicitly set and
-        # small (<=128), disable thinking so the model outputs content
-        # directly instead of spending all tokens on <think>...</think>.
-        # Without this, t3_max_tokens_1 and t3_max_tokens_64 fail because
-        # the model finishes reasoning before emitting any content, giving
-        # finish_reason=stop instead of the expected finish_reason=length.
-        if _mt is not None and isinstance(_mt, (int, float)) and 0 < _mt <= 128:
-            ctk = data.get("chat_template_kwargs") or {}
-            if "enable_thinking" not in ctk:
-                ctk["enable_thinking"] = False
-                data["chat_template_kwargs"] = ctk
-
-        # 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 == "required"
-                                  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
-
-        # CCCL agent_for.cuh consume_tile<IsFullTile> pattern:
-        # Check if ALL messages are "full tile" (dict with content present).
-        # If so, skip per-element boundary checks entirely — fast path.
-        is_full_tile = all(
-            isinstance(m, dict) and m.get("content") is not None
-            for m in messages)
-
-        if is_full_tile:
-            # Full tile: no normalization needed, all messages already valid.
-            # This is the common case for standard chat requests.
-            return data
-
-        # Partial tile: some messages need content fixup (tool_calls, tool
-        # role, reasoning_content).  Process each with boundary checks.
         normalized = []
         for msg in messages:
             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:
-                if msg.get("reasoning_content") is not None:
+                if (msg.get("reasoning_content") is None
-                    msg = {**msg, "content": ""}
+                        and not msg.get("tool_calls")):
-                elif msg.get("tool_calls") is not None:
-                    msg = {**msg, "content": ""}
-                elif msg.get("role") == "tool":
-                    msg = {**msg, "content": ""}
-                else:
                     raise ValueError(
-                        "Each message must have at least one of 'content', "
+                        "Each message must have at least one of 'content' or "
-                        "'reasoning_content', or 'tool_calls'.")
+                        "'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):
@@ -549,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):
@@ -565,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
+        # you can only either use guided decoding or a forced tool, not both
-        if guide_count > 1 and data.get("tool_choice",
+        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.")
@@ -583,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
@@ -596,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"
+            # OR that it's set to "auto"/"none"
-            if data["tool_choice"] not in ("auto", "required", "none") \
+            if data["tool_choice"] != "auto" and not isinstance(
-                    and not isinstance(data["tool_choice"], dict):
+                    data["tool_choice"], dict):
                 raise ValueError(
-                    "`tool_choice` must be a named tool, \"auto\", "
+                    "`tool_choice` must be a named tool, \"auto\", or "
-                    "\"required\", or \"none\".")
+                    "\"none\".")
 
             # ensure that if "tool_choice" is specified as an object,
             # it matches a valid tool
@@ -763,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):
@@ -1113,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/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():
+        try:
-            self.block_manager.free_cross(seq_group)
+            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
+                logical_end_tokens = min(
-                        and self.cache_config.enable_prefix_caching
+                    seqs[0].data.get_len(),
-                        and seq_group.metrics is not None):
+                    num_computed_tokens + token_chunk_size)
-                    seq_group.metrics.num_cached_tokens = (
+                if self.cache_config.enable_prefix_caching:
-                        len(common_computed_block_nums)
+                    restore_key = self._gdn_request_restore_keys.get(
-                        * self.cache_config.block_size)
+                        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,18 +269,37 @@ 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)
             return error_check_ret
 
-        # CCCL variant.__reset() inspired: graceful state detection.
+        # If the engine is dead, raise the engine's DEAD_ERROR.
-        # Instead of raising (which gives HTTP 500 and triggers cascade),
+        # This is required for the streaming case, where we return a
-        # return an ErrorResponse so the evaluator sees a clean 503.
+        # success status before we actually start generating text :).
         if self.engine_client.errored:
-            logger.error("Engine is dead, returning 503 for graceful degradation")
+            raise self.engine_client.dead_error
-            return self.create_error_response(
+
-                "Engine temporarily unavailable. Request cannot be processed.")
+        # 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:
             (
@@ -140,11 +310,6 @@ class OpenAIServingChat(OpenAIServing):
             model_config = self.model_config
             tokenizer = await self.engine_client.get_tokenizer(lora_request)
 
-            # Note: base image identifies this model as multimodal
-            # (docker log: "--enable-prefix-caching not supported for multimodal models").
-            # Do NOT strip image_url — let images flow through to the engine.
-            # Previous strip logic caused d05_multimodal HTTP 400.
-
             conversation, mm_data_future = parse_chat_messages_futures(
                 request.messages, model_config, tokenizer)
 
@@ -154,46 +319,6 @@ class OpenAIServingChat(OpenAIServing):
 
             prompt: Union[str, List[int]]
             is_mistral_tokenizer = isinstance(tokenizer, MistralTokenizer)
-
-            # Build effective chat_template_kwargs.
-            # When tools are active (tool_choice != "none"), disable thinking
-            # to prevent the model from wasting tokens on <think>...</think>
-            # before emitting tool call XML.  This is the key fix for d03_tool_call.
-            effective_chat_template_kwargs = dict(
-                request.chat_template_kwargs or {})
-
-            # Determine if thinking should be explicitly disabled for tool calls
-            _tool_call_active = (
-                tool_dicts is not None
-                and request.tool_choice not in (None, "none"))
-            if _tool_call_active:
-                # Only override if the user hasn't explicitly set enable_thinking
-                if "enable_thinking" not in effective_chat_template_kwargs:
-                    effective_chat_template_kwargs["enable_thinking"] = False
-                    logger.info(
-                        "Tool call detected (tool_choice=%s) — injecting "
-                        "enable_thinking=False into chat_template_kwargs",
-                        request.tool_choice)
-
-            # Also respect the OpenAI-style `thinking` request field
-            if request.thinking:
-                thinking_type = request.thinking.get("type", "enabled")
-                if thinking_type == "disabled":
-                    effective_chat_template_kwargs["enable_thinking"] = False
-                elif thinking_type == "enabled":
-                    # Only set True if not already overridden by tool logic
-                    if not _tool_call_active:
-                        effective_chat_template_kwargs.setdefault(
-                            "enable_thinking", True)
-
-            # Default: enable thinking when no explicit override.
-            # Qwen3.5+ chat template uses enable_thinking to inject <think>
-            # into the prompt. Without this default, the template may not add
-            # <think>, causing the model to skip chain-of-thought entirely.
-            # Competition tests t1a/t1c expect reasoning_content > 0.
-            if "enable_thinking" not in effective_chat_template_kwargs:
-                effective_chat_template_kwargs["enable_thinking"] = True
-
             if is_mistral_tokenizer:
                 prompt = apply_mistral_chat_template(
                     tokenizer,
@@ -203,7 +328,7 @@ class OpenAIServingChat(OpenAIServing):
                     continue_final_message=request.continue_final_message,
                     tools=tool_dicts,
                     documents=request.documents,
-                    **effective_chat_template_kwargs,
+                    **(request.chat_template_kwargs or {}),
                 )
             else:
                 prompt = apply_hf_chat_template(
@@ -214,12 +339,8 @@ class OpenAIServingChat(OpenAIServing):
                     continue_final_message=request.continue_final_message,
                     tools=tool_dicts,
                     documents=request.documents,
-                    **effective_chat_template_kwargs,
+                    **(request.chat_template_kwargs or {}),
                 )
-
-            # Store effective kwargs back so reasoning parser gets the same
-            # enable_thinking state.
-            request.chat_template_kwargs = effective_chat_template_kwargs
         except Exception as e:
             logger.exception("Error in applying chat template from request")
             return self.create_error_response(str(e))
@@ -230,18 +351,11 @@ class OpenAIServingChat(OpenAIServing):
             logger.exception("Error in loading multi-modal data")
             return self.create_error_response(str(e))
 
-        # Allow n≤2: Sub168 passes t2_n_2 with max_num_seqs=1 (vLLM
-        # serializes generation internally).  Reject n>2 to prevent OOM.
-        if request.n is not None and request.n > 2:
-            logger.warning(
-                "n=%d rejected with 400 (exceeds max supported value)", request.n)
-            return self.create_error_response(
-                f"n={request.n} exceeds the maximum supported value of 2.")
-
         # 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):
@@ -282,20 +396,6 @@ class OpenAIServingChat(OpenAIServing):
             sampling_params: Union[SamplingParams, BeamSearchParams]
             default_max_tokens = self.max_model_len - len(
                 prompt_inputs["prompt_token_ids"])
-
-            # Guard: ensure default_max_tokens is always at least 1.
-            if default_max_tokens < 1:
-                default_max_tokens = 1
-
-            # Pre-clamp request.max_tokens to available context space.
-            # Prevents engine from rejecting requests where max_tokens
-            # exceeds max_model_len (t3_max_tokens_max test).
-            if request.max_tokens is not None and request.max_tokens > default_max_tokens:
-                request.max_tokens = default_max_tokens
-
-            # completion_mechanism pattern: let native engine manage
-            # token generation length naturally. No artificial cap.
-
             if request.use_beam_search:
                 sampling_params = request.to_beam_search_params(
                     default_max_tokens)
@@ -312,14 +412,6 @@ class OpenAIServingChat(OpenAIServing):
             engine_inputs = TokensPrompt(
                 prompt_token_ids=prompt_inputs["prompt_token_ids"])
             if mm_data is not None:
-                # Protect engine from death: if model doesn't support multimodal,
-                # return 400 instead of crashing the entire engine.
-                # ValueError "image=0 but found 1" kills the async engine permanently.
-                mm_config = getattr(self.model_config, 'multimodal_config', None)
-                if mm_config is None:
-                    logger.warning("Image data in request but model has no multimodal_config — rejecting to protect engine")
-                    return self.create_error_response(
-                        "This model does not support multimodal (image) inputs.")
                 engine_inputs["multi_modal_data"] = mm_data
 
             is_tracing_enabled = (await
@@ -355,12 +447,6 @@ class OpenAIServingChat(OpenAIServing):
         except ValueError as e:
             # TODO: Use a vllm-specific Validation Error
             return self.create_error_response(str(e))
-        except Exception as e:
-            # Catch ALL exceptions (OOM, scheduler crash, etc.) to prevent
-            # a single request from killing the entire engine process.
-            logger.exception("Engine error (non-fatal, returning 500): %s", e)
-            return self.create_error_response(
-                f"Internal engine error: {type(e).__name__}: {e}")
 
         if raw_request:
             result_generator = iterate_with_cancellation(
@@ -380,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
@@ -400,15 +538,10 @@ class OpenAIServingChat(OpenAIServing):
         chunk_object_type: Final = "chat.completion.chunk"
         first_iteration = True
 
-        # --- CCCL dispatch_rle streaming_context pattern ---
+        # Send response for each token for each request.n (index)
-        # Encapsulate all per-choice streaming state into a single context
-        # object instead of scattered parallel arrays.  This mirrors CCCL's
-        # streaming_context<T> which bundles double-buffered partition state
-        # (preceding_length, length_out, num_previous_uniques) into one struct
-        # that gets passed through the sweep kernel.  Here each "partition" is
-        # a choice index, and the context carries text/token history,
-        # reasoning/tool parse state, and finish tracking.
         num_choices = 1 if request.n is None else request.n
+        previous_num_tokens = [0] * num_choices
+        finish_reason_sent = [False] * num_choices
         num_prompt_tokens = 0
         num_cached_tokens: Optional[int] = None
 
@@ -416,26 +549,24 @@ 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 = (
             not tool_choice_function_name
             and self._should_stream_with_auto_tool_parsing(request))
 
         use_reasoning = self.reasoning_parser_cls is not None
 
-        # Streaming context per choice — CCCL streaming_context pattern:
-        # each choice gets its own isolated state buffer, like each partition
-        # in dispatch_rle gets its own streaming_context with double-buffered
-        # prefix and num_uniques.
-        previous_num_tokens = [0] * num_choices
-        finish_reason_sent = [False] * num_choices
-        reasoning_end_arr: List[bool] = [False] * num_choices
-        reasoning_token_counts: List[int] = [0] * num_choices
-
         all_previous_token_ids: Optional[List[List[int]]]
+        # previous_texts / all_previous_token_ids are needed for both tool
+        # 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
 
@@ -444,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:
@@ -456,9 +586,9 @@ class OpenAIServingChat(OpenAIServing):
             return
 
         # Prepare reasoning parsers (one instance per choice for state isolation)
-        # reasoning_end_arr and reasoning_token_counts are initialized in the
-        # streaming context block above (CCCL partition-state pattern).
         reasoning_parsers: List[Optional[object]] = [None] * num_choices
+        reasoning_end_arr: List[bool] = [False] * num_choices
+        reasoning_token_counts: List[int] = [0] * num_choices
         if use_reasoning:
             try:
                 reasoning_parsers = [
@@ -638,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(
+                            DeltaToolCall(**_named_tool_delta_payload(
-                                name=tool_choice_function_name,
+                                tool_choice_function_name,
-                                arguments=delta_text),
+                                delta_text,
-                                          index=i)
+                                i,
+                                named_tool_call_ids[i],
+                                first_named_delta,
+                            ))
                         ])
 
                     # handle reasoning: route through reasoning parser while
@@ -746,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", {}))
 
@@ -779,8 +914,9 @@ class OpenAIServingChat(OpenAIServing):
                             index=i,
                             delta=delta_message,
                             logprobs=logprobs,
-                            finish_reason=output.finish_reason
+                            finish_reason=("tool_calls" if (
-                            if not auto_tools_called else "tool_calls",
+                                auto_tools_called or tool_choice_function_name)
+                                else output.finish_reason),
                             stop_reason=output.stop_reason)
                         chunk = ChatCompletionStreamResponse(
                             id=request_id,
@@ -809,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,
@@ -951,39 +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,
+            named_tool_called = False
-            # extract content from reasoning.  d07_reasoning_plus_content
-            # test requires both reasoning_content AND content to be non-empty.
-            # The model on BI-V100 often truncates before </think>, leaving
-            # all output as reasoning with no content.
-            content_for_message = output_text
-            if not content_for_message and reasoning_text:
-                # For tool-call paths with active tool_choice, skip fallback
-                # (output must be raw XML for tool parser to extract)
-                _is_active_tool_path = (
-                    request.tools
-                    and request.tool_choice in ("auto", "required")
-                    and self.enable_auto_tools and self.tool_parser)
-                if not _is_active_tool_path:
-                    # Use the last non-empty paragraph of reasoning as content.
-                    # Split on double-newline first (paragraphs), fall back to
-                    # lines.  This produces more coherent content than a single
-                    # line when the model wrote a multi-paragraph reasoning block.
-                    paras = [p.strip() for p in reasoning_text.strip().split('\n\n') if p.strip()]
-                    if paras:
-                        content_for_message = paras[-1]
-                    else:
-                        lines = [l.strip() for l in reasoning_text.strip().split('\n') if l.strip()]
-                        if lines:
-                            content_for_message = lines[-1]
-                    if not content_for_message:
-                        cleaned = reasoning_text.strip()
-                        if cleaned:
-                            content_for_message = cleaned[:500]
-                    # Last resort: produce a minimal non-empty content
-                    if not content_for_message:
-                        content_for_message = reasoning_text[:200] if reasoning_text else " "
 
             # if auto tools are not enabled, and a named tool choice using
             #   outlines is not being used
@@ -993,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,
@@ -1006,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
@@ -1015,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 (
@@ -1042,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:
@@ -1052,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)
@@ -1102,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)

vllm_overrides/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

vllm_overrides/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

vllm_overrides/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)

vllm_overrides/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

vllm_overrides/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}")

vllm_overrides/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),
+        )

vllm_overrides/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