wudixzy stack docker build test: 2615-line qwen3_5.py + 12 prebuilt .so + 251-line patch_ops.sh

Dockerfile

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

computility-run.yaml

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

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

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

qwen3_6_scripts/cli_args.py

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

qwen3_6_scripts/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/mamba_cache.py

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

qwen3_6_scripts/patch_block_major_cache_engine.py

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

qwen3_6_scripts/patch_block_major_worker_capacity.py

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

qwen3_6_scripts/patch_block_manager_cache_trace.py

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

qwen3_6_scripts/patch_corex_swap_blocks.py

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

qwen3_6_scripts/patch_executor_startup_debug.py

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

qwen3_6_scripts/patch_model_runner.py

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

qwen3_6_scripts/patch_ops.sh

@@ -0,0 +1,251 @@
+#!/usr/bin/env bash
+# BI-V100 patch script for Qwen3.6-35B-A3B (Qwen3_5 MoE architecture)
+#
+# Triton situation on BI-V100:
+#   - Standard Triton 2.3.1 is already present in the image.
+#   - HAS_TRITON = False (hardcoded in vendor vllm), but Triton is still used
+#     for TP-mode cache management (custom_cache_manager / libentry).
+#   - The vendor's triton_utils/__init__.py, custom_cache_manager.py, libentry.py
+#     are already correct for standard Triton 2.3.1 — do NOT overwrite them.
+#   - DO NOT install BI-V150 corex Triton 2.1.0 (pkgs/triton): that causes
+#     GPU hang on BI-V100 because the Triton CUDA PTX kernels are incompatible.
+
+# Recommended server start command for TP=4 support 256K, needs chunked prefill
+# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
+#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
+#     --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
+#     --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
+#     --tool-call-parser qwen3_coder --reasoning-parser qwen3
+#
+# With prefix caching (GDN align-mode, requires chunked prefill):
+# CUDA_VISIBLE_DEVICES="4,5,6,7" VLLM_ENGINE_ITERATION_TIMEOUT_S=3600 python3 -m vllm.entrypoints.openai.api_server \
+#     --model /workspace/models/Qwen3.6-35B-A3B --port 1111 --served-model-name llm \
+#     --max-model-len 262144 --trust-remote-code -tp 4 --gpu-memory-utilization 0.90 \
+#     --max-num-seqs 1 --disable-log-requests --disable-frontend-multiprocessing \
+#     --max-num-batched-tokens 8192 --enable-chunked-prefill --enable-prefix-caching \
+#     --max-seq-len-to-capture 32768 --enable-auto-tool-choice \
+#     --tool-call-parser qwen3_coder --reasoning-parser qwen3
+
+set -eo pipefail
+
+# cd into this script's directory so ./relative paths work
+cd "$(dirname "${BASH_SOURCE[0]}")"
+echo "[patch_ops] working directory: $(pwd)"
+
+build_stage() { printf '[BI100 BUILD] %s\n' "$1" >&2; }
+require_file() {
+    local path=$1
+    [[ -f "$path" ]] || {
+        printf 'required patch source is missing: %s\n' "$path" >&2
+        exit 2
+    }
+}
+install_patch_file() {
+    local source=$1
+    local target=$2
+
+    require_file "$source"
+    mkdir -p "$(dirname "$target")"
+    install -m 0644 "$source" "$target"
+}
+
+build_stage "patch script entered"
+
+build_stage "checking offline transformers dependency"
+# --- transformers: Qwen3_5 tokenizer / model files --------------------------
+TRANSFORMERS_REQUIRED_VERSION="4.55.3"
+if ! python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
+import importlib.metadata
+import sys
+
+required = sys.argv[1]
+try:
+    installed = importlib.metadata.version("transformers")
+except importlib.metadata.PackageNotFoundError:
+    raise SystemExit(1)
+raise SystemExit(0 if installed == required else 1)
+PY
+then
+  WHEEL_DIR="./wheels"
+  if ! ls "${WHEEL_DIR}/transformers-${TRANSFORMERS_REQUIRED_VERSION}"*.whl >/dev/null 2>&1; then
+    echo "transformers ${TRANSFORMERS_REQUIRED_VERSION} is required, but no offline wheel was found in ${WHEEL_DIR}" >&2
+    exit 2
+  fi
+  python3 -m pip install --no-index --no-deps --find-links="${WHEEL_DIR}" \
+    "transformers==${TRANSFORMERS_REQUIRED_VERSION}"
+fi
+
+python3 - "$TRANSFORMERS_REQUIRED_VERSION" <<'PY'
+import importlib.metadata
+import sys
+
+required = sys.argv[1]
+installed = importlib.metadata.version("transformers")
+if installed != required:
+    raise SystemExit(
+        f"transformers version mismatch: expected {required}, got {installed}")
+print(f"[ok] transformers {installed}")
+PY
+
+build_stage "discovering Python package roots"
+python3 - <<'PY' > /tmp/qwen36_patch_paths.env
+from patch_utils import package_root, shell_env_line
+
+print(shell_env_line("VLLM_ROOT", package_root("vllm")))
+print(shell_env_line("TRANSFORMERS_ROOT", package_root("transformers")))
+PY
+source /tmp/qwen36_patch_paths.env
+
+echo "VLLM_ROOT=${VLLM_ROOT}"
+echo "TRANSFORMERS_ROOT=${TRANSFORMERS_ROOT}"
+[[ -d "$VLLM_ROOT" ]] || {
+    printf 'vLLM root does not exist: %s\n' "$VLLM_ROOT" >&2
+    exit 2
+}
+
+VLLM_OVERRIDE_ROOT="./vendor_overrides/vllm"
+[[ -d "$VLLM_OVERRIDE_ROOT" ]] || {
+    printf 'vLLM override directory missing: %s\n' "$VLLM_OVERRIDE_ROOT" >&2
+    exit 2
+}
+
+build_stage "installing authoritative vLLM core block overrides"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/evictor_v2.py" \
+    "${VLLM_ROOT}/core/evictor_v2.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_kv_content_cache.py" \
+    "${VLLM_ROOT}/core/block/cpu_kv_content_cache.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/cpu_gpu_block_allocator.py" \
+    "${VLLM_ROOT}/core/block/cpu_gpu_block_allocator.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/prefix_caching_block.py" \
+    "${VLLM_ROOT}/core/block/prefix_caching_block.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block/block_table.py" \
+    "${VLLM_ROOT}/core/block/block_table.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/core/block_manager_v2.py" \
+    "${VLLM_ROOT}/core/block_manager_v2.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/sampling_params.py" \
+    "${VLLM_ROOT}/sampling_params.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/model_executor/sampling_metadata.py" \
+    "${VLLM_ROOT}/model_executor/sampling_metadata.py"
+install_patch_file \
+    "${VLLM_OVERRIDE_ROOT}/model_executor/layers/sampler.py" \
+    "${VLLM_ROOT}/model_executor/layers/sampler.py"
+
+build_stage "installing hash-pinned CoreX 3.2.3 extensions"
+bash ./install_prebuilt_corex.sh "${VLLM_ROOT}"
+
+build_stage "installing BI100 runtime modules"
+cp ./bi100_env.py "${VLLM_ROOT}/bi100_env.py"
+cp ./bi100_profile.py "${VLLM_ROOT}/bi100_profile.py"
+cp ./block_major_kv_cache.py "${VLLM_ROOT}/block_major_kv_cache.py"
+cp ./gdn_prefix.py "${VLLM_ROOT}/gdn_prefix.py"
+
+build_stage "installing CoreX paged-KV swap compatibility"
+python3 ./patch_corex_swap_blocks.py
+python3 ./patch_block_major_cache_engine.py
+python3 ./patch_worker_cache_transfer_order.py
+
+# --- paged_attn.py: replace forward_prefix with pure-PyTorch fallback -------
+# The Triton context_attention_fwd kernel hangs BI-V100 GPUs permanently
+# (standard Triton 2.3.1 PTX is not supported by the corex runtime either).
+# Our paged_attn.py bypasses it entirely via _forward_prefix_pytorch, which
+# utilizes K-tiling techniques, and also have _forward_decode_pytorch to bypass kernel
+# when context length is high
+cp ./paged_attn.py "${VLLM_ROOT}/attention/ops/paged_attn.py"
+
+# --- model_runner.py: fix prefix_cache_hit stays True in chunked-prefill chunk 2+ ---
+# Bug: _compute_for_prefix_cache_hit Case 1 (prefix_cache_len <= context_len)
+# leaves prefix_cache_hit=True. Then _add_seq_group uses block_table=computed_block_nums
+# (only the original prefix blocks), ignoring chunk-1 KV cache blocks.
+# _forward_prefix_pytorch then gets an undersized block_tables and crashes with
+# "amax(): Expected reduction dim -1 to have non-zero size" on the 2nd tile.
+# Fix: set prefix_cache_hit=False for Case 1 so the full block_tables is used.
+python3 ./patch_model_runner.py
+
+build_stage "installing executor startup diagnostics"
+python3 ./patch_executor_startup_debug.py
+python3 ./patch_worker_startup_profile_guard.py
+python3 ./patch_block_major_worker_capacity.py
+
+build_stage "installing transformers Qwen3.5 model support"
+cp -r ./qwen3_5 "${TRANSFORMERS_ROOT}/models/"
+cp -r ./qwen3_5_moe "${TRANSFORMERS_ROOT}/models/"
+python3 ./patch_transformers_qwen3_5.py
+
+build_stage "installing vLLM Qwen3.6 model implementation"
+# --- vllm model: Qwen3.6-35B-A3B (Qwen3_5 MoE arch) -------------------------
+cp ./mamba_cache.py "${VLLM_ROOT}/model_executor/models/"
+cp ./qwen3_5.py "${VLLM_ROOT}/model_executor/models/qwen3_5.py"
+python3 ./patch_vllm_qwen3_5.py
+
+# --- 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

@@ -0,0 +1,100 @@
+"""
+Patches transformers 4.55.3 to register qwen3_5 and qwen3_5_moe model types.
+
+Deploy steps on the remote machine:
+  1. patch_ops.sh locates transformers with importlib.util.find_spec.
+  2. cp -r modified_scripts/qwen3_5* into the detected transformers/models.
+  3. python3 modified_scripts/patch_transformers_qwen3_5.py
+"""
+
+import sys
+
+from patch_utils import package_root, replace_once, replace_one_of
+
+TRANSFORMERS_ROOT = package_root("transformers")
+AUTO_CONFIG = TRANSFORMERS_ROOT / "models" / "auto" / "configuration_auto.py"
+MODELS_INIT = TRANSFORMERS_ROOT / "models" / "__init__.py"
+
+
+def main():
+    print(f"=== Patching {AUTO_CONFIG} ===")
+    replace_one_of(AUTO_CONFIG, [
+        # CONFIG_MAPPING_NAMES: insert qwen3_5 + qwen3_5_moe right after qwen3
+        (
+            '("qwen3", "Qwen3Config"),',
+            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n        ("qwen3_5_moe", "Qwen3_5MoeConfig"),',
+        ),
+        (
+            '("qwen3", "Qwen3Config")\n',
+            '("qwen3", "Qwen3Config"),\n        ("qwen3_5", "Qwen3_5Config"),\n        ("qwen3_5_moe", "Qwen3_5MoeConfig"),\n',
+        ),
+    ], required=True, already_contains='("qwen3_5_moe", "Qwen3_5MoeConfig")')
+    replace_one_of(AUTO_CONFIG, [
+        # MODEL_NAMES_MAPPING (model_type -> human readable name)
+        (
+            '("qwen3", "Qwen3"),',
+            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n        ("qwen3_5_moe", "Qwen3_5_MoE"),',
+        ),
+        (
+            '("qwen3", "Qwen3")\n',
+            '("qwen3", "Qwen3"),\n        ("qwen3_5", "Qwen3_5"),\n        ("qwen3_5_moe", "Qwen3_5_MoE"),\n',
+        ),
+    ], required=True, already_contains='("qwen3_5_moe", "Qwen3_5_MoE")')
+
+    print(f"\n=== Patching {MODELS_INIT} ===")
+    replace_once(
+        MODELS_INIT,
+        "from .qwen3 import *\n",
+        "from .qwen3 import *\n    from .qwen3_5 import *\n    from .qwen3_5_moe import *\n",
+        required=True,
+        already_contains="from .qwen3_5_moe import *")
+
+    # Verification
+    print("\n=== Verification ===")
+    try:
+        import importlib.util, types
+
+        def _load_config_mod(module_name, file_path):
+            spec = importlib.util.spec_from_file_location(module_name, file_path)
+            mod = importlib.util.module_from_spec(spec)
+            mod.__package__ = ".".join(module_name.split(".")[:-1])
+            pkg = sys.modules.setdefault("transformers", types.ModuleType("transformers"))
+            pkg.__path__ = [str(TRANSFORMERS_ROOT)]
+            cu = sys.modules.setdefault(
+                "transformers.configuration_utils", types.ModuleType("transformers.configuration_utils"))
+            class _PC:
+                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__ = [str(TRANSFORMERS_ROOT)]
+            spec.loader.exec_module(mod)
+            return mod
+
+        mod27 = _load_config_mod(
+            "transformers.models.qwen3_5.configuration_qwen3_5",
+            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",
+            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})")
+        t = moe_cfg.text_config
+        print(f"    num_experts={t.num_experts}, top_k={t.num_experts_per_tok}, "
+              f"shared={t.shared_expert_intermediate_size}, layers={t.num_hidden_layers}")
+    except Exception as e:
+        print(f"  [optional] smoke-test failed (may be fine at runtime): {e}")
+
+    print("\nDone.")
+
+
+if __name__ == "__main__":
+    main()

qwen3_6_scripts/patch_utils.py

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

qwen3_6_scripts/patch_vllm_qwen3_5.py

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

qwen3_6_scripts/patch_vllm_tool_parser.py

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

qwen3_6_scripts/patch_worker_cache_transfer_order.py

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

qwen3_6_scripts/patch_worker_profile_override.py

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

qwen3_6_scripts/patch_worker_startup_profile_guard.py

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

qwen3_6_scripts/patch_xformers_profile.py

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

qwen3_6_scripts/patch_xformers_sdpa_batch.py

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

qwen3_6_scripts/patch_xformers_sdpa_batch_kernel.py

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

qwen3_6_scripts/patch_xformers_sdpa_seq.py

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

qwen3_6_scripts/patch_xformers_sdpa_seq_kernel.py

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

qwen3_6_scripts/prebuilt/corex-3.2.3-ivcore10/SHA256SUMS

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

qwen3_6_scripts/qwen3_5/__init__.py

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

qwen3_6_scripts/qwen3_5/configuration_qwen3_5.py

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

qwen3_6_scripts/qwen3_5_moe/__init__.py

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

qwen3_6_scripts/qwen3_5_moe/configuration_qwen3_5_moe.py

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

qwen3_6_scripts/qwen3coder_tool_parser.py

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

qwen3_6_scripts/reasoning/__init__.py

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

qwen3_6_scripts/reasoning/abs_reasoning_parsers.py

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

qwen3_6_scripts/reasoning/qwen3_reasoning_parser.py

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

qwen3_6_scripts/serving_tokenization.py

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

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

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

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

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

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

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

qwen3_6_scripts/vendor_overrides/vllm/core/block_manager_v2.py

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

qwen3_6_scripts/vendor_overrides/vllm/core/evictor_v2.py

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

qwen3_6_scripts/vendor_overrides/vllm/model_executor/sampling_metadata.py

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

qwen3_6_scripts/vendor_overrides/vllm/sampling_params.py

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