diff --git a/CMakeLists.txt b/CMakeLists.txt index ca3f6c48e3..197f0702b9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -39,6 +39,28 @@ if(CMAKE_CUDA_COMPILER_VERSION VERSION_LESS 13.1) "${CMAKE_CUDA_COMPILER_VERSION}") endif() +# CUDA 13's CCCL headers require MSVC's standard-conforming preprocessor. +if(MSVC) + add_compile_options($<$:-Xcompiler=/Zc:preprocessor>) + # windows.h defines max/min as macros that break std::max/std::min. + add_compile_definitions(NOMINMAX) +endif() + +# Media (vision) decode and acquisition need FFMPEG and libcurl. Those are not +# part of the default Windows toolchain, so the text-only build +# (NINFER_BUILD_MEDIA=OFF) compiles API-compatible stubs instead and rejects +# vision requests at runtime. +if(WIN32) + set(NINFER_MEDIA_DEFAULT OFF) +else() + set(NINFER_MEDIA_DEFAULT ON) +endif() +option(NINFER_BUILD_MEDIA "Build FFMPEG/libcurl media decode and acquisition" + ${NINFER_MEDIA_DEFAULT}) + +# ninfer_serve and prompt_input link ninfer_media_acquire unconditionally, so +# the target must exist whenever apps/tests build; NINFER_BUILD_MEDIA only +# selects the real (libcurl) vs stub implementation. set(NINFER_BUILD_MEDIA_ACQUIRE OFF) if(NINFER_BUILD_APPS OR BUILD_TESTING) set(NINFER_BUILD_MEDIA_ACQUIRE ON) @@ -55,11 +77,50 @@ if(NINFER_BUILD_APPS OR BUILD_TESTING) endif() find_package(CUDAToolkit REQUIRED) -find_package(PkgConfig REQUIRED) -pkg_check_modules(FFMPEG REQUIRED IMPORTED_TARGET - libavformat>=60 libavcodec>=60 libavutil>=58 libswscale>=7) -if(NINFER_BUILD_MEDIA_ACQUIRE) - pkg_check_modules(LIBCURL REQUIRED IMPORTED_TARGET libcurl>=7.85) +if(NINFER_BUILD_MEDIA) + if(WIN32) + # Windows: use find_path/find_library for FFMPEG + libcurl (no pkg-config) + set(FFMPEG_ROOT "${PROJECT_SOURCE_DIR}/../third_party/ffmpeg/ffmpeg-master-latest-win64-gpl-shared") + set(CURL_ROOT "${PROJECT_SOURCE_DIR}/../third_party/curl-inst") + + find_path(FFMPEG_INCLUDE_DIR NAMES libavformat/avformat.h PATHS "${FFMPEG_ROOT}/include" NO_DEFAULT_PATH) + find_library(AVFORMAT_LIBRARY NAMES avformat.lib PATHS "${FFMPEG_ROOT}/lib" NO_DEFAULT_PATH) + find_library(AVCODEC_LIBRARY NAMES avcodec.lib PATHS "${FFMPEG_ROOT}/lib" NO_DEFAULT_PATH) + find_library(AVUTIL_LIBRARY NAMES avutil.lib PATHS "${FFMPEG_ROOT}/lib" NO_DEFAULT_PATH) + find_library(SWSCALE_LIBRARY NAMES swscale.lib PATHS "${FFMPEG_ROOT}/lib" NO_DEFAULT_PATH) + + if(NOT FFMPEG_INCLUDE_DIR OR NOT AVFORMAT_LIBRARY OR NOT AVCODEC_LIBRARY OR NOT AVUTIL_LIBRARY OR NOT SWSCALE_LIBRARY) + message(FATAL_ERROR "FFMPEG libraries not found. Install BtbN ffmpeg-win64-gpl-shared into third_party/ffmpeg/") + endif() + + add_library(PkgConfig::FFMPEG INTERFACE IMPORTED) + set_target_properties(PkgConfig::FFMPEG PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${FFMPEG_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "${AVFORMAT_LIBRARY};${AVCODEC_LIBRARY};${AVUTIL_LIBRARY};${SWSCALE_LIBRARY}" + ) + + if(NINFER_BUILD_MEDIA_ACQUIRE) + find_path(CURL_INCLUDE_DIR NAMES curl/curl.h PATHS "${CURL_ROOT}/include" NO_DEFAULT_PATH) + find_library(CURL_LIBRARY NAMES libcurl_imp.lib PATHS "${CURL_ROOT}/lib" NO_DEFAULT_PATH) + + if(NOT CURL_INCLUDE_DIR OR NOT CURL_LIBRARY) + message(FATAL_ERROR "libcurl not found. Build curl with MSVC into third_party/curl-inst/") + endif() + + add_library(PkgConfig::LIBCURL INTERFACE IMPORTED) + set_target_properties(PkgConfig::LIBCURL PROPERTIES + INTERFACE_INCLUDE_DIRECTORIES "${CURL_INCLUDE_DIR}" + INTERFACE_LINK_LIBRARIES "${CURL_LIBRARY}" + ) + endif() + else() + find_package(PkgConfig REQUIRED) + pkg_check_modules(FFMPEG REQUIRED IMPORTED_TARGET + libavformat>=60 libavcodec>=60 libavutil>=58 libswscale>=7) + if(NINFER_BUILD_MEDIA_ACQUIRE) + pkg_check_modules(LIBCURL REQUIRED IMPORTED_TARGET libcurl>=7.85) + endif() + endif() endif() find_package(Threads REQUIRED) diff --git a/Dockerfile b/Dockerfile index bc3dc5db9a..5f269a13cb 100644 --- a/Dockerfile +++ b/Dockerfile @@ -13,6 +13,7 @@ RUN apt-get update \ libswscale-dev \ ninja-build \ pkg-config \ + python3 \ && rm -rf /var/lib/apt/lists/* WORKDIR /src @@ -31,6 +32,7 @@ ARG DEBIAN_FRONTEND=noninteractive RUN apt-get update \ && apt-get install --yes --no-install-recommends \ ca-certificates \ + curl \ libavcodec60 \ libavformat60 \ libavutil58 \ @@ -45,4 +47,19 @@ WORKDIR /workspace EXPOSE 8080 STOPSIGNAL SIGTERM +# /health is the only unauthenticated endpoint, so this needs no API key. It answers +# 503 once the inference executor has failed (issue #10) and 200 otherwise, which is +# what makes an alive-but-unusable server visible from outside the process. +# +# NINFER_PORT must match --port; HEALTHCHECK is static but the port is not. The +# server binds before loading the model and only listens afterwards, so expect +# connection-refused for the whole load — that is what start-period covers. +# +# NOTE: a restart policy does NOT act on health. Docker restarts on EXIT, not on +# unhealthy, so this makes the failed state observable and alertable; it does not +# by itself recover the container. Recovery needs a watcher (Swarm, an external +# monitor, or the native supervisor's health-restart path). +HEALTHCHECK --interval=30s --timeout=5s --start-period=180s --retries=3 \ + CMD curl -fsS "http://127.0.0.1:${NINFER_PORT:-8080}/health" || exit 1 + CMD ["ninfer-serve", "--help"] diff --git a/RUNBOOK.md b/RUNBOOK.md new file mode 100644 index 0000000000..92c080b4db --- /dev/null +++ b/RUNBOOK.md @@ -0,0 +1,645 @@ +# GPU runbook: vLLM-dialect thinking off + +Schema and request-parse tests on this branch cover the dialect without loading a +model. This runbook is the remaining full-size check: greedy byte-identity across +the three thinking-off spellings, plus the 400 paths that a live HTTP server must +echo. + +Do not run this until the coordinator schedules a GPU window. The registered +Qwen3.8-27B NVFP4 artifact is about 20 GiB of weights and cannot fit the 16 GiB +compact allocation cap. + +## Constraints + +- Do not stop, pause, or restart docker containers. +- Before any process that allocates GPU memory: + 1. `nvidia-smi` must show at least 20 GiB free. + 2. Acquire `C:\Users\igorl\.ninfer-gpu.lock` with `mkdir` (atomic). If it + exists, wait 60 s and retry for up to 30 minutes, then stop. + 3. Remove the lock directory immediately after, success or failure. +- Rebuild of `ninfer:seedstore` and the `:8018` lane is coordinator-owned. + +## Server + +Full-size serving flags (no `--preserve-thinking`): + +```text +--max-context 131072 --kv-capacity 1048576 --max-concurrency 8 --spec mtp +--draft-tokens 5 --lm-head-draft --kv-dtype int8 --prefill-chunk 2048 --vision +--cors --prefix-cache-mib 4096 +``` + +Public model ID: `qwen3.8-27b`. Base URL in this runbook: `http://127.0.0.1:8018`. + +Use temperature 0, seed 0, and a cold prefix (new prompt text, or restart so the +prefix cache does not hide a template mismatch). + +## Prompt set + +Three one-shot Chat Completions bodies. Only the thinking-off spelling changes. + +Shared fields: +| Name | Extra fields | +|---|---| +| `effort_none` | `"reasoning_effort": "none"` | +| `kwargs_off` | `"chat_template_kwargs": {"enable_thinking": false}` | +| `top_off` | `"enable_thinking": false` | + +Also send one Responses request with `"input": "Reply with the single word ping."`, +`"max_output_tokens": 32`, `"temperature": 0`, and +`"chat_template_kwargs": {"enable_thinking": false}`. + +## Expected generation + +For each thinking-off spelling: + +- HTTP 200. +- `choices[0].message.content` is byte-identical across `effort_none`, + `kwargs_off`, and `top_off` on a cold prompt. +- `choices[0].message.reasoning_content` is absent or empty. +- `usage.prompt_tokens` is identical across the three Chat Completions spellings + (the chat template must have taken the same thinking-off branch). + +A live server with thinking left on (omit all three spellings, default) must +differ: `reasoning_content` is non-empty or `prompt_tokens` is larger. + +## Expected 400s + +`chat_template_kwargs.enable_thinkng: false` (misspelling): + +```json +{"error":{"message":"chat_template_kwargs.enable_thinkng is not supported","type":"invalid_request_error","param":"chat_template_kwargs","code":"chat_template_option_not_supported"}} +``` + +`enable_thinking: true` together with `chat_template_kwargs.enable_thinking: false`: + +```json +{"error":{"message":"conflicting enable_thinking values","type":"invalid_request_error","param":"enable_thinking","code":"conflicting_template_option"}} +``` + +`reasoning_effort: "high"` against the registered template: + +HTTP 400 `reasoning_effort_not_supported`. Do not treat this as a dialect +failure; `high` is a parsed protocol value and is not an alias of `xhigh`. + +## Commands + +```bash +BASE=http://127.0.0.1:8018 +MODEL=qwen3.8-27b + +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Reply with the single word ping.\"}], + \"max_completion_tokens\": 32, + \"temperature\": 0, + \"seed\": 0, + \"reasoning_effort\": \"none\" +}" + +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Reply with the single word ping.\"}], + \"max_completion_tokens\": 32, + \"temperature\": 0, + \"seed\": 0, + \"chat_template_kwargs\": {\"enable_thinking\": false} +}" + +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Reply with the single word ping.\"}], + \"max_completion_tokens\": 32, + \"temperature\": 0, + \"seed\": 0, + \"enable_thinking\": false +}" + +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Reply with the single word ping.\"}], + \"chat_template_kwargs\": {\"enable_thinkng\": false} +}" + +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [{\"role\": \"user\", \"content\": \"Reply with the single word ping.\"}], + \"enable_thinking\": true, + \"chat_template_kwargs\": {\"enable_thinking\": false} +}" +``` + +Compare the three 200 bodies with `jq -S '.choices[0].message'` (or equivalent). +Pass only when content, reasoning_content, and prompt_tokens match. + +--- + +# GPU runbook: native Windows vs WSL2-container tax + +Native `ninfer-serve.exe` now compiles on this box (MSVC 19.44.35228 + CUDA 13.3.33, +`sm_120a`, `NINFER_BUILD_MEDIA=OFF`). Do not boot it or acquire the GPU lock until +the coordinator schedules an exclusive window. The Qwen3.8-27B NVFP4 artifact is +about 20 GiB and does not fit the 16 GiB compact cap. + +This runbook measures the WSL2 tax. Decode is expected to be similar (GPU-resident, +bandwidth-bound). Prefill, TTFT, weight load, boot wall time, and seed-store +captures all cross the WSL2 boundary on the container arm. + +## Arms + +| Arm | Runtime | Binary | +|---|---|---| +| A container | `ninfer:seedstore` under WSL2/docker | container `ninfer-serve` from `feat/prefix-seed-store` @ 352a49c3 plus this Windows port | +| B native | `build-win/apps/ninfer-serve.exe` from `task/issue-6-native-windows` | same git tree, MSVC+CUDA 13.3, text-only (`NINFER_BUILD_MEDIA=OFF`) | + +Same checkpoint: `neroued/Qwen3.8-27B-nvfp4-NInfer` / public model id `qwen3.8-27b`. +Same serving flags both arms. Native currently **cannot** honor `--vision` until +FFmpeg+libcurl are installed. Until then, drop `--vision` on **both** arms so the +A/B stays matched, or install the media prefix and rebuild native with +`-DNINFER_BUILD_MEDIA=ON` before the window. + +Never stop production containers (`sglang-qwen38` on `:8016`, embeddings, whisper). + +## Server flags (matched) + +No `--preserve-thinking`. JSONL log required. Port 8018 native or container, one +at a time. + +```text +--host 127.0.0.1 --port 8018 +--max-context 131072 --kv-capacity 1048576 --max-concurrency 8 +--spec mtp --draft-tokens 5 --lm-head-draft +--kv-dtype int8 --prefill-chunk 2048 --cors +--prefix-cache-mib 4096 +--request-log-jsonl -ninfer.jsonl +``` + +Add `--vision` only when both arms actually load Vision. + +Native launch (after vcvars64 + CUDA 13.3 on PATH): + +```bat +build-win\apps\ninfer-serve.exe --host 127.0.0.1 --port 8018 ... +``` + +Record wall time from process start to first `GET /health` 200. That is boot +wall time. Weight-load time is the `load_progress` / startup log span until the +server accepts connections. + +## Metrics + +Collect on every request from `request_done` JSONL: + +- prefill tok/s = `computed_prefill_tokens / timings_seconds.prefill` +- decode tok/s = `(completion_tokens - 1) / timings_seconds.decode` +- TTFT = `timings_seconds.ttft` +- `prefix_reuse_path`, `prefix_cache_hit_tokens` + +### Prefill (~6k and ~57k) + +Two prompt lengths, serial, c=1, `--greedy` (prefill is not MTP-luck bound): + +- ~6k: a real ~6k-token chat from the serving corpus or long-niah fixture. +- ~57k: a long-context body in the same family. Report `prompt_tokens` from + `request_done` so the two buckets are actual, not nominal. + +Three repetitions each length per arm. Report median prefill tok/s. + +### Decode c=1, >= 3 boots + +Three full process boots per arm, interleaved A/B/A/B/A/B, c=1 n=16, +`examples/cli/messages/scenario_*.json` cycled to 16, MTP-5, **not** greedy. +Arm score = median of 3 boot medians. Single-boot deltas under ~10% are noise. + +### TTFT cold vs seeded + +`--greedy`, prefix-cache on. Two identical temp=0 seed=0 requests: + +1. Cold: `prefix_reuse_path=full_reset`. Record TTFT. +2. Seeded: `prefix_reuse_path=seed_prefix` (or restore_*). Record TTFT. + +Pass the seed-store oracle if content is byte-identical and the second path is +not `full_reset`. Compare TTFT native vs container on both the cold and seeded +requests. + +### Weight-load and boot wall time + +Three boots per arm. Median seconds from process start to `/health` ok, and +median seconds of the weight-load phase from the startup log. + +### Soak (after the A/B, not instead of it) + +A multi-hour mixed-traffic soak on native, production-shaped chat + tools, before +native earns any default-local role. The container lane has already survived +hundreds of live agentic requests; native must match that bar. Out of scope for +the first exclusive window if time is short — schedule separately, do not skip. + +## Configure (native, already proven on this machine) + +```bat +call "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat" +set "CUDA_PATH=C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v13.3" +set "PATH=%CUDA_PATH%\bin;%PATH%" +cmake -S . -B build-win -G Ninja -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES=120a -DCMAKE_CUDA_COMPILER="%CUDA_PATH%\bin\nvcc.exe" -DNINFER_BUILD_MEDIA=OFF -DNINFER_BUILD_APPS=ON -DBUILD_TESTING=OFF -DNINFER_BUILD_BENCHMARKS=OFF +cmake --build build-win -j --target ninfer-serve +``` + +## GPU lock + +Coordinator only. `nvidia-smi` >= 20 GiB free; `mkdir C:\Users\igorl\.ninfer-gpu.lock` +(retry 60 s, up to 30 min); remove the lock directory after, success or failure. + +--- + +# GPU runbook: prefix-cache usage observability + +Drive one conversation so all four Engine reuse paths appear, then assert Chat +Completions `usage` matches `--request-log-jsonl` `request_done.result` for the +same request. Schema tests already cover field shape; this is the live match. + +Server: production flags including `--prefix-cache-mib 4096` and +`--request-log-jsonl /tmp/ninfer-usage.jsonl`. Model `qwen3.8-27b`. +`--greedy`. `enable_thinking: false`. + +Conversation (serial, same process): + +1. Unique user prompt A → expect `prefix_reuse_path=full_reset`, + `prefix_cache_hit_tokens=0`. +2. Repeat prompt A as a new request → expect `seed_prefix` and + `prefix_cache_hit_tokens` > 0. +3. Prompt A + assistant reply + new user turn B → expect + `restore_turn_checkpoint` or `append_frontier` (record whichever the log + prints; both are valid). +4. Append another user turn C on that history → expect `append_frontier` or + `restore_turn_checkpoint`. + +For each request, parse the HTTP `usage` object and the matching +`request_done` JSONL event. Pass only if: + +- `usage.prompt_tokens`, `usage.completion_tokens` match `result.prompt_tokens` + / `result.completion_tokens` +- `usage.prefix_cache_hit_tokens` == `result.prefix_cache_hit_tokens` == + `usage.prompt_tokens_details.cached_tokens` +- `usage.prefix_reuse_path` == `result.prefix_reuse_path` (string equality) +- `usage.total_tokens` == prompt + completion (cached_tokens is not an addend) +- `GET /v1/models` `data[0].max_model_len` equals the process `--max-context` + +If a listed path does not appear, do not invent a fifth name; record the +observed path from the log and fail only if usage disagrees with that log. + +--- + +# Runbook: Warmup Fail-Fast, Exception Logging, and Boot Watchdog (Issue #4) + +## Overview + +This runbook documents the operational verification procedures for: +1. **Crash / Terminate Logging**: Ensuring unhandled exceptions escaping thread or server boundaries print `typeid(error).name()` and `error.what()` directly through the console logger before calling `std::abort()` (preventing uninformative silent aborts under container PID 1). +2. **Warmup Fail-Fast**: Ensuring any exception during startup warmup (e.g., CUDA OOM, corrupted prefix cache, invalid batch allocation) terminates the process immediately with non-zero exit code (1) instead of leaving the process listening in an alive-but-503 zombie state. +3. **Pre-Listen Boot Watchdog**: A detached timer thread armed before warmup that terminates the process via `std::_Exit(1)` if boot does not reach the listening state within a configurable budget (default 120 s via `--boot-watchdog-timeout-s`), protecting against uncooperative GPU/driver wedges. +4. **Warmup Timeout Decoupling**: Ensuring startup warmup uses an explicit 60-second budget rather than the short client-facing `--pending-timeout-ms`. +5. **Auto KV-Capacity Bounding**: Clarifying `--kv-capacity auto` description in `--help` to state `(bounded by max-context * max-concurrency)`. + +--- + +## 1. Automated Unit Tests (CPU Container Lane) + +All serve unit tests run and pass inside the build container without requiring GPU hardware: + +```bash +docker run --rm -v "P:\NInfer.gemini:/workspace" -w /workspace \ + -e LD_LIBRARY_PATH=/usr/local/cuda-13.1/compat:/usr/local/cuda-13.1/targets/x86_64-linux/lib/stubs \ + ninfer:test-build bash -c \ + "cd /workspace/build && ctest --output-on-failure -R 'ninfer_(serve_options|http_error_handler|openai_schema|responses_schema|response_store|anthropic_schema|tool_call_parser|request_log|kv_capacity)_test'" +``` + +### Verified Test Cases: +- `ninfer_serve_options_test`: Verifies `--help` text contains `(bounded by max-context * max-concurrency)` for `--kv-capacity auto` and `--boot-watchdog-timeout-s` options. +- `ninfer_http_error_handler_test`: Verifies HTTP error JSON mapping. +- `ninfer_kv_capacity_test`: Verifies sequence capacity curve and page allocation bounds. +- `ninfer_openai_schema_test`, `ninfer_responses_schema_test`, `ninfer_response_store_test`, `ninfer_anthropic_schema_test`, `ninfer_tool_call_parser_test`, `ninfer_request_log_test`: 100% passing. + +--- + +## 2. Induced Failure & Error Path Procedures (GPU Maintenance Window) + +When executed in a coordinator-scheduled GPU maintenance window under the cross-agent lock protocol (`C:\Users\igorl\.ninfer-gpu.lock`): + +### Procedure A: Induce Warmup Failure (OOM / Allocation Fault) +Run `ninfer-serve` with `--prefix-cache-mib` set higher than available GPU VRAM: +```bash +./build/apps/ninfer-serve /path/to/qwen3_8_27b_nvfp4.ninfer --kv-capacity auto --prefix-cache-mib 60000 --port 8018 +``` +**Expected Observable Reality**: +- `httplib` binds the port and sets up the socket backlog at startup. +- Engine initialization or warmup throws `std::runtime_error("warmup generation failed: ...")` or allocation exception. +- Stderr log output: + ``` + [YYYY-MM-DD HH:MM:SS.mmm] [error] ninfer-serve: warmup generation failed: ... + ``` +- The process does NOT enter `server.listen()` (the HTTP accept loop) and terminates immediately with exit code 1. +- Socket is closed upon process exit; no zombie 503 HTTP server remains running. + +### Procedure B: Verify Clean Warmup & Normal Boot +Run `ninfer-serve` with standard production options: +```bash +./build/apps/ninfer-serve /path/to/qwen3_8_27b_nvfp4.ninfer --kv-capacity auto --max-context 131072 --port 8018 +``` +**Expected Observable Reality**: +- Console logs: + ``` + [YYYY-MM-DD HH:MM:SS.mmm] [info] ninfer-serve: loading model... + [YYYY-MM-DD HH:MM:SS.mmm] [info] ninfer-serve: model loaded in ... s + [YYYY-MM-DD HH:MM:SS.mmm] [info] ninfer-serve: KV capacity auto resolved=... + [YYYY-MM-DD HH:MM:SS.mmm] [info] ninfer-serve: warming up... + [YYYY-MM-DD HH:MM:SS.mmm] [info] ninfer-serve: listening on http://127.0.0.1:8018 (model id: ..., auth: disabled) + ``` +- Boot watchdog is cleanly disarmed upon reaching the listening state. +- `curl http://127.0.0.1:8018/v1/models` returns HTTP 200 OK with model descriptor. + +### Procedure C: Verify PID 1 Terminate Logging Handler +In a test container running without a custom init system: +- Trigger an unhandled exception escaping a thread boundary. +**Expected Observable Reality**: +- Stderr log output: + ``` + [YYYY-MM-DD HH:MM:SS.mmm] [error] ninfer-serve: terminate called after throwing : + ``` +- `std::abort()` terminates the process via `SIGABRT` (exit code 134, or container protection fault). + +### Procedure D: Verify Boot Watchdog Hang Protection +To test the uncooperative hang fallback, run with a short watchdog timeout: +```bash +./build/apps/ninfer-serve /path/to/qwen3_8_27b_nvfp4.ninfer --boot-watchdog-timeout-s 1 --port 8018 +``` +**Expected Observable Reality**: +- If model loading or warmup exceeds 1 second: + ``` + [YYYY-MM-DD HH:MM:SS.mmm] [error] ninfer-serve: boot watchdog timeout (1 s) exceeded before reaching listening state; terminating process + ``` +- The watchdog terminates the process immediately via `std::_Exit(1)`. + +--- + +## Issue #5: Tool-calling robustness and multimodal tool results + +Schema and parser test suites (`ninfer_tool_call_parser_test`, `ninfer_openai_schema_test`) +validate the tolerant parser and content parts on CPU without GPU memory. Live-model +validation for multimodal tool results and parallel tool-calling loops is documented below. + +### 1. Multimodal tool result (screenshot in tool message) + +Start server with `--vision` and `--tolerant-tool-calls`: + +```bash +BASE=http://127.0.0.1:8018 +MODEL=qwen3.8-27b + +# Turn 1 + 2: User requests screenshot, assistant calls tool, tool returns image data part +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [ + {\"role\": \"user\", \"content\": \"Take a screenshot and describe what you see.\"}, + { + \"role\": \"assistant\", + \"content\": null, + \"tool_calls\": [ + { + \"id\": \"call_screenshot_001\", + \"type\": \"function\", + \"function\": {\"name\": \"take_screenshot\", \"arguments\": \"{}\"} + } + ] + }, + { + \"role\": \"tool\", + \"tool_call_id\": \"call_screenshot_001\", + \"content\": [ + {\"type\": \"text\", \"text\": \"Screenshot taken successfully:\"}, + { + \"type\": \"image_url\", + \"image_url\": {\"url\": \"data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg==\"} + } + ] + } + ], + \"max_completion_tokens\": 128, + \"temperature\": 0, + \"seed\": 0 +}" +``` + +**Expected Observable Outcome**: +- HTTP 200 OK. +- The Engine decodes the base64 PNG in the tool message turn, preprocesses image patches via the Vision pipeline, and generates a description of the image content. +- `finish_reason` is `"stop"`. +- Token usage reports both text tokens and vision patch tokens in `prompt_tokens`. + +### 2. Live tolerant tool-call recovery + +When `--tolerant-tool-calls` is enabled on the server, the parser recovers complete tool calls even when the model emits duplicate closing tags or trailing suffixes (e.g. duplicate ``, ``, or trailing explanatory text after a complete function), or when the model omits the outer `` closing tag. + +```bash +curl -sS "$BASE/v1/chat/completions" -H 'Content-Type: application/json' -d "{ + \"model\": \"$MODEL\", + \"messages\": [ + {\"role\": \"user\", \"content\": \"Search for weather in Tokyo using the get_weather tool.\"} + ], + \"tools\": [ + { + \"type\": \"function\", + \"function\": { + \"name\": \"get_weather\", + \"description\": \"Get current weather for a city\", + \"parameters\": { + \"type\": \"object\", + \"properties\": { + \"city\": {\"type\": \"string\"} + }, + \"required\": [\"city\"] + } + } + } + ], + \"tool_choice\": \"auto\", + \"max_completion_tokens\": 256, + \"temperature\": 0 +}" +``` + +**Expected Observable Outcome**: +- HTTP 200 OK. +- If the model generation produces a valid `` block with duplicate closing suffixes (e.g. ``) or an unclosed ``: + - `choices[0].finish_reason` is `"tool_calls"`. + - `choices[0].message.tool_calls` contains the parsed function call (`get_weather` with `{"city":"Tokyo"}`). + - `choices[0].message.content` contains any text prefix before the `` tag (or null/empty if none). +- If the output contains near-miss tag syntax (e.g. `` or ``) or is cut off mid-parameter by token limits: + - The turn gracefully degrades to a plain text response with `finish_reason` `"stop"` or `"length"`. + - No internal 500 errors occur, and no phantom tool calls with empty/corrupted arguments are fabricated. + +--- + +# GPU runbook: decode micro-opts A/B (#69 + #67) + +This round is implementation and compile only. Do not boot a server or acquire the +GPU lock until the coordinator schedules an exclusive window. The Qwen3.8-27B +NVFP4 artifact is about 20 GiB and does not fit the 16 GiB compact cap. + +## Arms + +| Arm | Git | Binary | +|---|---|---| +| A baseline | `9dda66511c81e72686ba6b610256625a8af603a7` (`feat/prefix-seed-store` without the thinking dialect) | rebuild `ninfer-serve` from that commit | +| B treatment | `task/issue-1-decode-micro-opts` HEAD | rebuild `ninfer-serve` from this branch | + +Do not base either arm on the vLLM-dialect merge. Decode A/B must stay pure. + +Rebuild image `ninfer:seedstore` from the arm under test. Coordinator owns the +rebuild and the `:8018` lane. Never stop production containers +(`sglang-qwen38` on `:8016`, embeddings, whisper). + +## Server flags (identical both arms) + +No `--preserve-thinking`. Model ID `qwen3.8-27b`. JSONL log required. + +```text +--host 127.0.0.1 --port 8018 +--max-context 131072 --kv-capacity 1048576 --max-concurrency 8 +--spec mtp --draft-tokens 5 --lm-head-draft +--kv-dtype int8 --prefill-chunk 2048 --vision --cors +--prefix-cache-mib 4096 +--request-log-jsonl /tmp/ninfer-decode-ab.jsonl +``` + +## A/B protocol (binding) + +- Identical config both arms. +- At least **3 full process boots per arm** (6 boots total). Interleave + A/B/A/B/A/B in one exclusive GPU session so thermal/clock drift is visible. +- Client load: **c=1, n=16**. One in-flight request. Sixteen serial Chat + Completions per boot. +- Prompts: production-shaped chat from `examples/cli/messages/scenario_*.json` + (code, story, translation, structured). Cycle the twelve scenario fixtures + and repeat the first four to make sixteen. `max_completion_tokens=256`. + Do not send `enable_thinking` / `chat_template_kwargs` (this arm predates + the dialect). +- Speculative: MTP-5 as in the server flags above. Do **not** pass `--greedy` + on the A/B throughput boots (MTP acceptance luck is why boot variance is + large). +- Metric: per-request decode tok/s from the process log line + `decode=` which is `(completion_tokens - 1) / timings_seconds.decode`. + Boot score = median of the 16 request rates. Arm score = median of the 3 + boot scores. Also record all three boot scores so spread is visible. +- **Noise floor:** boot-to-boot decode on this card has been 133.5–155.9 tok/s + from stochastic MTP acceptance. A **single-boot** delta under about 10% is + noise. Do not accept or reject on one boot. +- Pass: arm B median is not a regression versus arm A after three boots. + The upstream claims (+3.7% MoE L2 prefetch, +2.3% node removal) were + measured on an RTX 5090 at T=1; they are unverified on the PRO 6000. + Adopt only if the 3-boot median does not regress. Drop the pick if it + regresses. + +### Throughput commands + +```bash +BASE=http://127.0.0.1:8018 +MODEL=qwen3.8-27b +PROMPTS=( + examples/cli/messages/scenario_code_cuda.json + examples/cli/messages/scenario_code_python.json + examples/cli/messages/scenario_code_typescript.json + examples/cli/messages/scenario_story_zh_scifi.json + examples/cli/messages/scenario_story_en_mystery.json + examples/cli/messages/scenario_story_zh_dialogue.json + examples/cli/messages/scenario_translation_zh_en.json + examples/cli/messages/scenario_translation_en_zh.json + examples/cli/messages/scenario_translation_markdown.json + examples/cli/messages/scenario_structured_jsonl.json + examples/cli/messages/scenario_structured_csv.json + examples/cli/messages/scenario_structured_sql.json +) +# Repeat first four to reach n=16. +for i in $(seq 0 15); do + msg="${PROMPTS[$((i % 12))]}" + python3 - "$BASE" "$MODEL" "$msg" <<'PY' +import json, sys, urllib.request +base, model, path = sys.argv[1], sys.argv[2], sys.argv[3] +body = { + "model": model, + "messages": json.load(open(path, encoding="utf-8")), + "max_completion_tokens": 256, + "stream": False, +} +req = urllib.request.Request( + base + "/v1/chat/completions", + data=json.dumps(body).encode(), + headers={"Content-Type": "application/json"}, + method="POST", +) +with urllib.request.urlopen(req, timeout=600) as r: + json.load(r) +PY +done +``` + +Parse `/tmp/ninfer-decode-ab.jsonl` events with `"event":"request_done"`: + +```text +decode_tok_s = (result.completion_tokens - 1) / timings_seconds.decode +``` + +## Greedy bit-identity: cold vs seeded + +Run **once per arm**, not as part of the 3-boot throughput. Use `--greedy` +in addition to the flags above (same prefix-cache). Cold-start the process. + +Prompt: a two-message chat so the seed frontier is the first user turn. +Send the same body twice. +Pass only if all of: + +1. Both HTTP 200. +2. `choices[0].message.content` is byte-identical across the two responses. +3. First `request_done.prefix_reuse_path` is `full_reset`. +4. Second `request_done.prefix_reuse_path` is `seed_prefix` (or a restore_* + path). `full_reset` on the second request means the seed-store oracle + failed — fail the pick even if content matches by chance. +5. Repeat the same pair on arm A and arm B. Content must match **across + arms** as well (prefetch and node folding must not change tokens). + +If (2) or (5) fails, drop the pick. If (4) fails, the seed-store contract +regressed and the pick is not adoptable. + +### Long greedy generation + +A 32-token probe cannot catch the 1-ulp class recorded in `src/ops/kernel/rope.cuh`: +in-kernel rotary drift that "surfaces as a diverged token deep inside long greedy +generations." Add one long greedy request per arm, same cold-start process as +the short probe (or a third request after it). + +```json +{ + "model": "qwen3.8-27b", + "messages": [{"role": "user", "content": "Write a detailed technical explanation of speculative decoding with MTP, including a worked numeric example."}], + "max_completion_tokens": 2048, + "temperature": 0, + "seed": 0 +} +``` + +Send it twice (cold, then seeded) on arm A and arm B. Pass only if: + +1. All four HTTP 200, `finish_reason` is `stop` or `length`. +2. `choices[0].message.content` is byte-identical across the four responses + (both arms, cold and seeded). +3. Seeded `prefix_reuse_path` is not `full_reset`. + +A mismatch anywhere in the 2048-token body fails the pick even if the 32-token +probe passed. + +## GPU lock + +Only the coordinator runs this. Before any process that allocates GPU +memory: `nvidia-smi` shows at least 20 GiB free; `mkdir` +`C:\Users\igorl\.ninfer-gpu.lock` (retry 60 s up to 30 min); remove the +lock directory after, success or failure. Never stop or restart the +production docker stack. diff --git a/apps/CMakeLists.txt b/apps/CMakeLists.txt index fe80287518..81e30aa112 100644 --- a/apps/CMakeLists.txt +++ b/apps/CMakeLists.txt @@ -17,3 +17,5 @@ target_include_directories(ninfer-serve PRIVATE ${PROJECT_SOURCE_DIR}/third_party ${PROJECT_SOURCE_DIR}/third_party/cpp-httplib) target_link_libraries(ninfer-serve PRIVATE ninfer_serve ninfer_product_load_progress) + +add_subdirectory(ninfer-supervisor) diff --git a/apps/cli/main.cpp b/apps/cli/main.cpp index eb7a029c62..e816189a9a 100644 --- a/apps/cli/main.cpp +++ b/apps/cli/main.cpp @@ -191,6 +191,7 @@ void print_generation_summary(const ninfer::GenerationResult& result, print_metric("kv cache payload", format_bytes(memory.kv_payload_bytes)); print_metric("gpu workspace peak", format_arena_peak(memory.workspace)); print_metric("runtime reservation", format_bytes(memory.runtime_reservation_bytes)); + print_metric("prefix cache", format_bytes(memory.prefix_cache_bytes)); print_metric("free after weights", format_bytes(memory.available_after_weights_bytes)); print_metric("free after startup", format_bytes(memory.available_after_startup_bytes)); print_metric("KV capacity headroom", format_bytes(memory.kv_capacity_headroom_bytes)); diff --git a/apps/cli/options.cpp b/apps/cli/options.cpp index 8016922c97..8b61eae4ab 100644 --- a/apps/cli/options.cpp +++ b/apps/cli/options.cpp @@ -92,7 +92,7 @@ std::string usage_text(const char* argv0) { "--vision enables image/video input and loads the fixed Vision GPU allocations.\n" "--kv-capacity auto leaves " + std::to_string(kDefaultKvCapacityHeadroomBytes / (1024ULL * 1024ULL)) + - " MiB of sizing headroom.\n" + " MiB of sizing headroom (bounded by max-context).\n" "Sampling defaults come from the loaded model and thinking mode; flags override " "individual fields.\n"; } diff --git a/apps/ninfer-supervisor/CMakeLists.txt b/apps/ninfer-supervisor/CMakeLists.txt new file mode 100644 index 0000000000..79e05177bf --- /dev/null +++ b/apps/ninfer-supervisor/CMakeLists.txt @@ -0,0 +1,28 @@ +if(NOT WIN32) + return() +endif() + +add_executable(ninfer-supervisor + main.cpp + engine_child.cpp + collector.cpp + server.cpp + tray.cpp) +target_include_directories(ninfer-supervisor PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/third_party + ${PROJECT_SOURCE_DIR}/third_party/cpp-httplib) +target_link_libraries(ninfer-supervisor PRIVATE CUDA::cudart) +target_link_options(ninfer-supervisor PRIVATE + "LINKER:/DEFAULTLIB:dxgi" + "LINKER:/DEFAULTLIB:ole32" + "LINKER:/DEFAULTLIB:shell32" + "LINKER:/DEFAULTLIB:user32" + "LINKER:/DEFAULTLIB:gdi32" + "LINKER:/DEFAULTLIB:advapi32" + "LINKER:/DEFAULTLIB:ws2_32" + "LINKER:/DEFAULTLIB:crypt32") +target_compile_definitions(ninfer-supervisor PRIVATE + WIN32_LEAN_AND_MEAN + NOMINMAX + _WIN32_WINNT=0x0A00) diff --git a/apps/ninfer-supervisor/collector.cpp b/apps/ninfer-supervisor/collector.cpp new file mode 100644 index 0000000000..18be89fcce --- /dev/null +++ b/apps/ninfer-supervisor/collector.cpp @@ -0,0 +1,482 @@ +#include "collector.hpp" +#include "insights.hpp" + +#ifdef _WIN32 +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#endif +#define CPPHTTPLIB_NO_EXCEPTIONS +#include + +#include +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { +namespace { + +std::string load_key(const std::string& path) { + try { + return read_api_key(path); + } catch (...) { return {}; } +} + +httplib::Client engine_client(const EngineSpec& spec) { + httplib::Client cli(spec.engine_host, spec.engine_port); + cli.set_connection_timeout(1, 0); + cli.set_read_timeout(2, 0); + const std::string key = load_key(spec.api_key_file); + if (!key.empty()) { cli.set_bearer_token_auth(key); } + return cli; +} + +} // namespace + +void Collector::poll_health(Collected& out) { + auto cli = engine_client(spec_); + if (auto res = cli.Get("/health")) { + out.health_status = res->status; + out.health_body = res->body; + } else { + out.health_status = 0; + out.health_body = "unreachable"; + } +} + +void Collector::poll_admin(Collected& out) { + auto cli = engine_client(spec_); + if (auto res = cli.Get("/admin/vram")) { + if (res->status == 200) { + try { + out.admin_vram = nlohmann::json::parse(res->body); + } catch (...) { + out.admin_vram_note = "admin/vram returned unreadable JSON"; + } + } else if (res->status == 401 || res->status == 403) { + out.admin_vram_note = "admin VRAM unavailable (enable --admin-vram and --api-key)"; + } else if (res->status == 404) { + out.admin_vram_note = "admin VRAM not registered on this engine"; + } else { + out.admin_vram_note = "admin/vram HTTP " + std::to_string(res->status); + } + } else { + out.admin_vram_note = "engine unreachable for admin/vram"; + } +} + +void Collector::poll_nvidia_smi(Collected& out) { + FILE* pipe = _popen( + "nvidia-smi --query-gpu=index,memory.used,memory.total " + "--format=csv,noheader,nounits", + "rt"); + if (pipe == nullptr) { + out.nvidia.error = "nvidia-smi not found"; + return; + } + std::string csv; + char buf[512]; + while (fgets(buf, sizeof(buf), pipe) != nullptr) { csv += buf; } + const int rc = _pclose(pipe); + if (rc != 0 && csv.empty()) { + out.nvidia.error = "nvidia-smi exited " + std::to_string(rc); + return; + } + out.nvidia = parse_nvidia_smi_memory_csv(csv, spec_.device); +} + +void Collector::poll_request_log(Collected& out) { + if (spec_.request_log.empty()) { + out.requests.log_error = "request log path not configured"; + return; + } + std::ifstream in(spec_.request_log); + if (!in) { + out.requests.log_error = "request log not present"; + return; + } + out.requests.log_available = true; + std::vector lines; + std::string last_start; + std::string line; + while (std::getline(in, line)) { + if (jsonl_event_is(line, "request_done")) { lines.push_back(line); } + if (jsonl_event_is(line, "server_start")) { last_start = std::move(line); } + } + if (!last_start.empty()) { + try { + const auto j = nlohmann::json::parse(last_start); + const auto& eng = j.at("engine"); + const auto& mem = j.at("memory"); + auto gib = [](const nlohmann::json& obj, const char* key) { + const auto n = obj.value(key, std::uint64_t{0}); + return std::to_string(n / 1048576) + " MiB"; + }; + out.engine_capacity_line = + std::string("KV capacity ") + eng.value("kv_capacity_mode", std::string("?")) + + " resolved=" + std::to_string(eng.value("kv_capacity", 0)) + + " tokens pages=" + std::to_string(eng.value("kv_capacity_page_groups", 0)) + "/" + + std::to_string(eng.value("kv_capacity_max_page_groups", 0)) + + " runtime=" + gib(mem, "runtime_reservation_bytes") + + " prefix-cache=" + gib(mem, "prefix_cache_bytes") + + " free-after-weights=" + gib(mem, "available_after_weights_bytes") + + " free-after-startup=" + gib(mem, "available_after_startup_bytes") + + " headroom=" + gib(mem, "kv_capacity_headroom_bytes") + + " slack=" + gib(mem, "planned_slack_bytes") + + " graphs=" + gib(mem, "cuda_graph_observed_bytes") + "/" + + gib(mem, "cuda_graph_allowance_bytes") + " (from request-log server_start)"; + } catch (...) {} + } + const std::size_t start = lines.size() > 32 ? lines.size() - 32 : 0; + double ttft_sum = 0; + double decode_sum = 0; + int n_ttft = 0; + int n_dec = 0; + for (std::size_t i = start; i < lines.size(); ++i) { + try { + const auto j = nlohmann::json::parse(lines[i]); + // The engine writes {"event":"request_done"}, not "type". Reading the wrong + // key made every record fall through and the panel read a permanent 0. + if (j.value("event", "") != "request_done") { continue; } + ++out.requests.done; + if (j.contains("speculative") && j.at("speculative").is_object()) { + const auto& sp = j.at("speculative"); + out.requests.mtp_backend = sp.value("backend", out.requests.mtp_backend); + out.requests.mtp_draft_window = sp.value("draft_window", out.requests.mtp_draft_window); + const auto drafted = sp.value("drafted_tokens", 0); + const auto accepted = sp.value("accepted_tokens", 0); + out.requests.mtp_drafted += drafted; + out.requests.mtp_accepted += accepted; + out.requests.mtp_fallback_steps += sp.value("fallback_steps", 0); + out.requests.mtp_rounds += sp.value("rounds", 0); + if (drafted > 0) { + out.requests.mtp_last_accept_rate = + static_cast(accepted) / static_cast(drafted); + } + if (sp.contains("accepted_per_position") && sp.at("accepted_per_position").is_array()) { + const auto& pos = sp.at("accepted_per_position"); + if (out.requests.mtp_accepted_per_position.size() < pos.size()) { + out.requests.mtp_accepted_per_position.resize(pos.size(), 0); + } + for (std::size_t p = 0; p < pos.size(); ++p) { + out.requests.mtp_accepted_per_position[p] += pos.at(p).get(); + } + } + } + if (j.contains("timings_seconds") && j.at("timings_seconds").contains("ttft")) { + ttft_sum += j.at("timings_seconds").at("ttft").get() * 1000.0; + ++n_ttft; + } + const auto& result = j.at("result"); + const double dec_s = + j.contains("timings_seconds") ? j.at("timings_seconds").value("decode", 0.0) : 0.0; + const int gen = result.value("completion_tokens", 0); + if (dec_s > 0.0 && gen > 1) { + decode_sum += static_cast(gen - 1) / dec_s; + ++n_dec; + } + const std::string reuse = result.value("prefix_reuse_path", ""); + out.requests.last_reuse = reuse; + if (reuse == "full_reset") { + ++out.requests.reuse_full_reset; + } else if (reuse.find("append") != std::string::npos) { + ++out.requests.reuse_append; + } else if (reuse.find("seed") != std::string::npos || + reuse.find("restore") != std::string::npos) { + ++out.requests.reuse_seed; + } else if (!reuse.empty()) { + ++out.requests.reuse_other; + } + } catch (...) {} + } + if (n_ttft != 0) { out.requests.ttft_ms_mean = ttft_sum / n_ttft; } + if (n_dec != 0) { out.requests.decode_tok_s_mean = decode_sum / n_dec; } +} + +std::int64_t Collector::now_ms() { + using namespace std::chrono; + return duration_cast(system_clock::now().time_since_epoch()).count(); +} + +void Collector::load_persisted_series() { + if (logs_dir_.empty()) { return; } + std::filesystem::create_directories(logs_dir_); + series_path_ = (std::filesystem::path(logs_dir_) / "series.jsonl").string(); + std::ifstream in(series_path_, std::ios::binary); + if (in) { + in.seekg(0, std::ios::end); + const auto sz = static_cast(in.tellg()); + const std::int64_t keep = 2 * 1024 * 1024; + if (sz > keep) { in.seekg(sz - keep, std::ios::beg); } + else { + in.seekg(0, std::ios::beg); + } + std::string body((std::istreambuf_iterator(in)), std::istreambuf_iterator()); + auto nl = body.find('\n'); + if (sz > keep && nl != std::string::npos) { body.erase(0, nl + 1); } + series_.load_jsonl(body); + } + series_file_.open(series_path_, std::ios::app); +} + +void Collector::persist_sample(const VramSample& s) { + if (!series_file_.is_open()) { return; } + series_file_ << format_series_sample_line(s) << '\n'; + series_file_.flush(); +} + +void Collector::persist_event(const VramSeriesEvent& e) { + if (!series_file_.is_open()) { return; } + series_file_ << format_series_event_line(e) << '\n'; + series_file_.flush(); +} + +void Collector::start_series() { + bool expected = false; + if (!series_run_.compare_exchange_strong(expected, true)) { return; } + load_persisted_series(); + series_thread_ = std::thread([this] { series_loop(); }); + observe_thread_ = std::thread([this] { observe_loop(); }); +} + +void Collector::stop_series() { + series_run_ = false; + if (series_thread_.joinable()) { series_thread_.join(); } + if (observe_thread_.joinable()) { observe_thread_.join(); } +} + +void Collector::observe_loop() { + // Health and /admin/vram are HTTP. They do not belong on the 10 Hz DXGI + // loop. They also cannot live only inside snapshot() — that is demand-driven + // by /api/state, so a dashboard that is closed records nothing. 1 Hz is the + // heartbeat: slow enough not to compete with the engine, fast enough that a + // 5 s reclaim is visible even with nobody watching. + while (series_run_.load()) { + const auto t0 = std::chrono::steady_clock::now(); + try { + Collected tmp; + poll_health(tmp); + poll_admin(tmp); + record_transitions(tmp); + std::lock_guard lock(mu_); + detector_last_ran_ms_ = now_ms(); + } catch (...) { + std::lock_guard lock(mu_); + detector_last_ran_ms_ = now_ms(); + } + const auto elapsed = std::chrono::steady_clock::now() - t0; + const auto period = std::chrono::milliseconds(1000); + if (elapsed < period) { std::this_thread::sleep_for(period - elapsed); } + } +} + +void Collector::series_loop() { + // DXGI is an in-process API call, cheap enough to sample at the full rate -- + // and the budget oscillation IS the finding, so it must not be decimated. + // nvidia-smi is a PROCESS SPAWN measured at ~51 ms on this box; polling it + // every tick cost ~10 spawns/s and ~48% of one core, continuously. That does + // not just waste CPU, it perturbs the machine this series exists to observe -- + // the game-test workload it is meant to measure would be competing with it. + // Device totals move slowly, so sample them at 1 Hz and carry the last + // reading forward into the fast series. + constexpr int kNvidiaEvery = 10; + int nvidia_tick = 0; + NvidiaSmiMemory nvidia_last; + while (series_run_.load()) { + const auto t0 = std::chrono::steady_clock::now(); + try { + VramSample sample; + sample.t_ms = now_ms(); + DxgiSnapshot dxgi = query_dxgi_local(spec_.device); + if (nvidia_tick == 0 && dxgi.ok) { + Collected nv; + poll_nvidia_smi(nv); + nvidia_last = nv.nvidia; + } + nvidia_tick = (nvidia_tick + 1) % kNvidiaEvery; + sample.budget_bytes = dxgi.budget_bytes; + sample.nvidia_used_bytes = mib_to_bytes(nvidia_last.used_mib); + { + std::lock_guard lock(mu_); + last_dxgi_ = dxgi; + last_nvidia_ = nvidia_last; + series_.push(sample); + persist_sample(sample); + } + } catch (...) {} + const auto elapsed = std::chrono::steady_clock::now() - t0; + const auto period = std::chrono::milliseconds(100); + if (elapsed < period) { std::this_thread::sleep_for(period - elapsed); } + } +} + +void Collector::record_transitions(const Collected& snap) { + const auto t = now_ms(); + std::lock_guard lock(mu_); + if (last_health_status_ != -1 && last_health_status_ != snap.health_status) { + if (snap.health_status == 200) { + const VramSeriesEvent ev{t, "engine_up", "health 200"}; + series_.push_event(ev); + persist_event(ev); + } else if (last_health_status_ == 200) { + const VramSeriesEvent ev{t, "engine_down", + "health " + std::to_string(snap.health_status)}; + series_.push_event(ev); + persist_event(ev); + } + } + last_health_status_ = snap.health_status; + last_health_body_ = snap.health_body; + last_admin_vram_ = snap.admin_vram; + last_admin_note_ = snap.admin_vram_note; + if (snap.admin_vram.is_object()) { + const std::string trans = snap.admin_vram.value("last_transition", ""); + const std::string reason = snap.admin_vram.value("last_reason", ""); + std::string kind; + if (admin_cursor_.observe(trans, reason, kind)) { + std::string released; + if (snap.admin_vram.contains("tiers") && snap.admin_vram.at("tiers").is_array()) { + for (const auto& tier : snap.admin_vram.at("tiers")) { + if (tier.value("released", false)) { + if (!released.empty()) { released += ","; } + released += tier.value("name", "?"); + } + } + } + std::string label = std::string(trans); + if (!reason.empty()) { + if (!label.empty()) { label += " "; } + label += reason; + } + if (!released.empty()) { label += " released=" + released; } + const VramSeriesEvent ev{t, kind, label}; + series_.push_event(ev); + persist_event(ev); + if (kind == "vram_release") { last_release_ms_ = t; } + if (kind == "vram_reclaim") { last_release_ms_ = 0; } + } + bool any_released = false; + if (snap.admin_vram.contains("tiers") && snap.admin_vram.at("tiers").is_array()) { + for (const auto& tier : snap.admin_vram.at("tiers")) { + if (tier.value("released", false)) { any_released = true; } + } + } + if (!any_released && trans == "reclaim") { last_release_ms_ = 0; } + } +} + +void Collector::note_engine_state(const std::string& state, const std::string& last_event) { + std::lock_guard lock(mu_); + if (!last_engine_state_.empty() && state != last_engine_state_) { + const auto t = now_ms(); + if (state == "Running" || state == "Starting") { + const VramSeriesEvent ev{t, "engine_start", last_event.empty() ? state : last_event}; + series_.push_event(ev); + persist_event(ev); + } else if (state == "Stopped" || state == "Stopping" || state == "Halted") { + const VramSeriesEvent ev{t, "engine_stop", last_event.empty() ? state : last_event}; + series_.push_event(ev); + persist_event(ev); + } + } + last_engine_state_ = state; +} + +nlohmann::json Collector::series_json() { + std::lock_guard lock(mu_); + const auto samples = series_.samples(); + nlohmann::json t_ms = nlohmann::json::array(); + nlohmann::json budget = nlohmann::json::array(); + nlohmann::json nvidia_used = nlohmann::json::array(); + for (const auto& s : samples) { + t_ms.push_back(s.t_ms); + budget.push_back(s.budget_bytes); + nvidia_used.push_back(s.nvidia_used_bytes); + } + nlohmann::json events = nlohmann::json::array(); + for (const auto& e : series_.events()) { + events.push_back({{"t_ms", e.t_ms}, {"kind", e.kind}, {"label", e.label}}); + } + return {{"hz", 10}, + {"raw", true}, + {"t_ms", std::move(t_ms)}, + {"budget_bytes", std::move(budget)}, + {"nvidia_used_bytes", std::move(nvidia_used)}, + {"events", std::move(events)}, + {"detector_last_ran_ms", detector_last_ran_ms_}}; +} + +nlohmann::json Collector::vram_control_json() { + std::lock_guard lock(mu_); + nlohmann::json tiers = nlohmann::json::array(); + bool any_released = false; + if (last_admin_vram_.is_object() && last_admin_vram_.contains("tiers") && + last_admin_vram_.at("tiers").is_array()) { + for (const auto& tier : last_admin_vram_.at("tiers")) { + const bool released = tier.value("released", false); + if (released) { any_released = true; } + tiers.push_back({{"name", tier.value("name", "")}, + {"released", released}, + {"held_bytes", tier.value("held_bytes", 0)}, + {"min_bytes", tier.value("min_bytes", 0)}, + {"max_bytes", tier.value("max_bytes", 0)}, + {"reclaimable_bytes", tier.value("reclaimable_bytes", 0)}}); + } + } + const auto now = now_ms(); + nlohmann::json out = { + {"last_transition", admin_cursor_.last_transition}, + {"last_reason", admin_cursor_.last_reason}, + {"any_released", any_released}, + {"since_release_s", + (any_released && last_release_ms_ > 0) + ? nlohmann::json((now - last_release_ms_) / 1000) + : nlohmann::json(nullptr)}, + {"tiers", std::move(tiers)}, + {"note", last_admin_note_}, + {"detector_last_ran_ms", detector_last_ran_ms_}, + {"detector_age_s", + detector_last_ran_ms_ > 0 ? nlohmann::json((now - detector_last_ran_ms_) / 1000) + : nlohmann::json(nullptr)}, + }; + return out; +} + +nlohmann::json Collector::insights_report() { + auto report = insights_from_request_log_path(spec_.request_log); + nlohmann::json admin; + std::string note; + { + std::lock_guard lock(mu_); + admin = last_admin_vram_; + note = last_admin_note_; + } + append_admin_vram_insights(report, admin, note); + return report; +} + +Collected Collector::snapshot() { + Collected out; + poll_request_log(out); + if (series_run_.load()) { + std::lock_guard lock(mu_); + out.health_status = last_health_status_ < 0 ? 0 : last_health_status_; + out.health_body = last_health_body_; + out.admin_vram = last_admin_vram_; + out.admin_vram_note = last_admin_note_; + out.dxgi = last_dxgi_; + out.nvidia = last_nvidia_; + return out; + } + poll_health(out); + poll_admin(out); + out.dxgi = query_dxgi_local(spec_.device); + poll_nvidia_smi(out); + record_transitions(out); + return out; +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/collector.hpp b/apps/ninfer-supervisor/collector.hpp new file mode 100644 index 0000000000..914b612ca0 --- /dev/null +++ b/apps/ninfer-supervisor/collector.hpp @@ -0,0 +1,101 @@ +#pragma once + +#include "config.hpp" +#include "dxgi_query.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { + +struct RequestMix { + std::uint64_t done = 0; + double ttft_ms_mean = 0; + double decode_tok_s_mean = 0; + std::uint64_t reuse_full_reset = 0; + std::uint64_t reuse_append = 0; + std::uint64_t reuse_seed = 0; + std::uint64_t reuse_other = 0; + std::string last_reuse; + bool log_available = false; + std::string log_error; + std::string mtp_backend; + int mtp_draft_window = 0; + std::uint64_t mtp_drafted = 0; + std::uint64_t mtp_accepted = 0; + std::uint64_t mtp_fallback_steps = 0; + std::uint64_t mtp_rounds = 0; + std::vector mtp_accepted_per_position; + double mtp_last_accept_rate = 0; +}; + +struct Collected { + DxgiSnapshot dxgi; + NvidiaSmiMemory nvidia; + nlohmann::json admin_vram = nullptr; + std::string admin_vram_note; + RequestMix requests; + std::string health_body; + std::string engine_capacity_line; + int health_status = 0; +}; + +class Collector { +public: + explicit Collector(EngineSpec spec, std::string logs_dir = {}) + : spec_(std::move(spec)), logs_dir_(std::move(logs_dir)), series_(6000) {} + ~Collector() { stop_series(); } + + Collector(const Collector&) = delete; + Collector& operator=(const Collector&) = delete; + + void start_series(); + void stop_series(); + Collected snapshot(); + nlohmann::json series_json(); + nlohmann::json vram_control_json(); + nlohmann::json insights_report(); + void note_engine_state(const std::string& state, const std::string& last_event); + +private: + void poll_health(Collected& out); + void poll_admin(Collected& out); + void poll_nvidia_smi(Collected& out); + void poll_request_log(Collected& out); + void series_loop(); + void observe_loop(); + void record_transitions(const Collected& snap); + void persist_sample(const VramSample& s); + void persist_event(const VramSeriesEvent& e); + void load_persisted_series(); + static std::int64_t now_ms(); + + EngineSpec spec_; + std::string logs_dir_; + std::string series_path_; + std::ofstream series_file_; + std::mutex mu_; + VramSeriesRing series_; + std::atomic series_run_{false}; + std::thread series_thread_; + std::thread observe_thread_; + std::int64_t detector_last_ran_ms_ = 0; + int last_health_status_ = -1; + std::string last_health_body_; + AdminVramCursor admin_cursor_; + std::int64_t last_release_ms_ = 0; + std::string last_engine_state_; + nlohmann::json last_admin_vram_ = nullptr; + std::string last_admin_note_; + DxgiSnapshot last_dxgi_; + NvidiaSmiMemory last_nvidia_; +}; + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/config.hpp b/apps/ninfer-supervisor/config.hpp new file mode 100644 index 0000000000..33e124f899 --- /dev/null +++ b/apps/ninfer-supervisor/config.hpp @@ -0,0 +1,108 @@ +#pragma once + +#include "logic.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { + +struct EngineSpec { + std::string executable; + std::vector args; + std::string workdir; + std::string api_key_file; + std::string engine_host = "127.0.0.1"; + int engine_port = 8010; + std::string request_log; + int device = 0; + bool unmanaged = false; // observe an engine this process did not spawn +}; + +struct SupervisorConfig { + EngineSpec engine; + std::string host = "127.0.0.1"; + int port = 8099; + bool bind_any = false; + bool monitor_only = false; // never spawn/stop/restart; HTTP observe only + std::string logs_dir; + bool run_at_login = false; + RestartPolicy restart; +}; + +inline bool manages_engine_process(const SupervisorConfig& cfg) noexcept { + return !cfg.monitor_only && !cfg.engine.unmanaged; +} + +inline std::string read_file_text(const std::string& path) { + std::ifstream in(path, std::ios::binary); + if (!in) { throw std::runtime_error("cannot read " + path); } + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +inline std::string read_api_key(const std::string& path) { + if (path.empty()) { return {}; } + std::string raw = read_file_text(path); + while (!raw.empty() && (raw.back() == '\n' || raw.back() == '\r' || raw.back() == ' ' || + raw.back() == '\t')) { + raw.pop_back(); + } + return raw; +} + +inline SupervisorConfig load_config_json(const std::string& json_text, + bool monitor_only_cli = false) { + const auto body = nlohmann::json::parse(json_text); + SupervisorConfig cfg; + if (body.contains("engine") && body.at("engine").is_object()) { + const auto& e = body.at("engine"); + cfg.engine.executable = e.value("executable", ""); + cfg.engine.workdir = e.value("workdir", ""); + cfg.engine.api_key_file = e.value("api_key_file", ""); + cfg.engine.engine_host = e.value("engine_host", "127.0.0.1"); + cfg.engine.engine_port = e.value("engine_port", 8010); + cfg.engine.request_log = e.value("request_log", ""); + cfg.engine.device = e.value("device", 0); + cfg.engine.unmanaged = e.value("unmanaged", false); + if (e.contains("args") && e.at("args").is_array()) { + for (const auto& a : e.at("args")) { + if (a.is_string()) { cfg.engine.args.push_back(a.get()); } + } + } + } + if (body.contains("supervisor") && body.at("supervisor").is_object()) { + const auto& s = body.at("supervisor"); + cfg.host = s.value("host", "127.0.0.1"); + cfg.port = s.value("port", 8099); + cfg.bind_any = s.value("bind_any", false); + cfg.monitor_only = s.value("monitor_only", false) || monitor_only_cli; + cfg.logs_dir = s.value("logs_dir", ""); + cfg.run_at_login = s.value("run_at_login", false); + if (s.contains("restart") && s.at("restart").is_object()) { + const auto& r = s.at("restart"); + cfg.restart.max_backoff_s = r.value("max_backoff_s", 60); + cfg.restart.crash_loop_window_s = r.value("crash_loop_window_s", 60); + cfg.restart.crash_loop_max = r.value("crash_loop_max", 5); + cfg.restart.health_fail_threshold = r.value("health_fail_threshold", 3); + } + } + if (monitor_only_cli) { cfg.monitor_only = true; } + if (manages_engine_process(cfg) && cfg.engine.executable.empty()) { + throw std::invalid_argument( + "engine.executable is required unless monitor_only or engine.unmanaged"); + } + if (!cfg.bind_any && !is_loopback_host(cfg.host)) { + throw std::invalid_argument( + "supervisor host must be loopback unless bind_any is true"); + } + return cfg; +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/dashboard.hpp b/apps/ninfer-supervisor/dashboard.hpp new file mode 100644 index 0000000000..262a528d5f --- /dev/null +++ b/apps/ninfer-supervisor/dashboard.hpp @@ -0,0 +1,252 @@ +#pragma once + +#include + +namespace ninfer::supervisor { + +inline constexpr std::string_view kDashboardHtml = R"HTML( + + + +NInfer supervisor + + + + +
+

NInfer supervisor

+
loopback control surface · live SSE
+
+
+
+

Engine

+
state…
+
health…
+
pid—
+
uptime—
+
restarts0
+
last event—
+
+ + + +
+
+
+

VRAM

+
adapter—
+
DXGI budget (system-wide WDDM pressure)—
+
device used (nvidia-smi)—
+
device total (nvidia-smi)—
+
supervisor process DXGI (not the engine)—
+
engine capacity (boot line)—
+
admin tiers—
+
released now—
+
time since release—
+
last vram action—
+
detector last ran—
+
admin note—
+
+
+

VRAM + DXGI budget (raw 10 Hz · no smoothing)

+ +
+ DXGI budget (WDDM pressure) + nvidia-smi used (physical) + engine / admin-vram events +
+
+
+

Recent requests

+
done (window)—
+
mean TTFT—
+
mean decode—
+
reuse mix—
+
MTP (captured)—
+
log—
+
+
+

Engine log tail

+ 
waiting…
+
+
+

Insights

+

same objects as GET /api/insights

+
+
+
+ + + +)HTML"; + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/dxgi_query.hpp b/apps/ninfer-supervisor/dxgi_query.hpp new file mode 100644 index 0000000000..d3c118a041 --- /dev/null +++ b/apps/ninfer-supervisor/dxgi_query.hpp @@ -0,0 +1,166 @@ +#pragma once + +// Cached DXGI QueryVideoMemoryInfo. CUDA is used at most once to bind an +// adapter LUID; the 10 Hz path is DXGI-only. Every HRESULT is checked. Failure +// degrades to ok=false rather than throwing or aborting. The factory is +// recreated only after a failed query so a TDR/device-removed does not take +// the process down. + +#include +#include + +#include +#include +#include +#include + +namespace ninfer::supervisor { + +struct DxgiSnapshot { + std::uint64_t budget_bytes = 0; + std::uint64_t current_usage_bytes = 0; + std::uint64_t available_for_reservation_bytes = 0; + std::uint64_t current_reservation_bytes = 0; + std::string adapter_name; + bool ok = false; + std::string error; +}; + +class DxgiBudgetSource { +public: + DxgiBudgetSource() = default; + ~DxgiBudgetSource() { reset(); } + + DxgiBudgetSource(const DxgiBudgetSource&) = delete; + DxgiBudgetSource& operator=(const DxgiBudgetSource&) = delete; + + DxgiSnapshot query(int device_index) { + DxgiSnapshot out; + try { + std::lock_guard lock(mu_); + if (!query_locked(device_index, out)) { + reset_locked(); + query_locked(device_index, out); + } + } catch (...) { + out = DxgiSnapshot{}; + out.error = "DXGI query threw"; + try { + std::lock_guard lock(mu_); + reset_locked(); + } catch (...) {} + } + return out; + } + + void reset() { + std::lock_guard lock(mu_); + reset_locked(); + } + +private: + static constexpr UINT kNvidia = 0x10DE; + + void reset_locked() { + if (adapter3_ != nullptr) { + adapter3_->Release(); + adapter3_ = nullptr; + } + if (factory_ != nullptr) { + factory_->Release(); + factory_ = nullptr; + } + bound_device_ = -1; + adapter_name_.clear(); + } + + bool bind_locked(int device_index, DxgiSnapshot& out) { + IDXGIFactory1* factory = nullptr; + HRESULT hr = CreateDXGIFactory1(__uuidof(IDXGIFactory1), reinterpret_cast(&factory)); + if (FAILED(hr) || factory == nullptr) { + out.error = "CreateDXGIFactory1 failed hr=" + std::to_string(static_cast(hr)); + return false; + } + + // Do not call CUDA here. cudaGetDeviceProperties during a TDR can abort + // the process (0xC0000409), which is the failure mode we exist to record. + UINT index = 0; + int gpu_seen = 0; + IDXGIAdapter3* chosen = nullptr; + std::string chosen_name; + for (;;) { + IDXGIAdapter1* a1 = nullptr; + hr = factory->EnumAdapters1(index, &a1); + ++index; + if (hr == DXGI_ERROR_NOT_FOUND) { break; } + if (FAILED(hr) || a1 == nullptr) { break; } + DXGI_ADAPTER_DESC1 desc{}; + if (FAILED(a1->GetDesc1(&desc))) { + a1->Release(); + continue; + } + if ((desc.Flags & DXGI_ADAPTER_FLAG_SOFTWARE) != 0) { + a1->Release(); + continue; + } + if (desc.VendorId != kNvidia || gpu_seen++ != device_index) { + a1->Release(); + continue; + } + IDXGIAdapter3* a3 = nullptr; + hr = a1->QueryInterface(__uuidof(IDXGIAdapter3), reinterpret_cast(&a3)); + a1->Release(); + if (FAILED(hr) || a3 == nullptr) { continue; } + char name[128]{}; + WideCharToMultiByte(CP_UTF8, 0, desc.Description, -1, name, + static_cast(sizeof(name)), nullptr, nullptr); + chosen = a3; + chosen_name = name; + break; + } + if (chosen == nullptr) { + factory->Release(); + out.error = "no DXGI adapter matched"; + return false; + } + factory_ = factory; + adapter3_ = chosen; + adapter_name_ = chosen_name; + bound_device_ = device_index; + return true; + } + + bool query_locked(int device_index, DxgiSnapshot& out) { + if (factory_ == nullptr || adapter3_ == nullptr || bound_device_ != device_index) { + if (!bind_locked(device_index, out)) { return false; } + } + DXGI_QUERY_VIDEO_MEMORY_INFO info{}; + const HRESULT hr = + adapter3_->QueryVideoMemoryInfo(0, DXGI_MEMORY_SEGMENT_GROUP_LOCAL, &info); + if (FAILED(hr)) { + out.error = "QueryVideoMemoryInfo hr=" + std::to_string(static_cast(hr)); + return false; + } + out.budget_bytes = info.Budget; + out.current_usage_bytes = info.CurrentUsage; + out.available_for_reservation_bytes = info.AvailableForReservation; + out.current_reservation_bytes = info.CurrentReservation; + out.adapter_name = adapter_name_; + out.ok = true; + out.error.clear(); + return true; + } + + std::mutex mu_; + IDXGIFactory1* factory_ = nullptr; + IDXGIAdapter3* adapter3_ = nullptr; + int bound_device_ = -1; + std::string adapter_name_; +}; + +inline DxgiSnapshot query_dxgi_local(int cuda_device) { + static DxgiBudgetSource source; + return source.query(cuda_device); +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/engine_child.cpp b/apps/ninfer-supervisor/engine_child.cpp new file mode 100644 index 0000000000..6e881bc8c8 --- /dev/null +++ b/apps/ninfer-supervisor/engine_child.cpp @@ -0,0 +1,327 @@ +#include "engine_child.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { +namespace { + +std::int64_t unix_ms() { + using namespace std::chrono; + return duration_cast(system_clock::now().time_since_epoch()).count(); +} + +std::wstring utf8_to_wide(const std::string& s) { + if (s.empty()) { return {}; } + const int n = MultiByteToWideChar(CP_UTF8, 0, s.c_str(), -1, nullptr, 0); + std::wstring out(static_cast(n), L'\0'); + MultiByteToWideChar(CP_UTF8, 0, s.c_str(), -1, out.data(), n); + if (!out.empty() && out.back() == L'\0') { out.pop_back(); } + return out; +} + +std::wstring quote_arg(const std::string& a) { + std::wstring w = utf8_to_wide(a); + if (w.find_first_of(L" \t\"") == std::wstring::npos) { return w; } + std::wstring q = L"\""; + for (wchar_t c : w) { + if (c == L'"') { q += L"\\\""; } else { q += c; } + } + q += L'"'; + return q; +} + +void close_handle(void*& h) { + if (h != nullptr) { + CloseHandle(static_cast(h)); + h = nullptr; + } +} + +} // namespace + +EngineChild::EngineChild(SupervisorConfig cfg) : cfg_(std::move(cfg)), gate_(cfg_.restart) { + if (cfg_.logs_dir.empty()) { cfg_.logs_dir = "ninfer-supervisor-logs"; } + std::filesystem::create_directories(cfg_.logs_dir); + log_path_ = (std::filesystem::path(cfg_.logs_dir) / "engine.log").string(); + if (!manages_engine_process(cfg_)) { + auto_restart_ = false; + st_.last_event = "monitor-only: not managing engine process"; + } +} + +EngineChild::~EngineChild() { + quit_ = true; + stop(); + close_handle(job_handle_); +} + +void EngineChild::request_quit() { quit_ = true; } + +EngineStatus EngineChild::status() const { + std::lock_guard lock(mu_); + EngineStatus s = st_; + s.crash_loop_halted = gate_.halted(); + s.health_fails = gate_.health_fails(); + return s; +} + +std::string EngineChild::log_tail(std::size_t max_bytes) const { + std::ifstream in(log_path_, std::ios::binary); + if (!in) { return {}; } + in.seekg(0, std::ios::end); + const auto size = static_cast(in.tellg()); + const std::size_t off = size > max_bytes ? size - max_bytes : 0; + in.seekg(static_cast(off)); + std::ostringstream out; + out << in.rdbuf(); + return out.str(); +} + +void EngineChild::append_log(const char* data, std::size_t n) { + rotate_logs_if_needed(); + std::ofstream out(log_path_, std::ios::binary | std::ios::app); + if (out) { out.write(data, static_cast(n)); } +} + +void EngineChild::rotate_logs_if_needed() { + std::error_code ec; + const auto sz = std::filesystem::file_size(log_path_, ec); + if (ec || sz < (8ULL << 20)) { return; } + const auto rotated = log_path_ + ".1"; + std::filesystem::remove(rotated, ec); + std::filesystem::rename(log_path_, rotated, ec); +} + +void EngineChild::start() { + if (!manages_engine_process(cfg_)) { return; } + auto_restart_ = true; + gate_.reset_halt(); + std::lock_guard lock(mu_); + if (st_.state == EngineState::Running || st_.state == EngineState::Starting) { return; } + st_.last_event = "start requested"; +} + +void EngineChild::stop() { + if (!manages_engine_process(cfg_)) { return; } + auto_restart_ = false; + stop_child_ = true; + HANDLE proc = nullptr; + { + std::lock_guard lock(mu_); + proc = static_cast(process_handle_); + st_.state = EngineState::Stopping; + st_.last_event = "stop requested"; + } + if (proc != nullptr) { TerminateProcess(proc, 1); } +} + +void EngineChild::observe_health(int http_status) { + if (!manages_engine_process(cfg_)) { + std::lock_guard lock(mu_); + if (http_status == 200) { + st_.health = "ok"; + st_.state = EngineState::Running; + st_.last_event = "unmanaged engine reachable"; + } else if (http_status == 503) { + st_.health = "unhealthy"; + st_.state = EngineState::Running; + st_.last_event = "unmanaged engine unhealthy"; + } else { + st_.health = "unreachable"; + st_.state = EngineState::Stopped; + st_.last_event = "unmanaged engine unreachable"; + } + return; + } + bool restart_now = false; + { + std::lock_guard lock(mu_); + if (http_status == 200) { + gate_.note_healthy(); + st_.health = "ok"; + return; + } + if (http_status == 503) { + st_.health = "unhealthy"; + if (gate_.note_health_fail() && auto_restart_.load()) { + st_.last_event = "health restart threshold"; + restart_now = true; + } + } else { + st_.health = "unreachable"; + } + } + if (restart_now) { restart(); } +} + +void EngineChild::restart() { + if (!manages_engine_process(cfg_)) { return; } + auto_restart_ = true; + gate_.reset_halt(); + stop_child_ = true; + HANDLE proc = nullptr; + { + std::lock_guard lock(mu_); + proc = static_cast(process_handle_); + st_.last_event = "restart requested"; + } + if (proc != nullptr) { TerminateProcess(proc, 1); } +} + +void EngineChild::spawn() { + stop_child_ = false; + std::wstring cmd = quote_arg(cfg_.engine.executable); + for (const auto& a : cfg_.engine.args) { + cmd += L' '; + cmd += quote_arg(a); + } + std::vector cmd_buf(cmd.begin(), cmd.end()); + cmd_buf.push_back(L'\0'); + + SECURITY_ATTRIBUTES sa{}; + sa.nLength = sizeof(sa); + sa.bInheritHandle = TRUE; + + HANDLE out_r = nullptr; + HANDLE out_w = nullptr; + if (!CreatePipe(&out_r, &out_w, &sa, 0)) { + throw std::runtime_error("CreatePipe failed"); + } + SetHandleInformation(out_r, HANDLE_FLAG_INHERIT, 0); + + STARTUPINFOW si{}; + si.cb = sizeof(si); + si.dwFlags = STARTF_USESTDHANDLES; + si.hStdOutput = out_w; + si.hStdError = out_w; + si.hStdInput = GetStdHandle(STD_INPUT_HANDLE); + + PROCESS_INFORMATION pi{}; + const std::wstring cwd = utf8_to_wide(cfg_.engine.workdir); + const wchar_t* cwd_ptr = cwd.empty() ? nullptr : cwd.c_str(); + if (!CreateProcessW(nullptr, cmd_buf.data(), nullptr, nullptr, TRUE, CREATE_NO_WINDOW, + nullptr, cwd_ptr, &si, &pi)) { + CloseHandle(out_r); + CloseHandle(out_w); + throw std::runtime_error("CreateProcessW failed"); + } + CloseHandle(out_w); + CloseHandle(pi.hThread); + + if (job_handle_ == nullptr) { + HANDLE job = CreateJobObjectW(nullptr, nullptr); + JOBOBJECT_EXTENDED_LIMIT_INFORMATION lim{}; + lim.BasicLimitInformation.LimitFlags = JOB_OBJECT_LIMIT_KILL_ON_JOB_CLOSE; + SetInformationJobObject(job, JobObjectExtendedLimitInformation, &lim, sizeof(lim)); + job_handle_ = job; + } + AssignProcessToJobObject(static_cast(job_handle_), pi.hProcess); + + { + std::lock_guard lock(mu_); + process_handle_ = pi.hProcess; + st_.pid = static_cast(pi.dwProcessId); + st_.state = EngineState::Running; + st_.started_unix_ms = unix_ms(); + ++st_.restart_count; + st_.last_event = "engine started"; + } + + std::thread reader([this, out_r] { + char buf[4096]; + DWORD n = 0; + while (ReadFile(out_r, buf, sizeof(buf), &n, nullptr) && n > 0) { append_log(buf, n); } + CloseHandle(out_r); + }); + reader.detach(); +} + +void EngineChild::capture_wait() { + HANDLE proc = nullptr; + { + std::lock_guard lock(mu_); + proc = static_cast(process_handle_); + } + if (proc == nullptr) { return; } + WaitForSingleObject(proc, INFINITE); + DWORD code = 0; + GetExitCodeProcess(proc, &code); + { + std::lock_guard lock(mu_); + st_.last_exit_code = static_cast(code); + st_.pid = 0; + process_handle_ = nullptr; + st_.state = EngineState::Stopped; + st_.last_event = "engine exited"; + } + CloseHandle(proc); +} + +void EngineChild::run_loop() { + while (!quit_.load()) { + const bool running = [&] { + std::lock_guard lock(mu_); + return process_handle_ != nullptr; + }(); + if (running) { + capture_wait(); + const bool intentional = stop_child_.exchange(false); + if (quit_.load() || !auto_restart_.load() || intentional) { continue; } + bool allow = false; + { + std::lock_guard lock(mu_); + allow = gate_.note_exit(std::chrono::steady_clock::now()); + } + if (!allow) { + std::lock_guard lock(mu_); + st_.state = EngineState::Halted; + st_.last_event = "crash-loop breaker: too many exits"; + auto_restart_ = false; + continue; + } + int wait_s = 0; + { + std::lock_guard lock(mu_); + gate_.advance_backoff(); + wait_s = gate_.backoff_seconds(); + } + { + std::lock_guard lock(mu_); + st_.state = EngineState::BackingOff; + st_.last_event = "backing off"; + } + for (int i = 0; i < wait_s * 10 && !quit_.load() && auto_restart_.load(); ++i) { + Sleep(100); + } + continue; + } + if (manages_engine_process(cfg_) && auto_restart_.load() && !gate_.halted() && + !quit_.load()) { + try { + { + std::lock_guard lock(mu_); + st_.state = EngineState::Starting; + } + spawn(); + } catch (const std::exception& ex) { + std::lock_guard lock(mu_); + st_.state = EngineState::Stopped; + st_.last_event = std::string("spawn failed: ") + ex.what(); + auto_restart_ = false; + } + continue; + } + Sleep(200); + } + stop(); +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/engine_child.hpp b/apps/ninfer-supervisor/engine_child.hpp new file mode 100644 index 0000000000..17169c22cd --- /dev/null +++ b/apps/ninfer-supervisor/engine_child.hpp @@ -0,0 +1,72 @@ +#pragma once + +#include "config.hpp" + +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { + +enum class EngineState : std::uint8_t { + Stopped, + Starting, + Running, + Stopping, + BackingOff, + Halted, +}; + +struct EngineStatus { + EngineState state = EngineState::Stopped; + std::uint64_t pid = 0; + std::int64_t started_unix_ms = 0; + int restart_count = 0; + int last_exit_code = 0; + bool crash_loop_halted = false; + std::string last_event; + std::string health; // ok / unhealthy / unreachable + int health_fails = 0; +}; + +class EngineChild { +public: + explicit EngineChild(SupervisorConfig cfg); + ~EngineChild(); + + EngineChild(const EngineChild&) = delete; + EngineChild& operator=(const EngineChild&) = delete; + + void start(); + void stop(); + void restart(); + void request_quit(); + void observe_health(int http_status); + + [[nodiscard]] EngineStatus status() const; + [[nodiscard]] std::string log_tail(std::size_t max_bytes) const; + [[nodiscard]] const SupervisorConfig& config() const noexcept { return cfg_; } + + void run_loop(); + +private: + void spawn(); + void capture_wait(); + void append_log(const char* data, std::size_t n); + void rotate_logs_if_needed(); + + SupervisorConfig cfg_; + RestartGate gate_; + mutable std::mutex mu_; + EngineStatus st_; + std::atomic stop_child_{false}; + std::atomic quit_{false}; + std::atomic auto_restart_{true}; + void* process_handle_ = nullptr; // HANDLE + void* job_handle_ = nullptr; + std::string log_path_; +}; + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/insights.hpp b/apps/ninfer-supervisor/insights.hpp new file mode 100644 index 0000000000..060ff1bb78 --- /dev/null +++ b/apps/ninfer-supervisor/insights.hpp @@ -0,0 +1,531 @@ +#pragma once + +#include "logic.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { + +inline nlohmann::json insight_unavailable(std::string id, std::string title, std::string statement, + nlohmann::json evidence = nlohmann::json::object(), + nlohmann::json measured_over = {{"requests", 0}}) { + return {{"id", std::move(id)}, + {"severity", "notice"}, + {"title", std::move(title)}, + {"statement", std::move(statement)}, + {"evidence", std::move(evidence)}, + {"confidence", "measured"}, + {"measured_over", std::move(measured_over)}, + {"availability", "unavailable"}}; +} + +inline nlohmann::json insight_available(std::string id, std::string severity, std::string title, + std::string statement, nlohmann::json evidence, + std::string recommendation, std::string confidence, + nlohmann::json measured_over) { + nlohmann::json out = {{"id", std::move(id)}, + {"severity", std::move(severity)}, + {"title", std::move(title)}, + {"statement", std::move(statement)}, + {"evidence", std::move(evidence)}, + {"confidence", std::move(confidence)}, + {"measured_over", std::move(measured_over)}, + {"availability", "available"}}; + if (!recommendation.empty()) { out["recommendation"] = std::move(recommendation); } + return out; +} + +inline std::int64_t json_i64(const nlohmann::json& j, const char* key, std::int64_t fallback = 0) { + if (!j.contains(key)) { return fallback; } + const auto& v = j.at(key); + if (v.is_number_integer()) { return v.get(); } + if (v.is_number()) { return static_cast(v.get()); } + return fallback; +} + +inline double json_f64(const nlohmann::json& j, const char* key, double fallback = 0) { + if (!j.contains(key)) { return fallback; } + const auto& v = j.at(key); + if (v.is_number()) { return v.get(); } + return fallback; +} + +// Analyze a JSONL blob. Does not invent content/reasoning_content — those keys are +// not written by the engine request log (schema_version 10). +inline nlohmann::json analyze_request_log_jsonl(std::string_view jsonl, std::string_view path) { + nlohmann::json report = { + {"source", + {{"request_log", jsonl.empty() ? "empty" : "ok"}, {"path", std::string(path)}}}, + {"insights", nlohmann::json::array()}, + }; + + std::unordered_map starts; + std::vector dones; + std::vector throughputs; + std::int64_t tmin = 0; + std::int64_t tmax = 0; + int parsed = 0; + + std::string line; + std::istringstream in{std::string(jsonl)}; + while (std::getline(in, line)) { + if (line.empty()) { continue; } + nlohmann::json j; + try { + j = nlohmann::json::parse(line); + } catch (...) { continue; } + const std::string event = j.value("event", ""); + if (event.empty()) { continue; } + ++parsed; + const auto ts = json_i64(j, "timestamp_unix_ms"); + if (tmin == 0 || ts < tmin) { tmin = ts; } + if (ts > tmax) { tmax = ts; } + if (event == "request_start" && j.contains("request")) { + const auto rid = json_i64(j.at("request"), "request_id"); + starts[j.value("server_instance_id", "") + ":" + std::to_string(rid)] = j; + } else if (event == "request_done") { + dones.push_back(std::move(j)); + } else if (event == "throughput") { + throughputs.push_back(std::move(j)); + } + } + + auto& insights = report["insights"]; + const double window_s = + (tmax > tmin) ? static_cast(tmax - tmin) / 1000.0 : 0.0; + + if (parsed == 0) { + insights.push_back(insight_unavailable( + "source.request_log", "Request log has no usable records", + "no request_done records in window", + {{"path", std::string(path)}, {"parsed_events", 0}})); + report["generated_note"] = "unavailable is not a clean zero"; + return report; + } + + if (dones.empty()) { + insights.push_back(insight_unavailable( + "source.request_done", "No completed requests in window", + "no request_done records in window", + {{"parsed_events", parsed}, {"request_start", starts.size()}}, + {{"requests", 0}, {"parsed_events", parsed}})); + return report; + } + + struct Bucket { + int queued = 0; + int prefill = 0; + int decode = 0; + int mixed = 0; + int unpaired = 0; + double queue_wait_sum = 0; + double prepare_sum = 0; + double prefill_sum = 0; + double decode_sum = 0; + double vision_sum = 0; + double ttft_sum = 0; + double total_sum = 0; + std::vector queued_ids; + std::vector prefill_ids; + std::vector decode_ids; + } b; + + int output_limit_thinking = 0; + int output_limit_hit_cap = 0; + int thinking_requests = 0; + int tools_declared = 0; + std::vector output_limit_ids; + std::vector output_limit_caps; + int reuse_reset_single = 0; + int reuse_reset_multi = 0; + int reuse_restore = 0; + int reuse_seed = 0; + int reuse_append = 0; + int reuse_other = 0; + std::uint64_t multi_prompt_tokens = 0; + std::uint64_t multi_hit_tokens = 0; + std::vector reset_multi_samples; + + for (const auto& done : dones) { + const auto& req = done.contains("request") ? done.at("request") : nlohmann::json::object(); + const auto id = json_i64(req, "request_id"); + if (req.value("enable_thinking", false)) { ++thinking_requests; } + if (json_i64(req, "tool_count") > 0) { ++tools_declared; } + const auto& result = done.contains("result") ? done.at("result") : nlohmann::json::object(); + const std::string finish = result.value("finish_reason", ""); + const int cap = static_cast(json_i64(req, "requested_output_tokens")); + const int completion = static_cast(json_i64(result, "completion_tokens")); + if (finish == "output_limit" && req.value("enable_thinking", false)) { + ++output_limit_thinking; + output_limit_ids.push_back(id); + output_limit_caps.push_back(cap); + if (cap > 0 && completion >= cap) { ++output_limit_hit_cap; } + } + const int messages = static_cast(json_i64(req, "message_count")); + const std::string reuse = result.value("prefix_reuse_path", ""); + const auto prompt_tokens = static_cast(json_i64(result, "prompt_tokens")); + const auto hit_tokens = static_cast(json_i64(result, "prefix_cache_hit_tokens")); + const bool multiturn = messages >= 2; + if (multiturn) { + multi_prompt_tokens += prompt_tokens; + multi_hit_tokens += hit_tokens; + } + if (reuse == "full_reset") { + if (multiturn) { + ++reuse_reset_multi; + if (reset_multi_samples.size() < 8) { + reset_multi_samples.push_back({{"request_id", id}, + {"message_count", messages}, + {"prompt_tokens", prompt_tokens}, + {"prefix_cache_hit_tokens", hit_tokens}}); + } + } else { + ++reuse_reset_single; + } + } else if (reuse.find("restore") != std::string::npos) { + ++reuse_restore; + } else if (reuse.find("seed") != std::string::npos) { + ++reuse_seed; + } else if (reuse.find("append") != std::string::npos) { + ++reuse_append; + } else if (!reuse.empty()) { + ++reuse_other; + } + + const auto& timings = + done.contains("timings_seconds") ? done.at("timings_seconds") : nlohmann::json::object(); + const double total = json_f64(timings, "total"); + const double prepare = json_f64(timings, "prepare"); + const double prefill = json_f64(timings, "prefill"); + const double decode = json_f64(timings, "decode"); + const double vision = json_f64(timings, "vision"); + const double ttft = json_f64(timings, "ttft"); + b.prepare_sum += prepare; + b.prefill_sum += prefill; + b.decode_sum += decode; + b.vision_sum += vision; + b.ttft_sum += ttft; + b.total_sum += total; + + const std::string join = + done.value("server_instance_id", "") + ":" + std::to_string(id); + auto it = starts.find(join); + if (it == starts.end()) { + ++b.unpaired; + continue; + } + const double wall_s = + static_cast(json_i64(done, "timestamp_unix_ms") - + json_i64(it->second, "timestamp_unix_ms")) / + 1000.0; + double queue_wait = wall_s - total; + if (queue_wait < 0.0) { queue_wait = 0.0; } + b.queue_wait_sum += queue_wait; + const bool queued = queue_wait >= 0.020 && wall_s > 0.0 && queue_wait >= 0.25 * wall_s; + if (queued) { + ++b.queued; + if (b.queued_ids.size() < 8) { b.queued_ids.push_back(id); } + } else if (total > 0.0 && prefill >= decode && prefill >= 0.4 * total) { + ++b.prefill; + if (b.prefill_ids.size() < 8) { b.prefill_ids.push_back(id); } + } else if (total > 0.0 && decode >= 0.4 * total) { + ++b.decode; + if (b.decode_ids.size() < 8) { b.decode_ids.push_back(id); } + } else { + ++b.mixed; + } + } + + const int paired = static_cast(dones.size()) - b.unpaired; + int max_waiting = 0; + int max_running = 0; + int max_prefill = 0; + for (const auto& tp : throughputs) { + if (!tp.contains("scheduler")) { continue; } + const auto& sch = tp.at("scheduler"); + max_waiting = std::max(max_waiting, static_cast(json_i64(sch, "waiting"))); + max_running = std::max(max_running, static_cast(json_i64(sch, "running"))); + max_prefill = std::max(max_prefill, static_cast(json_i64(sch, "prefilling"))); + } + + const auto over = nlohmann::json{{"requests", dones.size()}, + {"paired_requests", paired}, + {"unpaired_done", b.unpaired}, + {"window_s", window_s}, + {"throughput_events", throughputs.size()}}; + + const double mean_queue = paired > 0 ? b.queue_wait_sum / paired : 0.0; + const double mean_prefill = + dones.empty() ? 0.0 : b.prefill_sum / static_cast(dones.size()); + const double mean_decode = + dones.empty() ? 0.0 : b.decode_sum / static_cast(dones.size()); + const int cause_max = std::max({b.queued, b.prefill, b.decode, b.mixed}); + std::string cause = "mixed"; + std::string cause_id = "latency.mixed"; + std::vector cause_ids; + if (cause_max == b.queued && b.queued > 0) { + cause = "queued behind concurrency"; + cause_id = "latency.queued_behind_concurrency"; + cause_ids = b.queued_ids; + } else if (cause_max == b.prefill && b.prefill > 0) { + cause = "long prefill"; + cause_id = "latency.prefill_dominated"; + cause_ids = b.prefill_ids; + } else if (cause_max == b.decode && b.decode > 0) { + cause = "decode-dominated"; + cause_id = "latency.decode_dominated"; + cause_ids = b.decode_ids; + } + + std::ostringstream sat; + sat << paired << " paired of " << dones.size() << " request_done: " << b.queued + << " queued, " << b.prefill << " prefill-dominated, " << b.decode + << " decode-dominated, " << b.mixed << " mixed. Dominant cause: " << cause + << ". Mean queue wait " << (mean_queue * 1000.0) << " ms, mean prefill " + << (mean_prefill * 1000.0) << " ms, mean decode " << (mean_decode * 1000.0) + << " ms. Scheduler peak waiting=" << max_waiting << " running=" << max_running + << " prefilling=" << max_prefill << "."; + + const bool pressure = b.queued > 0 && (b.queued * 3 >= paired || max_waiting > 0); + insights.push_back(insight_available( + cause_id, pressure ? "warning" : "info", "Saturation vs latency", sat.str(), + {{"queued", b.queued}, + {"prefill_dominated", b.prefill}, + {"decode_dominated", b.decode}, + {"mixed", b.mixed}, + {"mean_queue_wait_s", mean_queue}, + {"mean_prefill_s", mean_prefill}, + {"mean_decode_s", mean_decode}, + {"scheduler_peak", + {{"waiting", max_waiting}, {"running", max_running}, {"prefilling", max_prefill}}}, + {"sample_request_ids", cause_ids}}, + pressure ? "Queued wait is a concurrency/backlog problem, not a slow kernel. " + "Raise --max-concurrency only if KV/headroom allows; otherwise the " + "engine is saturated." + : "", + "measured", over)); + + const double n_done = static_cast(dones.size()); + const double mean_prepare = n_done > 0 ? b.prepare_sum / n_done : 0.0; + const double mean_vision = n_done > 0 ? b.vision_sum / n_done : 0.0; + const double mean_ttft = n_done > 0 ? b.ttft_sum / n_done : 0.0; + const double ttft_body = mean_prepare + mean_prefill + mean_vision; + std::string ttft_cause = "mixed"; + std::string ttft_id = "latency.ttft_mixed"; + std::string ttft_rec; + if (mean_prefill >= mean_prepare && mean_prefill >= mean_vision && mean_prefill >= 0.4 * std::max(ttft_body, mean_ttft)) { + ttft_cause = "prefill"; + ttft_id = "latency.ttft_prefill_dominated"; + ttft_rec = "TTFT is prefill-dominated. --prefill-chunk is the lever, not decode kernels."; + } else if (mean_prepare >= mean_prefill && mean_prepare >= 0.4 * std::max(ttft_body, mean_ttft)) { + ttft_cause = "prepare"; + ttft_id = "latency.ttft_prepare_dominated"; + ttft_rec = "TTFT is prepare-dominated (tokenize/media), not GPU decode."; + } else if (mean_vision >= 0.4 * std::max(ttft_body, mean_ttft) && mean_vision > 0.0) { + ttft_cause = "vision"; + ttft_id = "latency.ttft_vision_dominated"; + ttft_rec = "TTFT is vision-preprocess dominated."; + } + std::ostringstream ttft_stmt; + ttft_stmt << "Mean TTFT " << (mean_ttft * 1000.0) << " ms over " << dones.size() + << " request_done: prepare " << (mean_prepare * 1000.0) << " ms, prefill " + << (mean_prefill * 1000.0) << " ms, vision " << (mean_vision * 1000.0) + << " ms (decode " << (mean_decode * 1000.0) + << " ms is after first token). Dominant TTFT component: " << ttft_cause << "."; + insights.push_back(insight_available( + ttft_id, "info", "TTFT decomposition", ttft_stmt.str(), + {{"mean_ttft_s", mean_ttft}, + {"mean_prepare_s", mean_prepare}, + {"mean_prefill_s", mean_prefill}, + {"mean_vision_s", mean_vision}, + {"mean_decode_s", mean_decode}, + {"dominant", ttft_cause}}, + ttft_rec, "measured", over)); + + const double multi_hit_ratio = + multi_prompt_tokens == 0 + ? 0.0 + : static_cast(multi_hit_tokens) / static_cast(multi_prompt_tokens); + std::ostringstream reuse_stmt; + reuse_stmt << "Reuse mix over " << dones.size() << " request_done: full_reset single-turn " + << reuse_reset_single << " (expected), full_reset multi-turn " << reuse_reset_multi + << ", restore " << reuse_restore << ", seed " << reuse_seed << ", append " + << reuse_append << ", other " << reuse_other << ". Multi-turn prefix-hit ratio " + << (multi_hit_ratio * 100.0) << "% (" << multi_hit_tokens << "/" + << multi_prompt_tokens << " tokens)."; + insights.push_back(insight_available( + "prefix.reuse_mix", reuse_reset_multi > 0 ? "notice" : "info", "Prefix-cache reuse mix", + reuse_stmt.str(), + {{"full_reset_single_turn", reuse_reset_single}, + {"full_reset_multi_turn", reuse_reset_multi}, + {"restore", reuse_restore}, + {"seed", reuse_seed}, + {"append", reuse_append}, + {"other", reuse_other}, + {"multi_turn_prompt_tokens", multi_prompt_tokens}, + {"multi_turn_hit_tokens", multi_hit_tokens}, + {"multi_turn_hit_ratio", multi_hit_ratio}}, + "", "measured", over)); + if (reuse_reset_multi > 0) { + std::ostringstream miss; + miss << reuse_reset_multi << " of " << dones.size() + << " request_done were multi-turn (message_count>=2) on full_reset with " + << "prefix_cache_hit_tokens often 0. A single-message full_reset is expected; " + << "a multi-turn full_reset is a miss that should have been restore or seed."; + insights.push_back(insight_available( + "prefix.multiturn_full_reset", "warning", "Multi-turn conversations resetting the prefix", + miss.str(), + {{"full_reset_multi_turn", reuse_reset_multi}, + {"full_reset_single_turn", reuse_reset_single}, + {"samples", reset_multi_samples}}, + "Check seed store / turn checkpoints. restore_turn_checkpoint or seed_prefix " + "should fire when message_count>=2.", + "measured", over)); + } + + // Content/reasoning_content are not in schema_version 10 request logs. + insights.push_back(insight_unavailable( + "client.content_fields", "Visitor content is not in the request log", + "result.content and reasoning_content are not written to request_done; " + "empty-reply-vs-reasoning cannot be confirmed from this source", + {{"schema_version", 10}, + {"looked_for", nlohmann::json::array({"content", "reasoning_content"})}, + {"requests", dones.size()}}, + over)); + + if (output_limit_thinking > 0) { + std::ostringstream stmt; + stmt << output_limit_thinking << " of " << dones.size() + << " request_done finished on output_limit with enable_thinking=true" + << " (" << output_limit_hit_cap << " also hit requested_output_tokens). " + << thinking_requests << " of " << dones.size() << " had thinking enabled."; + int cap_sum = 0; + for (int c : output_limit_caps) { cap_sum += c; } + const double cap_mean = + output_limit_caps.empty() + ? 0.0 + : static_cast(cap_sum) / static_cast(output_limit_caps.size()); + insights.push_back(insight_available( + "client.output_limit_while_thinking", + output_limit_thinking * 5 >= static_cast(dones.size()) ? "warning" : "notice", + "Thinking requests hitting output_limit", stmt.str(), + {{"output_limit_thinking", output_limit_thinking}, + {"hit_requested_cap", output_limit_hit_cap}, + {"thinking_requests", thinking_requests}, + {"mean_requested_output_tokens", cap_mean}, + {"sample_request_ids", output_limit_ids}}, + "Inferred: a thinking model with a small max_tokens can spend the budget on " + "reasoning and return an empty visitor reply. Content fields are not in this log, " + "so raise requested_output_tokens and compare finish_reason.", + "measured", over)); + } + + insights.push_back(insight_unavailable( + "client.narrated_tool_intent", "Narrated tool intent cannot be scored from JSONL", + "detecting narrated-intent-with-no-tools needs visitor-facing text; the request log " + "does not store content. tool_count is measurable and is reported in evidence.", + {{"requests_with_tools", tools_declared}, + {"requests", dones.size()}, + {"requests_without_tools", static_cast(dones.size()) - tools_declared}}, + over)); + + return report; +} + +inline void append_admin_vram_insights(nlohmann::json& report, const nlohmann::json& admin, + const std::string& note) { + if (!report.contains("insights") || !report["insights"].is_array()) { + report["insights"] = nlohmann::json::array(); + } + if (!admin.is_object()) { + report["insights"].push_back(insight_unavailable( + "vram.admin", "Admin VRAM is not available", + note.empty() ? "admin/vram was not readable; cannot tell if any tier is releasable" + : note, + {{"note", note}})); + return; + } + nlohmann::json pinned = nlohmann::json::array(); + nlohmann::json released = nlohmann::json::array(); + if (admin.contains("tiers") && admin.at("tiers").is_array()) { + for (const auto& tier : admin.at("tiers")) { + const auto min_b = json_i64(tier, "min_bytes"); + const auto max_b = json_i64(tier, "max_bytes"); + const bool rel = tier.value("released", false); + if (rel) { released.push_back(tier.value("name", "?")); } + if (min_b > 0 && min_b == max_b) { + pinned.push_back({{"name", tier.value("name", "")}, + {"min_bytes", min_b}, + {"max_bytes", max_b}, + {"reclaimable_bytes", json_i64(tier, "reclaimable_bytes")}, + {"released", rel}}); + } + } + } + const auto over = nlohmann::json{{"requests", 0}, + {"admin_tiers", pinned.size() + released.size()}, + {"last_transition", admin.value("last_transition", "")}, + {"last_reason", admin.value("last_reason", "")}}; + if (!pinned.empty()) { + report["insights"].push_back(insight_available( + "vram.tier_pinned_unreleasable", "warning", + "Admin VRAM is enabled but a tier cannot be released", + "A tier has min_bytes == max_bytes while --admin-vram is on. " + "--prefix-cache-mib N pins seed min=max=N, so reclaimable_bytes stays 0 " + "and the admin surface looks healthy while nothing can be released.", + {{"pinned_tiers", pinned}, + {"last_transition", admin.value("last_transition", "")}, + {"last_reason", admin.value("last_reason", "")}}, + "Omit --prefix-cache-mib or set a max above min if you want idle release.", + "measured", over)); + } + if (!released.empty()) { + report["insights"].push_back(insight_available( + "vram.tier_currently_released", "notice", + "A VRAM tier is currently released", + "Released tiers: " + released.dump() + + ". The engine is serving degraded (no cross-request prefix seeding) until reclaim. " + "Release is ~120x cheaper than reclaim on this hardware.", + {{"released", released}, + {"last_transition", admin.value("last_transition", "")}, + {"last_reason", admin.value("last_reason", "")}}, + "Reclaim before a traffic burst; a released seed store will full_reset more often.", + "measured", over)); + } +} + +inline nlohmann::json insights_from_request_log_path(const std::string& path) { + nlohmann::json report; + report["insights"] = nlohmann::json::array(); + if (path.empty()) { + report["source"] = {{"request_log", "unconfigured"}, {"path", ""}}; + report["insights"].push_back(insight_unavailable( + "source.request_log", "Request log is not configured", + "no request_done records in window", {{"path", ""}})); + return report; + } + std::ifstream in(path); + if (!in) { + report["source"] = {{"request_log", "missing"}, {"path", path}}; + report["insights"].push_back(insight_unavailable( + "source.request_log", "Request log is not present", + "no request_done records in window", {{"path", path}})); + return report; + } + std::ostringstream body; + body << in.rdbuf(); + return analyze_request_log_jsonl(body.str(), path); +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/logic.hpp b/apps/ninfer-supervisor/logic.hpp new file mode 100644 index 0000000000..cc78c27c9b --- /dev/null +++ b/apps/ninfer-supervisor/logic.hpp @@ -0,0 +1,414 @@ +#pragma once + +#include +#include +#include +#include +#include +#include + +namespace ninfer::supervisor { + +inline bool is_loopback_host(std::string_view host) { + return host == "127.0.0.1" || host == "::1" || host == "localhost" || host == "localhost."; +} + +inline bool is_loopback_peer(std::string_view addr) { + if (addr.empty()) { return false; } + if (is_loopback_host(addr)) { return true; } + // httplib may report IPv4-mapped IPv6. + return addr == "::ffff:127.0.0.1"; +} + +inline constexpr std::string_view kSupervisorControlHeader = "X-NInfer-Supervisor"; +inline constexpr std::string_view kSupervisorControlHeaderValue = "1"; + +inline std::string_view trim_sv(std::string_view s) { + while (!s.empty() && (s.front() == ' ' || s.front() == '\t' || s.front() == '\r' || + s.front() == '\n')) { + s.remove_prefix(1); + } + while (!s.empty() && (s.back() == ' ' || s.back() == '\t' || s.back() == '\r' || + s.back() == '\n')) { + s.remove_suffix(1); + } + return s; +} + +// Split Host into name and optional port. IPv6 literals must be bracketed when a +// port is present (`[::1]:8099`). +inline bool split_host_header(std::string_view host, std::string& name, int& port, bool& has_port) { + host = trim_sv(host); + name.clear(); + port = 0; + has_port = false; + if (host.empty()) { return false; } + if (host.front() == '[') { + const auto rb = host.find(']'); + if (rb == std::string_view::npos) { return false; } + name = std::string(host.substr(0, rb + 1)); + if (rb + 1 == host.size()) { return true; } + if (host[rb + 1] != ':') { return false; } + const auto p = host.substr(rb + 2); + if (p.empty()) { return false; } + int value = 0; + for (char c : p) { + if (c < '0' || c > '9') { return false; } + value = value * 10 + (c - '0'); + if (value > 65535) { return false; } + } + port = value; + has_port = true; + return true; + } + const auto colon = host.rfind(':'); + if (colon != std::string_view::npos && host.find(':') == colon) { + name = std::string(host.substr(0, colon)); + const auto p = host.substr(colon + 1); + if (p.empty() || name.empty()) { return false; } + int value = 0; + for (char c : p) { + if (c < '0' || c > '9') { return false; } + value = value * 10 + (c - '0'); + if (value > 65535) { return false; } + } + port = value; + has_port = true; + return true; + } + name = std::string(host); + return !name.empty(); +} + +inline bool is_loopback_host_name(std::string_view name) { + return is_loopback_host(name) || name == "[::1]"; +} + +// DNS-rebinding defense: only the listen port's loopback names, plus the +// configured bind host when --bind-any names a specific interface. Binding +// 0.0.0.0 does not open the Host allowlist. +inline bool host_header_allowed(std::string_view host_header, int listen_port, + std::string_view bind_host, bool bind_any) { + std::string name; + int port = 0; + bool has_port = false; + if (!split_host_header(host_header, name, port, has_port)) { return false; } + if (has_port && port != listen_port) { return false; } + if (is_loopback_host_name(name)) { return true; } + if (!bind_any) { return false; } + if (bind_host.empty() || bind_host == "0.0.0.0" || bind_host == "::" || bind_host == "[::]") { + return false; + } + return name == bind_host; +} + +inline bool supervisor_control_header_ok(std::string_view value) { + return trim_sv(value) == kSupervisorControlHeaderValue; +} + +struct NvidiaSmiMemory { + bool ok = false; + int index = -1; + std::uint64_t used_mib = 0; + std::uint64_t total_mib = 0; + std::string error; +}; + +// Parses `nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv,noheader,nounits`. +// Values are mebibytes. Picks the row whose index equals `device`. +inline NvidiaSmiMemory parse_nvidia_smi_memory_csv(std::string_view csv, int device) { + NvidiaSmiMemory out; + std::string_view rest = csv; + bool saw_row = false; + while (!rest.empty()) { + auto nl = rest.find_first_of("\n\r"); + auto line = trim_sv(nl == std::string_view::npos ? rest : rest.substr(0, nl)); + rest = nl == std::string_view::npos ? std::string_view{} + : rest.substr(nl + 1); + if (line.empty()) { continue; } + saw_row = true; + const auto c1 = line.find(','); + if (c1 == std::string_view::npos) { continue; } + const auto c2 = line.find(',', c1 + 1); + if (c2 == std::string_view::npos) { continue; } + const auto idx_s = trim_sv(line.substr(0, c1)); + const auto used_s = trim_sv(line.substr(c1 + 1, c2 - c1 - 1)); + const auto tot_s = trim_sv(line.substr(c2 + 1)); + int idx = 0; + std::uint64_t used = 0; + std::uint64_t tot = 0; + try { + idx = std::stoi(std::string(idx_s)); + used = std::stoull(std::string(used_s)); + tot = std::stoull(std::string(tot_s)); + } catch (...) { continue; } + if (idx != device) { continue; } + out.ok = true; + out.index = idx; + out.used_mib = used; + out.total_mib = tot; + return out; + } + out.error = saw_row ? "nvidia-smi csv has no row for the configured device" + : "nvidia-smi csv is empty"; + return out; +} + +inline std::uint64_t mib_to_bytes(std::uint64_t mib) { return mib * 1024ull * 1024ull; } + +// Pre-filter that agrees with the engine JSONL schema: the field is "event", +// not "type". A substring on the value without the key name is how a panel +// can look populated while every record is then discarded. +inline bool jsonl_event_is(std::string_view line, std::string_view event) { + const std::string compact = std::string("\"event\":\"") + std::string(event) + "\""; + const std::string spaced = std::string("\"event\": \"") + std::string(event) + "\""; + return line.find(compact) != std::string_view::npos || + line.find(spaced) != std::string_view::npos; +} + +struct VramSample { + std::int64_t t_ms = 0; + std::uint64_t budget_bytes = 0; + std::uint64_t nvidia_used_bytes = 0; +}; + +struct VramSeriesEvent { + std::int64_t t_ms = 0; + std::string kind; + std::string label; +}; + +// Diff last_transition/last_reason, not held_bytes. A 42 ms release finishes +// between 1 Hz polls; the persisted last_reason is what remains observable. +struct AdminVramCursor { + bool seen = false; + std::string last_transition; + std::string last_reason; + + // Returns true if this observation is an event after the baseline poll. + bool observe(std::string_view trans, std::string_view reason, std::string& kind) { + if (!seen) { + seen = true; + last_transition = std::string(trans); + last_reason = std::string(reason); + return false; + } + if (trans == last_transition && reason == last_reason) { return false; } + last_transition = std::string(trans); + last_reason = std::string(reason); + if (trans == "release") { + kind = "vram_release"; + } else if (trans == "reclaim" || trans == "reclaim-failed") { + kind = "vram_reclaim"; + } else { + kind = "admin_vram"; + } + return true; + } +}; + +inline bool parse_series_event_line(std::string_view line, VramSeriesEvent& ev); +inline bool parse_series_sample_line(std::string_view line, VramSample& s); + +// Ring of raw 10 Hz samples. No averaging. Oldest is dropped on overflow. +struct VramSeriesRing { + explicit VramSeriesRing(std::size_t cap = 6000) : cap_(cap), buf_(cap) {} + + void push(VramSample s) { + if (cap_ == 0) { return; } + buf_[head_] = s; + head_ = (head_ + 1) % cap_; + if (size_ < cap_) { ++size_; } + } + + void push_event(VramSeriesEvent e, std::size_t event_cap = 128) { + events_.push_back(std::move(e)); + while (events_.size() > event_cap) { events_.pop_front(); } + } + + [[nodiscard]] std::vector samples() const { + std::vector out; + out.reserve(size_); + const std::size_t start = size_ < cap_ ? 0 : head_; + for (std::size_t i = 0; i < size_; ++i) { out.push_back(buf_[(start + i) % cap_]); } + return out; + } + + [[nodiscard]] std::vector events() const { + return {events_.begin(), events_.end()}; + } + + [[nodiscard]] std::size_t size() const noexcept { return size_; } + + void load_jsonl(std::string_view jsonl) { + std::string_view rest = jsonl; + while (!rest.empty()) { + auto nl = rest.find('\n'); + auto line = trim_sv(nl == std::string_view::npos ? rest : rest.substr(0, nl)); + rest = nl == std::string_view::npos ? std::string_view{} : rest.substr(nl + 1); + if (line.empty()) { continue; } + VramSeriesEvent ev; + VramSample samp; + if (parse_series_event_line(line, ev)) { + push_event(std::move(ev)); + } else if (parse_series_sample_line(line, samp)) { + push(samp); + } + } + } + +private: + std::size_t cap_ = 0; + std::size_t head_ = 0; + std::size_t size_ = 0; + std::vector buf_; + std::deque events_; +}; + +inline bool extract_json_i64(std::string_view line, std::string_view key, std::int64_t& out) { + const std::string pat = "\"" + std::string(key) + "\":"; + auto pos = line.find(pat); + if (pos == std::string_view::npos) { return false; } + pos += pat.size(); + while (pos < line.size() && (line[pos] == ' ')) { ++pos; } + bool neg = false; + if (pos < line.size() && line[pos] == '-') { + neg = true; + ++pos; + } + if (pos >= line.size() || line[pos] < '0' || line[pos] > '9') { return false; } + std::int64_t v = 0; + while (pos < line.size() && line[pos] >= '0' && line[pos] <= '9') { + v = v * 10 + (line[pos] - '0'); + ++pos; + } + out = neg ? -v : v; + return true; +} + +inline bool extract_json_str(std::string_view line, std::string_view key, std::string& out) { + const std::string pat = "\"" + std::string(key) + "\":\""; + auto pos = line.find(pat); + if (pos == std::string_view::npos) { return false; } + pos += pat.size(); + std::string s; + while (pos < line.size() && line[pos] != '"') { + s.push_back(line[pos]); + ++pos; + } + out = std::move(s); + return true; +} + +inline bool parse_series_event_line(std::string_view line, VramSeriesEvent& ev) { + if (line.find("\"kind\"") == std::string_view::npos) { return false; } + std::int64_t t = 0; + if (!extract_json_i64(line, "t_ms", t)) { return false; } + ev.t_ms = t; + extract_json_str(line, "kind", ev.kind); + extract_json_str(line, "label", ev.label); + return !ev.kind.empty(); +} + +inline bool parse_series_sample_line(std::string_view line, VramSample& s) { + if (line.find("\"kind\"") != std::string_view::npos) { return false; } + std::int64_t t = 0, b = 0, n = 0; + if (!extract_json_i64(line, "t_ms", t)) { return false; } + extract_json_i64(line, "budget_bytes", b); + extract_json_i64(line, "nvidia_used_bytes", n); + s.t_ms = t; + s.budget_bytes = static_cast(b); + s.nvidia_used_bytes = static_cast(n); + return true; +} + +inline std::string format_series_sample_line(const VramSample& s) { + return std::string("{\"t_ms\":") + std::to_string(s.t_ms) + + ",\"budget_bytes\":" + std::to_string(s.budget_bytes) + + ",\"nvidia_used_bytes\":" + std::to_string(s.nvidia_used_bytes) + "}"; +} + +inline std::string format_series_event_line(const VramSeriesEvent& e) { + return std::string("{\"t_ms\":") + std::to_string(e.t_ms) + ",\"kind\":\"" + e.kind + + "\",\"label\":\"" + e.label + "\"}"; +} + +inline std::string extract_kv_capacity_line(std::string_view log) { + const auto key = std::string_view("KV capacity "); + const auto pos = log.rfind(key); + if (pos == std::string_view::npos) { return {}; } + auto start = log.find_last_of("\n", pos); + start = start == std::string_view::npos ? 0 : start + 1; + auto end = log.find('\n', pos); + auto line = log.substr(start, end == std::string_view::npos ? log.size() - start : end - start); + if (!line.empty() && line.back() == '\r') { line.remove_suffix(1); } + return std::string(trim_sv(line)); +} + +struct RestartPolicy { + int initial_backoff_s = 1; + int max_backoff_s = 60; + int crash_loop_max = 5; + int crash_loop_window_s = 60; + int health_fail_threshold = 3; +}; + +class RestartGate { +public: + explicit RestartGate(RestartPolicy policy = {}) : policy_(policy), backoff_s_(policy.initial_backoff_s) {} + + // Record an engine exit. Returns false if auto-restart is halted (crash loop). + bool note_exit(std::chrono::steady_clock::time_point now) { + if (halted_) { return false; } + const auto window = std::chrono::seconds(policy_.crash_loop_window_s); + while (!exits_.empty() && now - exits_.front() > window) { exits_.pop_front(); } + exits_.push_back(now); + if (static_cast(exits_.size()) >= policy_.crash_loop_max) { + halted_ = true; + return false; + } + return true; + } + + [[nodiscard]] int backoff_seconds() const noexcept { return backoff_s_; } + + void advance_backoff() { + if (backoff_s_ < policy_.max_backoff_s) { + const int next = backoff_s_ * 2; + backoff_s_ = next > policy_.max_backoff_s ? policy_.max_backoff_s : next; + } + } + + void note_healthy() { + backoff_s_ = policy_.initial_backoff_s; + health_fails_ = 0; + } + + bool note_health_fail() { + ++health_fails_; + return health_fails_ >= policy_.health_fail_threshold; + } + + void clear_health_fails() { health_fails_ = 0; } + + void reset_halt() { + halted_ = false; + exits_.clear(); + backoff_s_ = policy_.initial_backoff_s; + health_fails_ = 0; + } + + [[nodiscard]] bool halted() const noexcept { return halted_; } + [[nodiscard]] int recent_exits() const noexcept { return static_cast(exits_.size()); } + [[nodiscard]] int health_fails() const noexcept { return health_fails_; } + [[nodiscard]] const RestartPolicy& policy() const noexcept { return policy_; } + +private: + RestartPolicy policy_; + std::deque exits_; + int backoff_s_ = 1; + int health_fails_ = 0; + bool halted_ = false; +}; + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/main.cpp b/apps/ninfer-supervisor/main.cpp new file mode 100644 index 0000000000..e16923a361 --- /dev/null +++ b/apps/ninfer-supervisor/main.cpp @@ -0,0 +1,143 @@ +#include "collector.hpp" +#include "config.hpp" +#include "engine_child.hpp" +#include "logic.hpp" +#include "server.hpp" +#include "tray.hpp" + +#include + +#include +#include +#include +#include +#include + +namespace { + +void install_run_at_login(const std::string& command) { + HKEY key = nullptr; + if (RegCreateKeyExW(HKEY_CURRENT_USER, L"Software\\Microsoft\\Windows\\CurrentVersion\\Run", 0, + nullptr, 0, KEY_SET_VALUE, nullptr, &key, nullptr) != ERROR_SUCCESS) { + throw std::runtime_error("cannot open Run key"); + } + std::wstring w(command.begin(), command.end()); + const LONG st = + RegSetValueExW(key, L"NInferSupervisor", 0, REG_SZ, + reinterpret_cast(w.c_str()), + static_cast((w.size() + 1) * sizeof(wchar_t))); + RegCloseKey(key); + if (st != ERROR_SUCCESS) { throw std::runtime_error("cannot write Run key"); } +} + +void uninstall_run_at_login() { + HKEY key = nullptr; + if (RegOpenKeyExW(HKEY_CURRENT_USER, L"Software\\Microsoft\\Windows\\CurrentVersion\\Run", 0, + KEY_SET_VALUE, &key) != ERROR_SUCCESS) { + return; + } + RegDeleteValueW(key, L"NInferSupervisor"); + RegCloseKey(key); +} + +void usage() { + std::cout + << "usage: ninfer-supervisor --config FILE [--host 127.0.0.1] [--port 8099] [--bind-any]\n" + " [--monitor-only] [--install-login] [--uninstall-login]\n" + " Dashboard binds loopback by default. --bind-any is required for 0.0.0.0 and prints\n" + " a warning. Control POST /api/start|stop|restart is loopback-peer only, requires\n" + " header X-NInfer-Supervisor: 1, and returns 409 in --monitor-only / unmanaged mode.\n" + " Host is allowlisted on every route. Do not send CORS headers.\n"; +} + +} // namespace + +int main(int argc, char** argv) { + try { + std::string config_path; + std::string host_override; + int port_override = -1; + bool bind_any = false; + bool monitor_only = false; + bool install = false; + bool uninstall = false; + for (int i = 1; i < argc; ++i) { + const std::string a = argv[i]; + auto need = [&](const char* name) -> const char* { + if (i + 1 >= argc) { throw std::invalid_argument(std::string("missing ") + name); } + return argv[++i]; + }; + if (a == "--help" || a == "-h") { + usage(); + return 0; + } else if (a == "--config") { + config_path = need("--config"); + } else if (a == "--host") { + host_override = need("--host"); + } else if (a == "--port") { + port_override = std::stoi(need("--port")); + } else if (a == "--bind-any") { + bind_any = true; + } else if (a == "--monitor-only") { + monitor_only = true; + } else if (a == "--install-login") { + install = true; + } else if (a == "--uninstall-login") { + uninstall = true; + } else { + throw std::invalid_argument("unknown argument: " + a); + } + } + if (uninstall) { + uninstall_run_at_login(); + std::cout << "removed HKCU Run\\NInferSupervisor\n"; + return 0; + } + if (config_path.empty()) { + usage(); + return 2; + } + auto cfg = ninfer::supervisor::load_config_json( + ninfer::supervisor::read_file_text(config_path), monitor_only); + if (!host_override.empty()) { cfg.host = host_override; } + if (port_override > 0) { cfg.port = port_override; } + if (bind_any) { cfg.bind_any = true; } + if (monitor_only) { cfg.monitor_only = true; } + if (cfg.bind_any) { + std::cerr << "WARNING: binding beyond loopback; engine start/stop is exposed on " + << (cfg.host.empty() ? "0.0.0.0" : cfg.host) << ":" << cfg.port << "\n"; + } else if (!ninfer::supervisor::is_loopback_host(cfg.host)) { + throw std::invalid_argument("--host must be loopback without --bind-any"); + } + if (install) { + char module[MAX_PATH]{}; + GetModuleFileNameA(nullptr, module, MAX_PATH); + const std::string cmd = std::string("\"") + module + "\" --config \"" + config_path + "\""; + install_run_at_login(cmd); + std::cout << "installed HKCU Run\\NInferSupervisor\n"; + } + + ninfer::supervisor::EngineChild child(cfg); + ninfer::supervisor::Collector collector(cfg.engine, cfg.logs_dir); + collector.start_series(); + ninfer::supervisor::DashboardServer server(cfg, child, collector); + std::thread engine_thread([&] { child.run_loop(); }); + std::thread http_thread([&] { server.run(); }); + const std::string url = + "http://" + (cfg.bind_any ? std::string("127.0.0.1") : cfg.host) + ":" + + std::to_string(cfg.port) + "/"; + std::cout << "ninfer-supervisor dashboard " << url << "\n"; + ninfer::supervisor::TrayIcon tray(child, url, ninfer::supervisor::manages_engine_process(cfg)); + tray.run(); + server.stop(); + collector.stop_series(); + child.request_quit(); + child.stop(); + if (http_thread.joinable()) { http_thread.join(); } + if (engine_thread.joinable()) { engine_thread.join(); } + return 0; + } catch (const std::exception& ex) { + std::cerr << "ninfer-supervisor: " << ex.what() << "\n"; + return 1; + } +} diff --git a/apps/ninfer-supervisor/server.cpp b/apps/ninfer-supervisor/server.cpp new file mode 100644 index 0000000000..632447994e --- /dev/null +++ b/apps/ninfer-supervisor/server.cpp @@ -0,0 +1,199 @@ +#include "server.hpp" +#include "insights.hpp" + +#ifdef _WIN32 +#ifndef WIN32_LEAN_AND_MEAN +#define WIN32_LEAN_AND_MEAN +#endif +#endif +#define CPPHTTPLIB_NO_EXCEPTIONS +#include + +#include +#include +#include +#include + +namespace ninfer::supervisor { +namespace { + +const char* state_name(EngineState s) { + switch (s) { + case EngineState::Stopped: return "Stopped"; + case EngineState::Starting: return "Starting"; + case EngineState::Running: return "Running"; + case EngineState::Stopping: return "Stopping"; + case EngineState::BackingOff: return "BackingOff"; + case EngineState::Halted: return "Halted"; + } + return "?"; +} + +std::int64_t now_unix_s() { + using namespace std::chrono; + return duration_cast(system_clock::now().time_since_epoch()).count(); +} + +} // namespace + +DashboardServer::DashboardServer(SupervisorConfig cfg, EngineChild& child, Collector& collector) + : cfg_(std::move(cfg)), child_(child), collector_(collector) {} + +DashboardServer::~DashboardServer() { stop(); } + +bool DashboardServer::control_allowed(const std::string& remote) const { + return is_loopback_peer(remote); +} + +nlohmann::json DashboardServer::state_json() { + const EngineStatus st = child_.status(); + Collected snap = collector_.snapshot(); + nlohmann::json engine = { + {"state", state_name(st.state)}, + {"pid", st.pid}, + {"restart_count", st.restart_count}, + {"last_exit_code", st.last_exit_code}, + {"last_event", st.last_event}, + {"crash_loop_halted", st.crash_loop_halted}, + {"uptime_s", st.started_unix_ms == 0 + ? 0 + : now_unix_s() - st.started_unix_ms / 1000}, + }; + nlohmann::json dxgi = { + {"ok", snap.dxgi.ok}, + {"error", snap.dxgi.error}, + {"adapter_name", snap.dxgi.adapter_name}, + {"budget_bytes", snap.dxgi.budget_bytes}, + {"supervisor_usage_bytes", snap.dxgi.current_usage_bytes}, + {"supervisor_usage_note", "DXGI CurrentUsage of the supervisor process, not the engine"}, + }; + nlohmann::json nvidia = {{"ok", snap.nvidia.ok}, + {"error", snap.nvidia.error}, + {"index", snap.nvidia.index}, + {"used_bytes", mib_to_bytes(snap.nvidia.used_mib)}, + {"total_bytes", mib_to_bytes(snap.nvidia.total_mib)}}; + nlohmann::json req = {{"done", snap.requests.done}, + {"ttft_ms_mean", snap.requests.ttft_ms_mean}, + {"decode_tok_s_mean", snap.requests.decode_tok_s_mean}, + {"reuse_full_reset", snap.requests.reuse_full_reset}, + {"reuse_append", snap.requests.reuse_append}, + {"reuse_seed", snap.requests.reuse_seed}, + {"last_reuse", snap.requests.last_reuse}, + {"log_available", snap.requests.log_available}, + {"log_error", snap.requests.log_error}, + {"mtp_backend", snap.requests.mtp_backend}, + {"mtp_draft_window", snap.requests.mtp_draft_window}, + {"mtp_drafted", snap.requests.mtp_drafted}, + {"mtp_accepted", snap.requests.mtp_accepted}, + {"mtp_fallback_steps", snap.requests.mtp_fallback_steps}, + {"mtp_rounds", snap.requests.mtp_rounds}, + {"mtp_accepted_per_position", snap.requests.mtp_accepted_per_position}, + {"mtp_last_accept_rate", snap.requests.mtp_last_accept_rate}}; + nlohmann::json health = {{"status", snap.health_status}, {"body", snap.health_body}}; + child_.observe_health(snap.health_status); + collector_.note_engine_state(state_name(st.state), st.last_event); + const std::string log = child_.log_tail(16 * 1024); + std::string cap = extract_kv_capacity_line(log); + if (cap.empty()) { cap = snap.engine_capacity_line; } + nlohmann::json insights = collector_.insights_report(); + return {{"monitor_only", !manages_engine_process(cfg_)}, + {"engine", std::move(engine)}, + {"dxgi", std::move(dxgi)}, + {"nvidia_smi", std::move(nvidia)}, + {"engine_capacity_line", cap}, + {"admin_vram", snap.admin_vram}, + {"admin_vram_note", snap.admin_vram_note}, + {"vram_control", collector_.vram_control_json()}, + {"requests", std::move(req)}, + {"insights", insights}, + {"series", collector_.series_json()}, + {"health", std::move(health)}, + {"log_tail", log}}; +} + +void DashboardServer::stop() { + stop_ = true; + if (server_ != nullptr) { static_cast(server_)->stop(); } +} + +void DashboardServer::run() { + httplib::Server svr; + server_ = &svr; + // No Access-Control-Allow-Origin. The custom mutating header is a CSRF + // brake only because cross-origin preflight then fails closed. + svr.set_pre_routing_handler([this](const httplib::Request& req, httplib::Response& res) { + if (!host_header_allowed(req.get_header_value("Host"), cfg_.port, cfg_.host, cfg_.bind_any)) { + res.status = 403; + res.set_content(nlohmann::json{{"error", "host not allowed"}}.dump(), "application/json"); + return httplib::Server::HandlerResponse::Handled; + } + return httplib::Server::HandlerResponse::Unhandled; + }); + svr.Get("/", [](const httplib::Request&, httplib::Response& res) { + res.set_content(std::string(kDashboardHtml), "text/html; charset=utf-8"); + }); + svr.Get("/api/state", [this](const httplib::Request&, httplib::Response& res) { + res.set_content(state_json().dump(), "application/json"); + }); + svr.Get("/api/insights", [this](const httplib::Request&, httplib::Response& res) { + (void)collector_.snapshot(); + res.set_content(collector_.insights_report().dump(), "application/json"); + }); + svr.Get("/api/events", [this](const httplib::Request&, httplib::Response& res) { + res.set_header("Cache-Control", "no-cache"); + res.set_header("Connection", "keep-alive"); + res.set_chunked_content_provider("text/event-stream", [this](std::size_t, httplib::DataSink& sink) { + if (stop_.load()) { + sink.done(); + return false; + } + const std::string payload = "data: " + state_json().dump() + "\n\n"; + sink.write(payload.data(), payload.size()); + for (int i = 0; i < 10 && !stop_.load(); ++i) { + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + return !stop_.load(); + }); + }); + auto control = [this](const httplib::Request& req, httplib::Response& res, auto fn) { + if (!control_allowed(req.remote_addr)) { + res.status = 403; + res.set_content(nlohmann::json{{"error", "control is loopback-only"}}.dump(), + "application/json"); + return; + } + if (!supervisor_control_header_ok(req.get_header_value(std::string(kSupervisorControlHeader)))) { + res.status = 403; + res.set_content(nlohmann::json{{"error", "missing X-NInfer-Supervisor header"}}.dump(), + "application/json"); + return; + } + if (!manages_engine_process(cfg_)) { + res.status = 409; + res.set_content(nlohmann::json{{"error", "engine is unmanaged"}}.dump(), + "application/json"); + return; + } + fn(); + res.set_content(state_json().dump(), "application/json"); + }; + svr.Post("/api/start", [this, control](const httplib::Request& req, httplib::Response& res) { + control(req, res, [this] { child_.start(); }); + }); + svr.Post("/api/stop", [this, control](const httplib::Request& req, httplib::Response& res) { + control(req, res, [this] { child_.stop(); }); + }); + svr.Post("/api/restart", [this, control](const httplib::Request& req, httplib::Response& res) { + control(req, res, [this] { child_.restart(); }); + }); + + const std::string host = cfg_.bind_any ? "0.0.0.0" : cfg_.host; + if (cfg_.bind_any || !is_loopback_host(host)) { + std::cerr << "WARNING: ninfer-supervisor is binding " << host + << " — control endpoints start/stop the engine. Loopback is the default.\n"; + } + svr.listen(host, cfg_.port); + server_ = nullptr; +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/server.hpp b/apps/ninfer-supervisor/server.hpp new file mode 100644 index 0000000000..065d142bd9 --- /dev/null +++ b/apps/ninfer-supervisor/server.hpp @@ -0,0 +1,33 @@ +#pragma once + +#include "collector.hpp" +#include "dashboard.hpp" +#include "engine_child.hpp" + +#include +#include +#include +#include + +namespace ninfer::supervisor { + +class DashboardServer { +public: + DashboardServer(SupervisorConfig cfg, EngineChild& child, Collector& collector); + ~DashboardServer(); + + void run(); + void stop(); + +private: + nlohmann::json state_json(); + bool control_allowed(const std::string& remote) const; + + SupervisorConfig cfg_; + EngineChild& child_; + Collector& collector_; + std::atomic stop_{false}; + void* server_ = nullptr; // httplib::Server* +}; + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/supervisor.example.json b/apps/ninfer-supervisor/supervisor.example.json new file mode 100644 index 0000000000..89ebbd04d2 --- /dev/null +++ b/apps/ninfer-supervisor/supervisor.example.json @@ -0,0 +1,31 @@ +{ + "engine": { + "executable": "P:/NInfer/build-win/apps/ninfer-serve.exe", + "args": [ + "P:/models/qwen3_8_27b_nvfp4.ninfer", + "--host", "127.0.0.1", + "--port", "8010" + ], + "workdir": "P:/NInfer", + "api_key_file": "P:/models/ninfer-api-key.txt", + "engine_host": "127.0.0.1", + "engine_port": 8010, + "request_log": "", + "device": 0, + "unmanaged": false + }, + "supervisor": { + "host": "127.0.0.1", + "port": 8099, + "bind_any": false, + "monitor_only": false, + "logs_dir": "P:/NInfer/supervisor-logs", + "run_at_login": false, + "restart": { + "max_backoff_s": 60, + "crash_loop_window_s": 60, + "crash_loop_max": 5, + "health_fail_threshold": 3 + } + } +} diff --git a/apps/ninfer-supervisor/supervisor.monitor-only.example.json b/apps/ninfer-supervisor/supervisor.monitor-only.example.json new file mode 100644 index 0000000000..efdc1fb07a --- /dev/null +++ b/apps/ninfer-supervisor/supervisor.monitor-only.example.json @@ -0,0 +1,17 @@ +{ + "engine": { + "unmanaged": true, + "engine_host": "127.0.0.1", + "engine_port": 8010, + "api_key_file": "P:/models/ninfer-api-key.txt", + "request_log": "P:/NInfer/supervisor-logs/prod.jsonl", + "device": 0 + }, + "supervisor": { + "host": "127.0.0.1", + "port": 8099, + "bind_any": false, + "monitor_only": true, + "logs_dir": "P:/NInfer/supervisor-logs" + } +} diff --git a/apps/ninfer-supervisor/tray.cpp b/apps/ninfer-supervisor/tray.cpp new file mode 100644 index 0000000000..504b67c5fd --- /dev/null +++ b/apps/ninfer-supervisor/tray.cpp @@ -0,0 +1,287 @@ +#include "tray.hpp" + +#include +#include + +#include + +namespace ninfer::supervisor { +namespace { + +constexpr UINT kTrayMsg = WM_APP + 1; +constexpr UINT kIdOpen = 1; +constexpr UINT kIdStart = 2; +constexpr UINT kIdStop = 3; +constexpr UINT kIdRestart = 4; +constexpr UINT kIdQuit = 5; +constexpr UINT_PTR kTimer = 1; +constexpr UINT kTrayUid = 1; +constexpr wchar_t kClass[] = L"NInferSupervisorTray"; + +COLORREF status_fill(TrayStatus status) { + switch (status) { + case TrayStatus::Working: return RGB(0x2F, 0x9E, 0x54); + case TrayStatus::Pending: return RGB(0xD1, 0x8B, 0x12); + case TrayStatus::Failed: return RGB(0xC5, 0x3B, 0x33); + case TrayStatus::Idle: break; + } + return RGB(0x6B, 0x72, 0x80); +} + +const wchar_t* status_word(TrayStatus status) { + switch (status) { + case TrayStatus::Working: return L"running"; + case TrayStatus::Pending: return L"pending"; + case TrayStatus::Failed: return L"failed"; + case TrayStatus::Idle: break; + } + return L"idle"; +} + +int small_icon_size() { + const int size = GetSystemMetrics(SM_CXSMICON); + return size > 0 ? size : 16; +} + +// Drawn rather than shipped as an .ico resource: no build-system change, no +// binary asset in the tree, correct at whatever SM_CXSMICON the display scaling +// reports, and the fill colour can carry status. +// +// The colour bitmap is an explicit 24bpp DIB section, NOT CreateCompatibleBitmap. +// A screen-compatible DDB is 32bpp on any modern display, and CreateIconIndirect +// then reads its alpha channel as per-pixel alpha. GDI never writes alpha, so +// every pixel would come out fully transparent and the icon would vanish. At +// 24bpp there is no alpha channel to misread and the 1bpp mask decides shape. +HICON make_tray_icon(TrayStatus status, int size) { + HDC screen = GetDC(nullptr); + if (screen == nullptr) { return nullptr; } + + BITMAPINFO bi{}; + bi.bmiHeader.biSize = sizeof(BITMAPINFOHEADER); + bi.bmiHeader.biWidth = size; + bi.bmiHeader.biHeight = -size; // top-down + bi.bmiHeader.biPlanes = 1; + bi.bmiHeader.biBitCount = 24; + bi.bmiHeader.biCompression = BI_RGB; + + void* bits = nullptr; + HBITMAP color_bmp = CreateDIBSection(screen, &bi, DIB_RGB_COLORS, &bits, nullptr, 0); + HBITMAP mask_bmp = CreateBitmap(size, size, 1, 1, nullptr); + HDC color_dc = CreateCompatibleDC(screen); + HDC mask_dc = CreateCompatibleDC(screen); + if (color_bmp == nullptr || mask_bmp == nullptr || color_dc == nullptr || mask_dc == nullptr) { + if (color_dc != nullptr) { DeleteDC(color_dc); } + if (mask_dc != nullptr) { DeleteDC(mask_dc); } + if (color_bmp != nullptr) { DeleteObject(color_bmp); } + if (mask_bmp != nullptr) { DeleteObject(mask_bmp); } + ReleaseDC(nullptr, screen); + return nullptr; + } + auto* old_color = static_cast(SelectObject(color_dc, color_bmp)); + auto* old_mask = static_cast(SelectObject(mask_dc, mask_bmp)); + + // Full-bleed fill; the rounded corners are cut by the mask, so the glyph + // antialiases against the fill rather than fringing against a background. + RECT rc{0, 0, size, size}; + HBRUSH fill = CreateSolidBrush(status_fill(status)); + FillRect(color_dc, &rc, fill); + DeleteObject(fill); + + LOGFONTW lf{}; + lf.lfHeight = -(size * 3 / 4); + lf.lfWeight = FW_BOLD; + lf.lfQuality = ANTIALIASED_QUALITY; + lf.lfCharSet = DEFAULT_CHARSET; + lstrcpyW(lf.lfFaceName, L"Segoe UI"); + HFONT font = CreateFontIndirectW(&lf); + if (font != nullptr) { + auto* old_font = static_cast(SelectObject(color_dc, font)); + SetBkMode(color_dc, TRANSPARENT); + SetTextColor(color_dc, RGB(0xFF, 0xFF, 0xFF)); + DrawTextW(color_dc, L"N", 1, &rc, DT_CENTER | DT_VCENTER | DT_SINGLELINE | DT_NOCLIP); + SelectObject(color_dc, old_font); + DeleteObject(font); + } + + // 1bpp mask: white (1) transparent, black (0) opaque. + PatBlt(mask_dc, 0, 0, size, size, WHITENESS); + HPEN pen = CreatePen(PS_SOLID, 1, RGB(0, 0, 0)); + auto* old_brush = static_cast(SelectObject(mask_dc, GetStockObject(BLACK_BRUSH))); + auto* old_pen = static_cast(SelectObject(mask_dc, pen)); + const int radius = size / 3 < 2 ? 2 : size / 3; + RoundRect(mask_dc, 0, 0, size, size, radius, radius); + SelectObject(mask_dc, old_brush); + SelectObject(mask_dc, old_pen); + DeleteObject(pen); + + SelectObject(color_dc, old_color); + SelectObject(mask_dc, old_mask); + + ICONINFO info{}; + info.fIcon = TRUE; + info.hbmMask = mask_bmp; + info.hbmColor = color_bmp; + HICON icon = CreateIconIndirect(&info); + + DeleteObject(color_bmp); + DeleteObject(mask_bmp); + DeleteDC(color_dc); + DeleteDC(mask_dc); + ReleaseDC(nullptr, screen); + return icon; +} + +LRESULT CALLBACK tray_wnd(HWND hwnd, UINT msg, WPARAM wparam, LPARAM lparam) { + TrayIcon* self = nullptr; + if (msg == WM_NCCREATE) { + auto* cs = reinterpret_cast(lparam); + self = static_cast(cs->lpCreateParams); + SetWindowLongPtrW(hwnd, GWLP_USERDATA, reinterpret_cast(self)); + } else { + self = reinterpret_cast(GetWindowLongPtrW(hwnd, GWLP_USERDATA)); + } + if (self == nullptr) { return DefWindowProcW(hwnd, msg, wparam, lparam); } + if (msg == WM_TIMER && wparam == kTimer) { + self->refresh_icon(); + return 0; + } + if (msg == kTrayMsg && (LOWORD(lparam) == WM_RBUTTONUP || LOWORD(lparam) == WM_LBUTTONUP)) { + POINT pt{}; + GetCursorPos(&pt); + HMENU menu = CreatePopupMenu(); + AppendMenuW(menu, MF_STRING, kIdOpen, L"Open dashboard"); + AppendMenuW(menu, MF_STRING, kIdStart, L"Start engine"); + AppendMenuW(menu, MF_STRING, kIdStop, L"Stop engine"); + AppendMenuW(menu, MF_STRING, kIdRestart, L"Restart engine"); + AppendMenuW(menu, MF_SEPARATOR, 0, nullptr); + AppendMenuW(menu, MF_STRING, kIdQuit, L"Quit supervisor"); + SetForegroundWindow(hwnd); + const int cmd = + TrackPopupMenu(menu, TPM_RETURNCMD | TPM_NONOTIFY, pt.x, pt.y, 0, hwnd, nullptr); + DestroyMenu(menu); + if (cmd == kIdOpen) { self->open_dashboard(); } + if (cmd == kIdStart) { self->child().start(); } + if (cmd == kIdStop) { self->child().stop(); } + if (cmd == kIdRestart) { self->child().restart(); } + if (cmd == kIdQuit) { PostQuitMessage(0); } + return 0; + } + if (msg == WM_DESTROY) { + PostQuitMessage(0); + return 0; + } + return DefWindowProcW(hwnd, msg, wparam, lparam); +} + +} // namespace + +TrayIcon::TrayIcon(EngineChild& child, std::string dashboard_url, bool manages_engine) + : child_(child), dashboard_url_(std::move(dashboard_url)), manages_engine_(manages_engine) {} + +TrayIcon::~TrayIcon() { + if (hwnd_ != nullptr) { DestroyWindow(static_cast(hwnd_)); } + if (hicon_ != nullptr) { DestroyIcon(static_cast(hicon_)); } +} + +EngineChild& TrayIcon::child() { return child_; } + +void TrayIcon::open_dashboard() const { + ShellExecuteA(nullptr, "open", dashboard_url_.c_str(), nullptr, nullptr, SW_SHOWNORMAL); +} + +void TrayIcon::request_quit() { + if (hwnd_ != nullptr) { PostMessageW(static_cast(hwnd_), WM_CLOSE, 0, 0); } +} + +TrayStatus TrayIcon::current_status() { + const EngineStatus status = child_.status(); + if (!manages_engine_) { + // Unmanaged: no process of ours has a state, so observed health is the + // only thing actually measured. EngineChild keeps it current for + // unmanaged engines through observe_health(). + if (status.health == "ok") { return TrayStatus::Working; } + if (status.health == "unhealthy") { return TrayStatus::Failed; } + if (status.health == "unreachable") { return TrayStatus::Pending; } + return TrayStatus::Idle; + } + if (status.crash_loop_halted) { return TrayStatus::Failed; } + switch (status.state) { + case EngineState::Running: + return status.health == "unhealthy" ? TrayStatus::Failed : TrayStatus::Working; + case EngineState::Starting: + case EngineState::Stopping: + case EngineState::BackingOff: return TrayStatus::Pending; + case EngineState::Halted: return TrayStatus::Failed; + case EngineState::Stopped: break; + } + return TrayStatus::Idle; +} + +std::wstring TrayIcon::tooltip(TrayStatus status) const { + std::wstring tip = L"NInfer supervisor - "; + tip += status_word(status); + if (!manages_engine_) { tip += L" (monitor-only)"; } + return tip; +} + +void TrayIcon::refresh_icon() { + if (hwnd_ == nullptr) { return; } + const TrayStatus status = current_status(); + if (static_cast(status) == last_status_) { return; } + last_status_ = static_cast(status); + + HICON icon = make_tray_icon(status, small_icon_size()); + if (icon == nullptr) { return; } + + const std::wstring tip = tooltip(status); + NOTIFYICONDATAW nid{}; + nid.cbSize = sizeof(nid); + nid.hWnd = static_cast(hwnd_); + nid.uID = kTrayUid; + nid.uFlags = NIF_ICON | NIF_TIP; + nid.hIcon = icon; + lstrcpynW(nid.szTip, tip.c_str(), ARRAYSIZE(nid.szTip)); + Shell_NotifyIconW(NIM_MODIFY, &nid); + + if (hicon_ != nullptr) { DestroyIcon(static_cast(hicon_)); } + hicon_ = icon; +} + +void TrayIcon::run() { + WNDCLASSEXW wc{}; + wc.cbSize = sizeof(wc); + wc.lpfnWndProc = tray_wnd; + wc.hInstance = GetModuleHandleW(nullptr); + wc.lpszClassName = kClass; + RegisterClassExW(&wc); + HWND hwnd = CreateWindowExW(0, kClass, L"NInfer supervisor", 0, 0, 0, 0, 0, HWND_MESSAGE, + nullptr, wc.hInstance, this); + hwnd_ = hwnd; + + const TrayStatus status = current_status(); + last_status_ = static_cast(status); + HICON icon = make_tray_icon(status, small_icon_size()); + hicon_ = icon; + const std::wstring tip = tooltip(status); + + NOTIFYICONDATAW nid{}; + nid.cbSize = sizeof(nid); + nid.hWnd = hwnd; + nid.uID = kTrayUid; + nid.uFlags = NIF_MESSAGE | NIF_TIP | NIF_ICON; + nid.uCallbackMessage = kTrayMsg; + nid.hIcon = icon != nullptr ? icon : LoadIconW(nullptr, MAKEINTRESOURCEW(32512)); + lstrcpynW(nid.szTip, tip.c_str(), ARRAYSIZE(nid.szTip)); + Shell_NotifyIconW(NIM_ADD, &nid); + SetTimer(hwnd, kTimer, 1000, nullptr); + + MSG msg; + while (GetMessageW(&msg, nullptr, 0, 0) > 0) { + TranslateMessage(&msg); + DispatchMessageW(&msg); + } + KillTimer(hwnd, kTimer); + Shell_NotifyIconW(NIM_DELETE, &nid); +} + +} // namespace ninfer::supervisor diff --git a/apps/ninfer-supervisor/tray.hpp b/apps/ninfer-supervisor/tray.hpp new file mode 100644 index 0000000000..97e43b505c --- /dev/null +++ b/apps/ninfer-supervisor/tray.hpp @@ -0,0 +1,48 @@ +#pragma once + +#include "engine_child.hpp" + +#include +#include + +namespace ninfer::supervisor { + +// What the tray icon says at a glance. In monitor-only mode this still reflects +// the OBSERVED health of the engine rather than a neutral "not my process": +// EngineChild keeps st_.health current for unmanaged engines too, so the colour +// is measured, not invented. The tooltip carries the managed/observing +// distinction instead, because that belongs in words rather than in a hue. +enum class TrayStatus : std::uint8_t { + Idle, // grey — nothing running, or nothing known yet + Working, // green — engine running and answering /health + Pending, // amber — starting, stopping, backing off, or unreachable + Failed, // red — halted, crash-looped, or reporting unhealthy +}; + +class TrayIcon { +public: + TrayIcon(EngineChild& child, std::string dashboard_url, bool manages_engine); + ~TrayIcon(); + void run(); + void request_quit(); + void open_dashboard() const; + EngineChild& child(); + + // Repaints the tray icon when the status changes. Called on a timer; cheap + // because it reads EngineChild's in-memory status and never polls the + // engine or shells out to nvidia-smi. + void refresh_icon(); + +private: + TrayStatus current_status(); + std::wstring tooltip(TrayStatus status) const; + + EngineChild& child_; + std::string dashboard_url_; + bool manages_engine_ = true; + void* hwnd_ = nullptr; + void* hicon_ = nullptr; + int last_status_ = -1; +}; + +} // namespace ninfer::supervisor diff --git a/apps/serve/main.cpp b/apps/serve/main.cpp index 9263e100cd..48cd81bd23 100644 --- a/apps/serve/main.cpp +++ b/apps/serve/main.cpp @@ -8,21 +8,86 @@ #include #include #include +#include #include #include #include +#include #include +#include +#include +#include #include namespace { std::atomic g_server{nullptr}; +class BootWatchdog { +public: + explicit BootWatchdog(std::chrono::seconds timeout) + : timeout_(timeout), done_(std::make_shared>(false)) { + if (timeout_.count() <= 0) { return; } + auto done = done_; + std::thread([done, timeout = timeout_] { + const auto deadline = std::chrono::steady_clock::now() + timeout; + while (std::chrono::steady_clock::now() < deadline) { + if (done->load(std::memory_order_relaxed)) { return; } + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } + if (!done->load(std::memory_order_relaxed)) { + ninfer::serve::write_console_log( + ninfer::serve::ConsoleLogLevel::Error, + "boot watchdog timeout (" + std::to_string(timeout.count()) + + " s) exceeded before reaching listening state; terminating process"); + std::cerr.flush(); + std::_Exit(1); + } + }).detach(); + } + + void disarm() { + if (done_) { done_->store(true, std::memory_order_relaxed); } + } + + // No disarm on destruction: when warmup throws, this object unwinds before the + // service whose teardown can hang in device synchronization — the watchdog must + // stay armed through that unwind. A boot that fails fast exits the process (and + // the detached thread) before the deadline; the healthy path disarms explicitly + // once the server reaches the listening state. + ~BootWatchdog() = default; + +private: + std::chrono::seconds timeout_{0}; + std::shared_ptr> done_; +}; + void handle_signal(int) { ninfer::serve::HttpServer* server = g_server.load(); if (server != nullptr) { server->stop(); } } +// An exception that escapes a request boundary ends the process through +// std::terminate, and the default handler's message is the only record of which +// exception it was. That message is worth writing through the server's own log: +// under a container this process is pid 1, the kernel discards the SIGABRT that +// abort() raises against itself, glibc falls through to its abort instruction, +// and all the kernel reports is a bare protection fault inside libc. +[[noreturn]] void log_terminate() { + std::string detail = "terminate called with no active exception"; + if (std::current_exception() != nullptr) { + try { + std::rethrow_exception(std::current_exception()); + } catch (const std::exception& error) { + detail = std::string("terminate called after throwing ") + typeid(error).name() + ": " + + error.what(); + } catch (...) { detail = "terminate called after throwing a non-std exception"; } + } + ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Error, detail); + std::cerr.flush(); + std::abort(); +} + std::string format_bytes(std::size_t bytes) { constexpr double kMiB = 1024.0 * 1024.0; constexpr double kGiB = 1024.0 * kMiB; @@ -39,13 +104,21 @@ std::string format_bytes(std::size_t bytes) { } // namespace int main(int argc, char** argv) { + std::set_terminate(log_terminate); + ninfer::serve::ServeOptions options; try { - const ninfer::serve::ServeOptions options = ninfer::serve::parse_serve_options(argc, argv); - if (options.help_requested) { - std::cout << ninfer::serve::serve_usage_text(argv[0]); - return 0; - } + options = ninfer::serve::parse_serve_options(argc, argv); + } catch (const std::exception& exception) { + ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Error, exception.what()); + std::cerr << ninfer::serve::serve_usage_text(argv[0]); + return 1; + } + if (options.help_requested) { + std::cout << ninfer::serve::serve_usage_text(argv[0]); + return 0; + } + try { using Clock = std::chrono::steady_clock; ninfer::serve::HttpServer server(options); if (!server.bind()) { @@ -79,6 +152,7 @@ int main(int argc, char** argv) { << " tokens pages=" << memory.kv_capacity_page_groups << '/' << memory.kv_capacity_max_page_groups << " runtime=" << format_bytes(memory.runtime_reservation_bytes) + << " prefix-cache=" << format_bytes(memory.prefix_cache_bytes) << " free-after-weights=" << format_bytes(memory.available_after_weights_bytes) << " free-after-startup=" << format_bytes(memory.available_after_startup_bytes) << " headroom=" << format_bytes(memory.kv_capacity_headroom_bytes) @@ -93,6 +167,8 @@ int main(int argc, char** argv) { } ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Info, capacity.str()); + BootWatchdog watchdog(std::chrono::seconds(options.boot_watchdog_timeout_s)); + ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Info, "warming up..."); service.warmup(); @@ -106,18 +182,19 @@ int main(int argc, char** argv) { << ", auth: " << (options.api_key.empty() ? "disabled" : "bearer") << ')'; ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Info, listening.str()); + watchdog.disarm(); + const bool ok = server.listen(); g_server.store(nullptr); if (!ok) { ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Error, - "failed to bind " + options.host + ':' + + "accept loop failed on " + options.host + ':' + std::to_string(options.port)); return 1; } return 0; } catch (const std::exception& exception) { ninfer::serve::write_console_log(ninfer::serve::ConsoleLogLevel::Error, exception.what()); - std::cerr << ninfer::serve::serve_usage_text(argv[0]); return 1; } } diff --git a/bench/targets/qwen3_6_27b/ninfer_bench_support.cpp b/bench/targets/qwen3_6_27b/ninfer_bench_support.cpp index eb08e44d2c..11f265914d 100644 --- a/bench/targets/qwen3_6_27b/ninfer_bench_support.cpp +++ b/bench/targets/qwen3_6_27b/ninfer_bench_support.cpp @@ -666,6 +666,7 @@ std::string format_json(const BenchEnvironment& env, const std::string& command, << " \"planned_slack_bytes\": " << env.memory.planned_slack_bytes << ",\n" << " \"cuda_graph_allowance_bytes\": " << env.memory.cuda_graph_allowance_bytes << ",\n" << " \"cuda_graph_observed_bytes\": " << env.memory.cuda_graph_observed_bytes << ",\n" + << " \"prefix_cache_bytes\": " << env.memory.prefix_cache_bytes << ",\n" << " \"kv_payload_bytes\": " << env.memory.kv_payload_bytes << "\n" << " },\n" << " \"config\": {\n" diff --git a/docs/maintainer/paged-kv-cache.md b/docs/maintainer/paged-kv-cache.md index b2cea7d1ec..11f0925aed 100644 --- a/docs/maintainer/paged-kv-cache.md +++ b/docs/maintainer/paged-kv-cache.md @@ -1410,7 +1410,7 @@ contiguous-KV reference 只记录当时的 `B=1` paging migration，不是当前 - `EngineOptions.max_context=S` 是 per-sequence logical ceiling，`EngineOptions.kv_capacity` 是 `Explicit(K_main)` 或 `Automatic(R)`；令 `L=ceil(S/64)`、`M_min=max(L,max_concurrency)`、 `M_max=max_concurrency*L`，Explicit 取 `M=ceil(K_main/64)`，Automatic 根据完整 target physical - reservation curve 与权重加载后的空闲显存扣除 headroom `R` 后，直接求得区间内最大的 `M`； + reservation curve（含 prefix-seed arena）与权重加载后的空闲显存扣除 headroom `R` 后，直接求得区间内最大的 `M`； CLI/server 的 `R` 为 1 GiB；Main 与 DFlash Full 的 per-allocation logical capacity 均为 `L`、physical capacity 均为 `M` pages，MTP 的 logical capacity 为 `L`、physical capacity 为 `M + max_concurrency*ceil((K_draft-1)/64)` pages，其中 `K_draft` 是 diff --git a/docs/serving.md b/docs/serving.md index 3fe656cbf5..4328548d46 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -43,8 +43,11 @@ cannot be combined with `--vision`. A later request cannot enable a capability o | Method and path | Behavior | |---|---| | `GET /health` | process health | -| `GET /v1/models` | configured OpenAI model alias | -| `GET /v1/models/{id}` | lookup of the configured alias | +| `GET /admin/vram` | current VRAM tier state (seed/KV held bytes, range, last transition). Registered only with `--admin-vram`, which requires `--api-key` | +| `POST /admin/vram/release` | release named cache tiers (`{"tiers":["seed"],"target_mib":N}`). KV is refused in this phase. Same enable-gate as GET | +| `POST /admin/vram/reclaim` | re-acquire previously released cache tiers; failure stays degraded. Same enable-gate as GET | +| `GET /v1/models` | configured OpenAI model alias, including `max_model_len` = `--max-context` | +| `GET /v1/models/{id}` | lookup of the configured alias, same `max_model_len` | | `POST /v1/chat/completions` | OpenAI-style chat generation | | `POST /v1/responses` | OpenAI Responses Core generation, state, typed Items, and SSE | | `POST /v1/responses/input_tokens` | Responses prompt-token count without generation | @@ -80,7 +83,7 @@ The endpoint supports: - `stream_options.include_usage`; - function tools, tool choices, assistant tool-call history, and tool-result messages; - the top-level `reasoning_effort` field; -- the `enable_thinking` extension; +- top-level `enable_thinking` and `chat_template_kwargs.enable_thinking`; - `chat_template_kwargs.preserve_thinking` and the top-level `preserve_thinking` alias. The request `model` must equal the public model ID: the artifact `identity.model_id` by default, or @@ -102,13 +105,29 @@ not exposed by the loaded template returns HTTP 400 with code For Chat Completions, `reasoning_effort: "none"` disables thinking. `low`, `medium`, and `xhigh` select the corresponding template effort when available. The other OpenAI protocol values `minimal`, `high`, and `max` are parsed but rejected when the loaded template does not expose them. -`enable_thinking` controls the same new-turn thinking switch; a contradictory combination with -`reasoning_effort` returns `conflicting_template_option`. +`high` is not an alias of `xhigh`. Top-level `enable_thinking: false` and +`chat_template_kwargs.enable_thinking: false` disable the same new-turn thinking switch as +`reasoning_effort: "none"`. The two `enable_thinking` spellings must agree when both are present. A +contradictory combination of either spelling with `reasoning_effort` returns +`conflicting_template_option`. `preserve_thinking` controls whether reasoning from closed assistant turns remains in later -prompts. It defaults to the server setting, which is off unless `--preserve-thinking` is used. If -both OpenAI spellings are present they must carry the same boolean value. Unknown non-null -`chat_template_kwargs` are rejected. +prompts. It is independent of the new-turn thinking switch: a request may disable thinking on the +current turn while still preserving closed-turn reasoning, or the reverse. It defaults to the +server setting, which is off unless `--preserve-thinking` is used. If both OpenAI spellings are +present they must carry the same boolean value. Unknown non-null `chat_template_kwargs` return +HTTP 400 `chat_template_option_not_supported`; a misspelled `enable_thinking` key is rejected +rather than ignored, so thinking cannot remain on by default when a client intended to disable it. + +Chat Completions `usage` keeps the OpenAI totals (`prompt_tokens`, `completion_tokens`, +`total_tokens`) and adds `prompt_tokens_details.cached_tokens` as a subset of `prompt_tokens`, +never an addend. The same cached count is also emitted as `prefix_cache_hit_tokens`, and +`prefix_reuse_path` names the Engine path that served the prompt (`full_reset`, +`append_frontier`, `restore_turn_checkpoint`, `restore_response_checkpoint`, `seed_prefix`) — +the same strings as `request_done.result` in `--request-log-jsonl`. Anthropic Messages does not +emit these keys: its `input_tokens` already excludes cache reads, so reporting a cached subset +would change that field's meaning. Responses already reports `input_tokens_details.cached_tokens` +and adds the same two log-named fields. Streaming begins with an assistant-role chunk, sends separate reasoning and content deltas, then a finish-reason chunk and `[DONE]`. When `stream_options.include_usage` is true, a final empty @@ -211,7 +230,9 @@ wire response contains typed `output` Items. | `temperature` | finite number in `[0,2]` | | `top_p` | finite number in `[0,1]` | | `metadata` | at most 16 string pairs; keys at most 64 characters and values at most 512 | -| `reasoning.effort` | `none` disables thinking; `low`, `medium`, or `xhigh` selects an effort exposed by the loaded chat template; `minimal`, `high`, and `max` return `reasoning_effort_not_supported` for the registered templates | +| `reasoning.effort` | `none` disables thinking; `low`, `medium`, or `xhigh` selects an effort exposed by the loaded chat template; `minimal`, `high`, and `max` return `reasoning_effort_not_supported` for the registered templates. `high` is not mapped to `xhigh` | +| `enable_thinking` | optional boolean; `false` disables new-turn thinking the same way as `reasoning.effort: "none"` | +| `chat_template_kwargs.enable_thinking` | vLLM-dialect alias for the same option; conflicting values with top-level `enable_thinking` are rejected | | `chat_template_kwargs.preserve_thinking` | optional boolean controlling whether closed-turn reasoning remains in reconstructed prompts | | `preserve_thinking` | top-level alias for the same option; conflicting values are rejected | | `text.format` | omitted or `{"type":"text"}` only | @@ -226,7 +247,8 @@ wire response contains typed `output` Items. | `stream_options` | omitted or `{"include_obfuscation":false}` | Unknown top-level fields fail with `unknown_parameter`. Recognized but unsupported features fail -with a field-specific 400 error instead of being silently ignored. +with a field-specific 400 error instead of being silently ignored. Unknown non-null +`chat_template_kwargs` keys fail with `chat_template_option_not_supported`. ### Input Item contract @@ -480,6 +502,7 @@ curl http://127.0.0.1:8080/v1/models \ | `--no-prefix-reuse` | disable compatible-prefix caching | prefix reuse on | | `--no-thinking` | disable thinking by default | thinking on | | `--preserve-thinking` | preserve closed-turn assistant reasoning by default | off | +| `--tolerant-tool-calls` | recover complete Qwen tool calls with malformed wrapper/suffix output | off | | `--cors` | permissive browser CORS headers | off | | `--temperature F` | process-level temperature override | unset | | `--top-p F` | process-level top-p override | unset | @@ -577,17 +600,28 @@ network serialization run outside the GPU executor and do not delay formation of sequence's logical ceiling; the latter sizes the shared Main Text KV pool used by all active requests and retained prefixes. Both are represented with 64-token pages internally, while a sequence can never cross the exact `--max-context` frontier. `--kv-capacity N` requests an explicit -capacity; `--kv-capacity auto` chooses the largest legal capacity that fits the memory remaining +capacity (and means min==max==N). `--kv-capacity-min` / `--kv-capacity-max` set an elastic token +range and boot at max; this phase reports the range and does not yet shrink KV at runtime. +`--prefix-cache-mib N` is the same for the seed store (min==max==N). `--prefix-cache-mib-min` / +`--prefix-cache-mib-max` boot at max; min 0 is fully releasable via idle release or +`POST /admin/vram/release`. `--vram-idle-release-after-s N` (default 0, disabled) drops the seed +store after N seconds with no in-flight GPU work. `--vram-observe-only` logs would-be releases +without freeing memory. `--admin-vram` (default off) registers the `/admin/vram` routes and +requires `--api-key`; without both flags the mutation endpoints do not exist. Default flag-free +behaviour is unchanged: no idle release, no admin routes, seed store +fixed at `--prefix-cache-mib` (0 disables it). +`--kv-capacity auto` chooses the largest legal capacity that fits the memory remaining after weights are loaded while keeping 1 GiB of sizing headroom. When omitted it follows `--max-context`, preserving one full-length request's capacity. The shared pool is fixed at startup and is not divided evenly among request lanes. Automatic sizing evaluates the complete target runtime layout for the chosen concurrency, KV -dtype, speculative backend, draft window, Vision setting, workspace, and CUDA Graph allowance. It -uses a direct page-capacity calculation rather than allocation probing. Startup reports the policy, -resolved capacity, runtime reservation, free memory after weights, automatic headroom, planned -slack, actual free memory after complete startup, and observed Graph memory. An explicit capacity -is never silently reduced, and neither policy permits request-time pool growth. +dtype, speculative backend, draft window, Vision setting, workspace, CUDA Graph allowance, and +prefix-cache arena. It uses a direct page-capacity calculation rather than allocation probing. +Startup reports the policy, resolved capacity, runtime reservation, prefix-cache arena, free +memory after weights, automatic headroom, planned slack, actual free memory after complete +startup, and observed Graph memory. An explicit capacity is never silently reduced, and neither +policy permits request-time pool growth. Admission reserves the full prompt-plus-effective-output page entitlement, so an admitted request can finish within its declared bound. A later request waits in FIFO order when the remaining shared @@ -634,8 +668,25 @@ context-capacity finishes map to `length`/ `max_tokens`; ordinary model or strin `stop`/ `end_turn`. Function tools are rendered into the model prompt and generated calls are parsed into protocol -responses. NInfer does not execute tools and does not enforce client JSON Schema through constrained -decoding. +responses. Inside an explicit ``/`` wrapper the argument payload may be +either `` blocks or a single JSON object; JSON-object arguments are typed with the +same declared-schema rules as parameter blocks. Forms that only resemble tool syntax +(`[tool_use:…]`, bare `tool_call:`/`arguments:`, or a `` without ``) are not +parsed as calls. A successful parse leaves assistant `content` empty (OpenAI `content: null` when +`tool_calls` are present); a preamble before the wrapper is not returned as user-visible text. +NInfer does not execute tools and does not validate tool arguments against the full +client JSON Schema through constrained decoding; that remains the client's responsibility. When +parsing a generated call, NInfer does consult the top-level parameter `"type"` declared in each +tool's schema to decide whether a parameter value that is valid JSON may be deserialized into the +corresponding JSON type (number, boolean, array, object, null): only parameters whose declared +type(s) are all valid non-string JSON Schema types are deserialized. Parameters typed as `"string"` +(or declared via a type array that includes `"string"`), parameters with an unknown or misspelled +`"type"`, and parameters absent from the schema preserve the model's raw text so the string +contract reaches the client intact. Full JSON Schema validation (constraints, required sets, +formats, nested keywords) is not performed server-side and remains the client's job. +Duplicate tool names within a single request are rejected with a 400 on all three +protocol surfaces (OpenAI Chat Completions, OpenAI Responses, Anthropic Messages), +keeping the per-tool parameter type map unambiguous. Prompt-token usage includes chat-template and expanded media tokens. Generated-token usage comes from accepted output token IDs, including a stop token whose decoded text may be withheld. diff --git a/include/ninfer/engine.h b/include/ninfer/engine.h index c44c7592b3..26c4d3f098 100644 --- a/include/ninfer/engine.h +++ b/include/ninfer/engine.h @@ -4,6 +4,8 @@ #include #include +#include +#include namespace ninfer { @@ -92,8 +94,14 @@ class Engine { [[nodiscard]] MemorySummary memory_summary() const; [[nodiscard]] RuntimeStats runtime_stats() const; [[nodiscard]] MediaCacheSummary media_cache_summary() const; + [[nodiscard]] bool is_healthy() const noexcept; void reset_memory_peaks() noexcept; + [[nodiscard]] VramControlState vram_control_state() const; + // tiers: "seed" is implemented. "kv" is refused. target_mib is optional (0 = floor/min). + void vram_release(const std::vector& tiers, std::size_t target_mib = 0); + void vram_reclaim(); + private: class Impl; std::shared_ptr impl_; diff --git a/include/ninfer/ops/gqa_attention.h b/include/ninfer/ops/gqa_attention.h index 54bf16796a..90a04bd9b1 100644 --- a/include/ninfer/ops/gqa_attention.h +++ b/include/ninfer/ops/gqa_attention.h @@ -90,7 +90,8 @@ gqa_attention_workspace_capacity_bytes(std::int32_t q_heads, DType cache_dtype, void gqa_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& valid_columns, const Tensor& kv_table_rows, float scale, PagedKVBatchLayerView cache, GqaExecutionEnvelope envelope, - WorkspaceArena& workspace, Tensor& out, cudaStream_t stream); + WorkspaceArena& workspace, Tensor& out, cudaStream_t stream, + const Tensor* gate = nullptr); /** * A2: perform only the cache-write part of A1. k/v are contiguous BF16 `[256,4|2,T]`, positions is @@ -107,7 +108,9 @@ void gqa_kv_append(const Tensor& k, const Tensor& v, const Tensor& positions, * to A1. Caller workspace is reported by gqa_attention_workspace_capacity_bytes(). */ void gqa_attention_cached(const Tensor& q, const Tensor& positions, float scale, + /* optional fused sigmoid gate: see gqa_attention */ const PagedKVLayerView& cache, GqaExecutionEnvelope envelope, - WorkspaceArena& workspace, Tensor& out, cudaStream_t stream); + WorkspaceArena& workspace, Tensor& out, cudaStream_t stream, + const Tensor* gate = nullptr); } // namespace ninfer::ops diff --git a/include/ninfer/ops/rope.h b/include/ninfer/ops/rope.h index d308991f1e..522ce4b25a 100644 --- a/include/ninfer/ops/rope.h +++ b/include/ninfer/ops/rope.h @@ -40,4 +40,13 @@ void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& q, Tenso // from x; Q versus K role does not change the transformation. void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& x, cudaStream_t stream); +/** + * Fused q/k RMS-norm (unit-offset weights) + rotary. For the Text D=256 rotary-64 + * geometries (16Q/2K and 24Q/4K) this runs one kernel; any other geometry falls + * back to rmsnorm + rmsnorm + rope with identical results. + */ +void qk_norm_rope(const Tensor& positions, int rotary_dim, float theta, const Tensor& q_in, + const Tensor& q_weight, Tensor& q_out, const Tensor& k_in, + const Tensor& k_weight, Tensor& k_out, float eps, cudaStream_t stream); + } // namespace ninfer::ops diff --git a/include/ninfer/ops/sparse_moe.h b/include/ninfer/ops/sparse_moe.h index 73cc05b89c..d8ab1e6bd0 100644 --- a/include/ninfer/ops/sparse_moe.h +++ b/include/ninfer/ops/sparse_moe.h @@ -22,6 +22,11 @@ enum class SparseMoeEpilogue : std::uint8_t { AddResidual, }; +struct WeightPrefetchSpan { + const void* data = nullptr; + std::size_t bytes = 0; +}; + /** * Returns the transient capacity required by SparseMoe for every T in the inclusive * [min_tokens,max_tokens] interval. The routed QTypes are the fixed implementation profile. @@ -59,7 +64,15 @@ enum class SparseMoeEpilogue : std::uint8_t { * Execution is enqueued on stream without host synchronization. Workspace is caller-owned, * graph-stable transient storage and carries no state beyond the call. */ +/** + * `next_prefetch` names the weight payload the next decode-step consumer will stream. The D4 + * epilogue, whose tail runs on an otherwise idle bus, issues fire-and-forget L2 prefetches for it, + * clamped to the op's own cap. Purely a cache hint: no value is read through it and the emitted + * tokens are unaffected. An empty span disables the hint, which is what every non-decode route + * passes. + */ void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, - Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream); + Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream, + WeightPrefetchSpan next_prefetch = {}); } // namespace ninfer::ops diff --git a/include/ninfer/types.h b/include/ninfer/types.h index c074a3e04a..419cf5d2ea 100644 --- a/include/ninfer/types.h +++ b/include/ninfer/types.h @@ -89,7 +89,39 @@ struct EngineOptions { std::uint32_t media_preprocess_threads = 0; bool enable_vision = false; bool use_cuda_graph = true; + // Device bytes reserved at startup for the cross-request prefix-seed store; 0 disables it. + std::size_t prefix_cache_bytes = 0; + // Elastic seed-store range. Boot allocates prefix_cache_bytes (the max). Min 0 is fully + // releasable. When both extra fields are 0 they follow prefix_cache_bytes (fixed size). + std::size_t prefix_cache_min_bytes = 0; + std::size_t prefix_cache_max_bytes = 0; + // Elastic KV range in tokens. Boot uses kv_capacity (the max). 0 follows the resolved policy. + std::uint32_t kv_capacity_min_tokens = 0; + std::uint32_t kv_capacity_max_tokens = 0; + std::uint32_t vram_guarantee_context = 0; // 0 = max_context + std::uint32_t vram_guarantee_concurrency = 1; + std::size_t vram_floor_bytes = 0; // 0 = derive from the capability guarantee + std::uint32_t vram_idle_release_after_s = 0; // 0 disables idle release + bool vram_observe_only = false; LoadProgress load_progress; + std::function on_fatal_error; +}; + +struct VramTierState { + std::string name; + std::size_t held_bytes = 0; + std::size_t min_bytes = 0; + std::size_t max_bytes = 0; + std::size_t reclaimable_bytes = 0; + bool released = false; +}; + +struct VramControlState { + std::vector tiers; + std::size_t floor_bytes = 0; + bool observe_only = false; + std::string last_transition; + std::string last_reason; }; enum class SamplingMode : std::uint8_t { @@ -388,6 +420,7 @@ enum class PrefixReusePath : std::uint8_t { AppendAtFrontier, RestoreTurnCheckpoint, RestoreResponseCheckpoint, + SeedPrefixCache, }; struct GenerationResult { @@ -431,6 +464,8 @@ struct MemorySummary { std::size_t workspace_logical_peak_bytes = 0; std::size_t cuda_graph_allowance_bytes = 0; std::size_t cuda_graph_observed_bytes = 0; + std::size_t prefix_cache_bytes = 0; + std::size_t prefix_cache_held_bytes = 0; std::size_t kv_payload_bytes = 0; }; diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index f5590f3f77..e4d070143e 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -61,6 +61,7 @@ target_link_libraries(ninfer_nvfp4_tma PRIVATE ninfer_core CUDA::cudart CUDA::cu # Shared mathematical Ops. The list is intentionally explicit: adding a # source is a build-boundary decision, not an accidental recursive-glob side effect. add_library(ninfer_ops STATIC + ops/common/device_info.cu ops/launcher/add_bias.cu ops/launcher/argmax.cu ops/launcher/cast.cu @@ -269,18 +270,37 @@ add_library(ninfer_text STATIC text/unicode.cpp ${PROJECT_SOURCE_DIR}/third_party/utf8proc/utf8proc.c) ninfer_internal_includes(ninfer_text) +if(WIN32) + # utf8proc.h marks its API __declspec(dllimport) unless UTF8PROC_STATIC is set; + # compiling the source itself with dllimport is an error (C2491). + target_compile_definitions(ninfer_text PRIVATE UTF8PROC_STATIC) +endif() -add_library(ninfer_media_decode STATIC - media/decode/decode.cpp) +if(NINFER_BUILD_MEDIA) + add_library(ninfer_media_decode STATIC + media/decode/decode.cpp) + target_link_libraries(ninfer_media_decode PRIVATE PkgConfig::FFMPEG) +else() + # API-compatible stub: keeps the vision frontend compiling without FFMPEG. + add_library(ninfer_media_decode STATIC + media/decode/decode_stub.cpp) +endif() ninfer_internal_includes(ninfer_media_decode) -target_link_libraries(ninfer_media_decode PRIVATE PkgConfig::FFMPEG) if(NINFER_BUILD_MEDIA_ACQUIRE) # Product-only path/data/HTTP acquisition. No target package links this library. - add_library(ninfer_media_acquire STATIC - product/media_acquire/acquire.cpp) + if(NINFER_BUILD_MEDIA) + add_library(ninfer_media_acquire STATIC + product/media_acquire/acquire.cpp) + target_link_libraries(ninfer_media_acquire PRIVATE PkgConfig::LIBCURL) + if(WIN32) + target_link_libraries(ninfer_media_acquire PRIVATE ws2_32) + endif() + else() + add_library(ninfer_media_acquire STATIC + product/media_acquire/acquire_stub.cpp) + endif() ninfer_internal_includes(ninfer_media_acquire) - target_link_libraries(ninfer_media_acquire PRIVATE PkgConfig::LIBCURL) endif() if(NINFER_BUILD_PROMPT_INPUT) diff --git a/src/artifact/materializer.cpp b/src/artifact/materializer.cpp index 2df1305d13..88ed5a036d 100644 --- a/src/artifact/materializer.cpp +++ b/src/artifact/materializer.cpp @@ -1,8 +1,11 @@ #include "artifact/materializer.h" +#include "core/verbose.h" + #include #include +#include #include #include #include @@ -248,6 +251,36 @@ MaterializedArtifact materialize(const Reader& reader, const MaterializationPlan if (copied != total || next_range != ranges.size()) { throw ArtifactError("direct materialization did not cover every tensor byte"); } + if (ninfer::verbose_enabled()) { + for (const DeviceMaterialization& placement : plan.device_objects) { + if (placement.bytes > (1ULL << 20)) { continue; } + const ObjectHandle handle = placement.object; + const ObjectDescriptor& desc = reader.objects().at(handle.index); + const PayloadSpan payload = reader.payload(desc); + std::byte* dev = static_cast(out.objects_.at(handle.index).device); + const std::size_t n = + static_cast(std::min(64, placement.bytes)); + std::vector host(n); + (void)cudaMemcpy(host.data(), dev, n, cudaMemcpyDeviceToHost); + bool match = true; + for (std::size_t i = 0; i < n; ++i) { + if (host[i] != payload.data[i]) { match = false; break; } + } + std::fprintf(stderr, + "[verbose] materialize check: %s bytes=%llu dev=%p file_off=%llu match=%s\n", + std::string(object_name(desc)).c_str(), + (unsigned long long)placement.bytes, (void*)dev, + (unsigned long long)payload.absolute_offset, match ? "YES" : "NO"); + if (!match) { + std::fprintf(stderr, "[verbose] dev = "); + for (std::size_t i = 0; i < n; ++i) { std::fprintf(stderr, "%02x", (unsigned)host[i]); } + std::fprintf(stderr, "\n[verbose] file = "); + for (std::size_t i = 0; i < n; ++i) { std::fprintf(stderr, "%02x", (unsigned)payload.data[i]); } + std::fprintf(stderr, "\n"); + } + } + std::fflush(stderr); + } out.stats_.h2d_bytes = copied; out.stats_.upload_seconds = std::chrono::duration(std::chrono::steady_clock::now() - start).count(); diff --git a/src/artifact/reader.cpp b/src/artifact/reader.cpp index 1dc3afd1ea..ab2949ed52 100644 --- a/src/artifact/reader.cpp +++ b/src/artifact/reader.cpp @@ -1,10 +1,13 @@ #include "artifact/reader.h" +#include "core/verbose.h" + #include #include #include #include +#include #include #include #include @@ -15,10 +18,14 @@ #include #include +#if defined(_WIN32) +#include +#else #include #include #include #include +#endif namespace ninfer::artifact { namespace { @@ -181,6 +188,46 @@ struct TransparentStringHash { class MappedFile { public: explicit MappedFile(const std::filesystem::path& path) { +#if defined(_WIN32) + HANDLE handle = ::CreateFileW(path.wstring().c_str(), GENERIC_READ, FILE_SHARE_READ, + nullptr, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, nullptr); + if (handle == INVALID_HANDLE_VALUE) { + throw std::system_error(::GetLastError(), std::generic_category(), + "open " + path.string()); + } + + LARGE_INTEGER file_size {}; + if (!::GetFileSizeEx(handle, &file_size) || file_size.QuadPart < 0 || + static_cast(file_size.QuadPart) > + std::numeric_limits::max()) { + ::CloseHandle(handle); + throw ArtifactError("artifact size does not fit the process address space"); + } + + const auto size = static_cast(file_size.QuadPart); + void* mapping = nullptr; + if (size != 0) { + HANDLE file_mapping = + ::CreateFileMappingW(handle, nullptr, PAGE_READONLY, 0, 0, nullptr); + if (file_mapping == nullptr) { + const DWORD error = ::GetLastError(); + ::CloseHandle(handle); + throw std::system_error(error, std::generic_category(), + "CreateFileMapping " + path.string()); + } + mapping = ::MapViewOfFile(file_mapping, FILE_MAP_READ, 0, 0, 0); + ::CloseHandle(file_mapping); + if (mapping == nullptr) { + const DWORD error = ::GetLastError(); + ::CloseHandle(handle); + throw std::system_error(error, std::generic_category(), + "MapViewOfFile " + path.string()); + } + } + fd_ = handle; + data_ = static_cast(mapping); + size_ = size; +#else const int fd = ::open(path.c_str(), O_RDONLY | O_CLOEXEC | O_DIRECT); if (fd < 0) { throw std::system_error(errno, std::generic_category(), "open " + path.string()); @@ -212,11 +259,19 @@ class MappedFile { fd_ = fd; data_ = static_cast(mapping); size_ = size; +#endif + NINFER_VERBOSE("MappedFile: %s base=%p size=%zu", path.string().c_str(), + static_cast(data_), size_); } ~MappedFile() { +#if defined(_WIN32) + if (data_ != nullptr) { ::UnmapViewOfFile(const_cast(data_)); } + if (fd_ != INVALID_HANDLE_VALUE) { ::CloseHandle(fd_); } +#else if (data_ != nullptr) { ::munmap(const_cast(data_), size_); } if (fd_ >= 0) { ::close(fd_); } +#endif } MappedFile(const MappedFile&) = delete; @@ -232,6 +287,69 @@ class MappedFile { reinterpret_cast(destination.data()) % alignment != 0) { throw ArtifactError("direct artifact read is not 4096-byte aligned"); } + NINFER_VERBOSE("read_direct: offset=%llu size=%zu dest=%p", + static_cast(absolute_offset), destination.size(), + static_cast(destination.data())); +#if defined(_WIN32) + // LARGE_INTEGER is a union whose first member is the anonymous + // { DWORD LowPart; LONG HighPart; } struct, NOT QuadPart. Aggregate + // initialization with a single value therefore sets LowPart to the low + // 32 bits and HighPart to 0, silently truncating any offset >= 2^32. + // Set QuadPart explicitly so the full 64-bit offset is used. + LARGE_INTEGER position {}; + position.QuadPart = static_cast(absolute_offset); + LARGE_INTEGER moved {}; + if (!::SetFilePointerEx(fd_, position, &moved, FILE_BEGIN)) { + throw std::system_error(::GetLastError(), std::generic_category(), + "direct artifact seek"); + } + std::size_t total = 0; + while (total < destination.size()) { + DWORD got = 0; + const BOOL ok = ::ReadFile(fd_, destination.data() + total, + static_cast(destination.size() - total), &got, nullptr); + if (!ok) { + throw std::system_error(::GetLastError(), std::generic_category(), + "direct artifact read"); + } + if (got == 0) { + // EOF reached. Return partial count; caller's short-read check + // will decide whether this is acceptable. + break; + } + total += got; + } + if (ninfer::verbose_enabled()) { + static int verify_count = 0; + static int mismatch_count = 0; + ++verify_count; + const std::size_t n = destination.size() < 64 ? destination.size() : 64; + bool match = true; + for (std::size_t i = 0; i < n; ++i) { + if (destination.data()[i] != data_[absolute_offset + i]) { match = false; break; } + } + if (verify_count == 1) { + std::fprintf(stderr, + "[verbose] read_direct verify probe active (first: offset=%llu match=%s)\n", + (unsigned long long)absolute_offset, match ? "YES" : "NO"); + } + if (!match) { + ++mismatch_count; + std::fprintf(stderr, "[verbose] read_direct MISMATCH #%d offset=%llu\n", + mismatch_count, (unsigned long long)absolute_offset); + std::fprintf(stderr, "[verbose] read = "); + for (std::size_t i = 0; i < n; ++i) { + std::fprintf(stderr, "%02x", (unsigned)destination.data()[i]); + } + std::fprintf(stderr, "\n[verbose] mmap = "); + for (std::size_t i = 0; i < n; ++i) { + std::fprintf(stderr, "%02x", (unsigned)data_[absolute_offset + i]); + } + std::fprintf(stderr, "\n"); + } + } + return total; +#else if (absolute_offset > static_cast(std::numeric_limits::max()) || destination.size() > static_cast(std::numeric_limits::max())) { throw ArtifactError("direct artifact read exceeds platform I/O limits"); @@ -246,10 +364,15 @@ class MappedFile { throw std::system_error(errno, std::generic_category(), "direct artifact read"); } return static_cast(bytes); +#endif } private: - int fd_ = -1; +#if defined(_WIN32) + HANDLE fd_ = INVALID_HANDLE_VALUE; +#else + int fd_ = -1; +#endif const std::byte* data_ = nullptr; std::size_t size_ = 0; }; diff --git a/src/core/verbose.h b/src/core/verbose.h new file mode 100644 index 0000000000..2b7d2aee17 --- /dev/null +++ b/src/core/verbose.h @@ -0,0 +1,37 @@ +#pragma once + +// ninfer::core - toggleable verbose logging for debugging. +// +// Enable by setting the environment variable NINFER_VERBOSE to any value other +// than "0" or empty (e.g. NINFER_VERBOSE=1). The variable is read once and +// cached, so toggling it at runtime has no effect; set it before launch. +// +// Windows (PowerShell): $env:NINFER_VERBOSE="1"; .\serve.ps1 1 +// WSL / bash: NINFER_VERBOSE=1 ./serve.sh 1 +// +// All output goes to stderr with a "[verbose]" prefix so it does not interfere +// with the structured console log. + +#include +#include +#include + +namespace ninfer { + +[[nodiscard]] inline bool verbose_enabled() noexcept { + static const bool enabled = [] { + const char* v = std::getenv("NINFER_VERBOSE"); + return v != nullptr && v[0] != '\0' && std::strcmp(v, "0") != 0; + }(); + return enabled; +} + +} // namespace ninfer + +#define NINFER_VERBOSE(...) \ + do { \ + if (::ninfer::verbose_enabled()) { \ + std::fprintf(stderr, "[verbose] " __VA_ARGS__); \ + std::fputc('\n', stderr); \ + } \ + } while (0) diff --git a/src/media/decode/decode_stub.cpp b/src/media/decode/decode_stub.cpp new file mode 100644 index 0000000000..1a639393a6 --- /dev/null +++ b/src/media/decode/decode_stub.cpp @@ -0,0 +1,30 @@ +// API-compatible stand-in for media/decode/decode.cpp in builds configured +// with NINFER_BUILD_MEDIA=OFF (no FFMPEG). The public API is preserved so the +// vision frontend compiles unchanged; every entry point throws at runtime. +// Text-only servers never reach these calls: the generation service rejects +// media requests when started without --vision. + +#include "media/decode/decode.h" + +#include +#include + +namespace ninfer::media::decode { + +namespace { +[[noreturn]] void unavailable() { + throw std::runtime_error( + "media decode is unavailable in this build; configure with " + "NINFER_BUILD_MEDIA=ON (requires FFMPEG) to serve vision models"); +} +} // namespace + +Image decode_image(std::span, const Policy&) { + unavailable(); +} + +Video decode_video(std::span, const Policy&, double, int, int) { + unavailable(); +} + +} // namespace ninfer::media::decode diff --git a/src/ops/common/device_info.cu b/src/ops/common/device_info.cu new file mode 100644 index 0000000000..9891552b00 --- /dev/null +++ b/src/ops/common/device_info.cu @@ -0,0 +1,30 @@ +#include "ops/common/device_info.h" + +#include +#include + +namespace ninfer::ops { +namespace { + +constexpr int kReferenceSmCount = 170; // RTX 5090 + +int query_sm_count() { + int device = 0; + if (cudaGetDevice(&device) != cudaSuccess) { return kReferenceSmCount; } + int count = 0; + if (cudaDeviceGetAttribute(&count, cudaDevAttrMultiProcessorCount, device) != cudaSuccess) { + return kReferenceSmCount; + } + if (count <= 0) { return kReferenceSmCount; } + std::fprintf(stderr, "ninfer: persistent grids sized for %d SMs\n", count); + return count; +} + +} // namespace + +int device_sm_count() { + static const int count = query_sm_count(); + return count; +} + +} // namespace ninfer::ops diff --git a/src/ops/common/device_info.h b/src/ops/common/device_info.h new file mode 100644 index 0000000000..8517966695 --- /dev/null +++ b/src/ops/common/device_info.h @@ -0,0 +1,18 @@ +#pragma once + +namespace ninfer::ops { + +/** + * Multiprocessor count of the active CUDA device, queried once and cached. + * + * Persistent-grid launchers size one resident wave from this value. Sizing from + * a hardcoded reference-part count leaves multiprocessors idle on a device with + * a wider die (or oversubscribes a narrower one); both are sm_120a parts and + * differ only in enabled SM count. + * + * Returns the reference RTX 5090 count if the device query fails, so a launcher + * always receives a positive, usable value. + */ +int device_sm_count(); + +} // namespace ninfer::ops diff --git a/src/ops/kernel/gqa_attention_decode.cuh b/src/ops/kernel/gqa_attention_decode.cuh index 47953579be..3b8e13a4b6 100644 --- a/src/ops/kernel/gqa_attention_decode.cuh +++ b/src/ops/kernel/gqa_attention_decode.cuh @@ -147,7 +147,7 @@ __launch_bounds__(256) __global__ void gqa_attention_small_t_reduce_output_kerne const __nv_bfloat16* partial_acc, const float* partial_m, const float* partial_l, const std::int32_t* positions, const std::int32_t* valid_columns, std::int32_t tokens, std::int32_t full_width, std::int32_t column_begin, std::int32_t batch_size, - std::int32_t split_count, __nv_bfloat16* out) { + std::int32_t split_count, __nv_bfloat16* out, const __nv_bfloat16* __restrict__ gate) { static_assert(DChunk > 0 && DChunk <= kGqaHeadDim); const int q_head = static_cast(blockIdx.x); @@ -206,7 +206,13 @@ __launch_bounds__(256) __global__ void gqa_attention_small_t_reduce_output_kerne if (head_m == -CUDART_INF_F) { const int d = d_start + tid; if (tid < DChunk && d < kGqaHeadDim) { - out[gqa_q_index(q_head, d, output_column)] = __float2bfloat16(0.0f); + const auto zero_index = gqa_q_index(q_head, d, output_column); + if (gate == nullptr) { + out[zero_index] = __float2bfloat16(0.0f); + } else { + const float gated = 0.0f * sigmoid(__bfloat162float(gate[zero_index])); + out[zero_index] = __float2bfloat16_rn(gated); + } } return; } @@ -254,8 +260,18 @@ __launch_bounds__(256) __global__ void gqa_attention_small_t_reduce_output_kerne if constexpr (Offset) { absolute_column += column_begin; } valid = absolute_column < valid_columns[batch]; } - const float value = (valid && head_l > 0.0f) ? numerator / head_l : 0.0f; - out[gqa_q_index(q_head, d, output_column)] = __float2bfloat16(value); + const float value = (valid && head_l > 0.0f) ? numerator / head_l : 0.0f; + const auto out_index = gqa_q_index(q_head, d, output_column); + if (gate == nullptr) { + out[out_index] = __float2bfloat16(value); + } else { + // Fused sigmoid gate. The standalone elementwise kernel reads the BF16 value + // this store would have produced, so replicate its arithmetic exactly: round + // the reduce result to BF16 first, multiply in FP32, round-to-nearest store. + const __nv_bfloat16 reduced = __float2bfloat16(value); + const float gated = __bfloat162float(reduced) * sigmoid(__bfloat162float(gate[out_index])); + out[out_index] = __float2bfloat16_rn(gated); + } } } // namespace ninfer::ops diff --git a/src/ops/kernel/rope.cuh b/src/ops/kernel/rope.cuh index 73ed115d83..8de10ae07f 100644 --- a/src/ops/kernel/rope.cuh +++ b/src/ops/kernel/rope.cuh @@ -5,6 +5,8 @@ // D/R=128/128, plus packed Vision 16Q/16K at D/R=72/72. One CTA owns one token and shares its // rotary coefficients across heads. +#include "ops/kernel/rmsnorm.cuh" + #include #include @@ -114,6 +116,78 @@ __device__ __forceinline__ void apply_rope_head(__nv_bfloat16* data, std::int64_ __floats2bfloat162_rn(second.x * c0 + first.x * s0, second.y * c1 + first.y * s1); } +// Fused q/k RMS-norm for the Text D=256, rotary-64 decode geometries. One CTA per +// token, warp-per-head (QHeads query rows then KHeads key rows). The norm body +// replicates rmsnorm_warp_bf16x2_kernel (Offset epilogue, identical reduction and +// BF16x2 rounding), so the pair of standalone norm kernels it replaces is preserved +// bit-for-bit. +// +// The rotary step stays in its own kernel on purpose. Folding it in - reusing the +// freshly rounded pair and exchanging its partner through shfl_xor(16) - reproduces +// every formula, coefficient and rounding of apply_rope_head, yet still drifts from +// the standalone rope kernel by a last-bit amount that surfaces as a diverged token +// deep inside long greedy generations and costs about one percent of MTP throughput +// through a lower draft acceptance rate. +template +__launch_bounds__((QHeads + KHeads) * 32) __global__ void qk_norm_rope_text_kernel( + const std::int32_t* __restrict__ positions, const __nv_bfloat162* __restrict__ q_in, + const __nv_bfloat162* __restrict__ q_weight, __nv_bfloat162* __restrict__ q_out, + const __nv_bfloat162* __restrict__ k_in, const __nv_bfloat162* __restrict__ k_weight, + __nv_bfloat162* __restrict__ k_out, float eps) { + constexpr int kPairs = 128; + constexpr int kQStride = QHeads * kPairs; + constexpr int kKStride = KHeads * kPairs; + const int token = static_cast(blockIdx.x); + const int lane = static_cast(threadIdx.x) & 31; + const int warp = static_cast(threadIdx.x) >> 5; + + // Rotary coefficients: lane l owns scalar pair l, same inputs as fixed_sincos. + float lane_sin = 0.0f; + float lane_cos = 0.0f; + { + const float angle = static_cast(positions[token]) * kTextRopeInvFrequency[lane]; + sincosf(angle, &lane_sin, &lane_cos); + } + const int half = lane & 15; + const float c0 = __shfl_sync(kFullWarpMask, lane_cos, half * 2); + const float c1 = __shfl_sync(kFullWarpMask, lane_cos, half * 2 + 1); + const float s0 = __shfl_sync(kFullWarpMask, lane_sin, half * 2); + const float s1 = __shfl_sync(kFullWarpMask, lane_sin, half * 2 + 1); + + const bool is_q = warp < QHeads; + const int head = is_q ? warp : warp - QHeads; + const __nv_bfloat162* x = is_q ? q_in : k_in; + const __nv_bfloat162* w = is_q ? q_weight : k_weight; + __nv_bfloat162* out = is_q ? q_out : k_out; + const std::int64_t row_base = + static_cast(token) * (is_q ? kQStride : kKStride) + + static_cast(head) * kPairs; + + __nv_bfloat162 values[4]; + float sum = 0.0f; +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int pair = lane + item * 32; + values[item] = x[row_base + pair]; + const float2 xf = __bfloat1622float2(values[item]); + sum += xf.x * xf.x + xf.y * xf.y; + } + sum = warp_reduce_sum(sum); + float inv = lane == 0 ? rsqrtf(sum / static_cast(256) + eps) : 0.0f; + inv = __shfl_sync(kFullWarpMask, inv, 0); + +#pragma unroll + for (int item = 0; item < 4; ++item) { + const int pair = lane + item * 32; + const float2 xf = __bfloat1622float2(values[item]); + const float2 wf = __bfloat1622float2(w[pair]); + __nv_bfloat162 stored = + __floats2bfloat162_rn(rmsnorm_epilogue(xf.x, inv, wf.x, 0.0f), + rmsnorm_epilogue(xf.y, inv, wf.y, 0.0f)); + out[row_base + pair] = stored; + } +} + template __global__ void rope_fixed_kernel(const std::int32_t* positions, __nv_bfloat16* q, __nv_bfloat16* k, std::int32_t tokens, std::int64_t q_token_stride, diff --git a/src/ops/launcher/gqa_attention.h b/src/ops/launcher/gqa_attention.h index a05fe9975b..7915248f62 100644 --- a/src/ops/launcher/gqa_attention.h +++ b/src/ops/launcher/gqa_attention.h @@ -40,13 +40,13 @@ void gqa_attention_small_t_launch(const Tensor& q, const Tensor& k, const Tensor PagedKVBatchLayerView cache, GqaExecutionEnvelope envelope, std::int32_t column_begin, std::int32_t width, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, - Tensor& out, cudaStream_t stream); + Tensor& out, cudaStream_t stream, const void* gate = nullptr); void gqa_attention_cached_small_t_launch(const Tensor& q, const Tensor& positions, float scale, const PagedKVLayerView& cache, GqaExecutionEnvelope envelope, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, - cudaStream_t stream); + cudaStream_t stream, const void* gate = nullptr); void gqa_attention_prompt_launch(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& valid_columns, diff --git a/src/ops/launcher/gqa_attention_decode.cu b/src/ops/launcher/gqa_attention_decode.cu index ea286080cb..b44c591323 100644 --- a/src/ops/launcher/gqa_attention_decode.cu +++ b/src/ops/launcher/gqa_attention_decode.cu @@ -238,7 +238,7 @@ void gqa_attention_small_t_launch_for(const Tensor& q, CacheInput input, const T const GqaSmallTInvocation& invocation, GqaExecutionEnvelope envelope, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, - cudaStream_t stream) { + cudaStream_t stream, const void* gate) { const auto logical_capacity = static_cast(envelope.max_visible_keys); const auto implementation_window = static_cast(envelope.max_visible_keys); const auto splits = @@ -313,7 +313,8 @@ void gqa_attention_small_t_launch_for(const Tensor& q, CacheInput input, const T ? nullptr : static_cast(invocation.valid_columns->data), invocation.width, invocation.full_width, invocation.column_begin, - invocation.batch_size, splits, static_cast<__nv_bfloat16*>(out.data)); + invocation.batch_size, splits, static_cast<__nv_bfloat16*>(out.data), + static_cast(gate)); }; const bool masked = invocation.valid_columns != nullptr; const auto launch_profile = [&]() { @@ -350,7 +351,7 @@ void gqa_attention_small_t_launch(const Tensor& q, const Tensor& k, const Tensor PagedKVBatchLayerView cache, GqaExecutionEnvelope envelope, std::int32_t column_begin, std::int32_t width, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, - Tensor& out, cudaStream_t stream) { + Tensor& out, cudaStream_t stream, const void* gate) { const GqaAppendInput input{static_cast(k.data), static_cast(v.data)}; const GqaSmallTInvocation invocation{ @@ -364,19 +365,19 @@ void gqa_attention_small_t_launch(const Tensor& q, const Tensor& k, const Tensor if (q.ne[1] == Gqa27Geometry::QHeads) { gqa_attention_small_t_launch_for(q, input, pos, scale, cache, invocation, envelope, partial_acc, partial_m, partial_l, - out, stream); + out, stream, gate); return; } gqa_attention_small_t_launch_for(q, input, pos, scale, cache, invocation, envelope, partial_acc, partial_m, partial_l, - out, stream); + out, stream, gate); } void gqa_attention_cached_small_t_launch(const Tensor& q, const Tensor& pos, float scale, const PagedKVLayerView& cache, GqaExecutionEnvelope envelope, Tensor& partial_acc, Tensor& partial_m, Tensor& partial_l, Tensor& out, - cudaStream_t stream) { + cudaStream_t stream, const void* gate) { const GqaCachedInput input{}; const GqaSmallTInvocation invocation{ .valid_columns = nullptr, @@ -390,12 +391,12 @@ void gqa_attention_cached_small_t_launch(const Tensor& q, const Tensor& pos, flo if (q.ne[1] == Gqa27Geometry::QHeads) { gqa_attention_small_t_launch_for(q, input, pos, scale, batch_cache, invocation, envelope, partial_acc, - partial_m, partial_l, out, stream); + partial_m, partial_l, out, stream, gate); return; } gqa_attention_small_t_launch_for(q, input, pos, scale, batch_cache, invocation, envelope, partial_acc, partial_m, partial_l, - out, stream); + out, stream, gate); } } // namespace ninfer::ops::detail diff --git a/src/ops/launcher/rope.cu b/src/ops/launcher/rope.cu index 03ca1835a9..f66b31a723 100644 --- a/src/ops/launcher/rope.cu +++ b/src/ops/launcher/rope.cu @@ -197,4 +197,36 @@ void rope_single_launch(const Tensor& positions, int rotary_dim, float theta, Te CUDA_CHECK(cudaGetLastError()); } +namespace { + +template +void launch_qk_norm_rope(const Tensor& positions, const Tensor& q_in, const Tensor& q_weight, + Tensor& q_out, const Tensor& k_in, const Tensor& k_weight, Tensor& k_out, + float eps, cudaStream_t stream) { + const int tokens = positions.ne[0]; + qk_norm_rope_text_kernel<<>>( + static_cast(positions.data), + reinterpret_cast(q_in.data), + reinterpret_cast(q_weight.data), + reinterpret_cast<__nv_bfloat162*>(q_out.data), + reinterpret_cast(k_in.data), + reinterpret_cast(k_weight.data), + reinterpret_cast<__nv_bfloat162*>(k_out.data), eps); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace + +void qk_norm_rope_text_launch(const Tensor& positions, const Tensor& q_in, const Tensor& q_weight, + Tensor& q_out, const Tensor& k_in, const Tensor& k_weight, + Tensor& k_out, float eps, cudaStream_t stream) { + if (q_in.ne[1] == 24 && k_in.ne[1] == 4) { + launch_qk_norm_rope<24, 4>(positions, q_in, q_weight, q_out, k_in, k_weight, k_out, eps, + stream); + return; + } + launch_qk_norm_rope<16, 2>(positions, q_in, q_weight, q_out, k_in, k_weight, k_out, eps, + stream); +} + } // namespace ninfer::ops::detail diff --git a/src/ops/launcher/rope.h b/src/ops/launcher/rope.h index ba35837fc5..9c595a27ab 100644 --- a/src/ops/launcher/rope.h +++ b/src/ops/launcher/rope.h @@ -15,4 +15,8 @@ void rope_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& q void rope_single_launch(const Tensor& positions, int rotary_dim, float theta, Tensor& x, cudaStream_t stream); +void qk_norm_rope_text_launch(const Tensor& positions, const Tensor& q_in, const Tensor& q_weight, + Tensor& q_out, const Tensor& k_in, const Tensor& k_weight, + Tensor& k_out, float eps, cudaStream_t stream); + } // namespace ninfer::ops::detail diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cu b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cu index 19dacec69f..afa0237e9c 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cu +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cu @@ -71,8 +71,11 @@ void launch_tma(const std::uint8_t* activation_codes, const std::uint8_t* activa (void)kConfigured; const dim3 grid(Geometry::kOutputRows / Schedule::kBlockN, tokens / Schedule::kBlockM); + static_assert(sizeof(Nvfp4W4a4TmaDescriptors) == 512); + const std::uint64_t* descriptor_bytes = nvfp4_stage_tma_descriptor(descriptors, stream); nvfp4_w4a4_tma_kernel - <<>>(descriptors, alpha, epilogue, output); + <<>>(descriptor_bytes, alpha, epilogue, + output); CUDA_CHECK(cudaGetLastError()); } diff --git a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh index aa6914eca5..6b5f4b806a 100644 --- a/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh +++ b/src/ops/linear/nvfp4/nvfp4_w4a4_tma.cuh @@ -47,6 +47,34 @@ inline CUtensorMap nvfp4_make_tma_2d(void* address, CUtensorMapDataType data_typ return map; } +// Stage a TMA descriptor into a persistent device buffer and return the device +// pointer for the kernel to read. Passing a host (stack) pointer to the kernel +// relies on the GPU reading host memory over UVA, which is fragile (and the +// stack frame may not outlive an async launch). The 512-byte descriptor is +// copied on the given stream, so the copy is ordered before any kernel launched +// on that stream. Safe for single-stream use (the current deployment); a +// multi-stream caller must supply per-stream buffers. +inline const std::uint64_t* nvfp4_stage_tma_descriptor(const Nvfp4W4a4TmaDescriptors& descriptors, + cudaStream_t stream) { + static std::uint64_t* d_descriptor = [] { + std::uint64_t* p = nullptr; + const cudaError_t err = cudaMalloc(&p, sizeof(Nvfp4W4a4TmaDescriptors)); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("cudaMalloc TMA descriptor: ") + + cudaGetErrorString(err)); + } + return p; + }(); + const cudaError_t err = cudaMemcpyAsync(d_descriptor, &descriptors, + sizeof(Nvfp4W4a4TmaDescriptors), + cudaMemcpyHostToDevice, stream); + if (err != cudaSuccess) { + throw std::runtime_error(std::string("cudaMemcpyAsync TMA descriptor: ") + + cudaGetErrorString(err)); + } + return d_descriptor; +} + template Nvfp4W4a4TmaDescriptors make_nvfp4_w4a4_tma_descriptors(const std::uint8_t* activation_codes, const std::uint8_t* activation_scales, @@ -182,13 +210,20 @@ __device__ __forceinline__ void nvfp4_tma_load_2d(void* destination, const CUten template __global__ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4_tma_kernel( - const __grid_constant__ Nvfp4W4a4TmaDescriptors descriptors, float alpha, + const std::uint64_t descriptors[64], float alpha, const __grid_constant__ Epilogue epilogue, const __grid_constant__ OutputPolicy output) { static_assert((Geometry::kInputRows % Schedule::kBlockK) == 0); static_assert((Geometry::kOutputRows % Schedule::kBlockN) == 0); extern __shared__ __align__(128) unsigned char shared_bytes[]; auto& shared = *reinterpret_cast*>(shared_bytes); + // MSVC cannot pass the 128-aligned TMA descriptor by value (C2719), so it is + // passed as a pointer to a device buffer in global memory. That buffer is + // written by the host (cudaMemcpyAsync H2D in nvfp4_stage_tma_descriptor), so + // it is already visible to the TMA unit's tensormap proxy — no in-kernel + // staging and no tensormap fence are required. (Staging the descriptor into + // local/shared memory inside the kernel makes it invisible to the TMA unit + // without a fence.proxy.tensormap, which surfaces as "Illegal instruction".) const int token_begin = static_cast(blockIdx.y) * Schedule::kBlockM; const int row_begin = static_cast(blockIdx.x) * Schedule::kBlockN; @@ -201,6 +236,8 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 asm volatile("fence.mbarrier_init.release.cluster;" : : : "memory"); } __syncthreads(); + const Nvfp4W4a4TmaDescriptors* tma_desc = + reinterpret_cast(descriptors); constexpr int kKTiles = Geometry::kInputRows / Schedule::kBlockK; @@ -222,17 +259,17 @@ __launch_bounds__(Schedule::kThreads, Schedule::kMinBlocksPerSm) void nvfp4_w4a4 nvfp4_mbarrier_arrive_expect_tx(&shared.full[stage], kTransactionBytes); auto& tensors = shared.scratch.tensors; - nvfp4_tma_load_2d(tensors.a_codes[stage], &descriptors.a_codes, + nvfp4_tma_load_2d(tensors.a_codes[stage], &tma_desc->a_codes, k_tile * Schedule::kCodeRowBytes, token_begin, &shared.full[stage]); - nvfp4_tma_load_2d(tensors.b_codes[stage], &descriptors.b_codes, + nvfp4_tma_load_2d(tensors.b_codes[stage], &tma_desc->b_codes, k_tile * Schedule::kCodeRowBytes, row_begin, &shared.full[stage]); - nvfp4_tma_load_2d(tensors.a_scale4[stage], &descriptors.a_scales, (k_tile / 2) * 16, + nvfp4_tma_load_2d(tensors.a_scale4[stage], &tma_desc->a_scales, (k_tile / 2) * 16, token_begin, &shared.full[stage]); const int b_scale_row = ((row_begin / 128) * Geometry::kScaleTilesPerRow + k_tile * Schedule::kK64PerStage) * 32; - nvfp4_tma_load_2d(tensors.b_scales[stage], &descriptors.b_scales, 0, b_scale_row, + nvfp4_tma_load_2d(tensors.b_scales[stage], &tma_desc->b_scales, 0, b_scale_row, &shared.full[stage]); } } diff --git a/src/ops/linear_attention/gated_delta_net/chunked/output.cu b/src/ops/linear_attention/gated_delta_net/chunked/output.cu index 25dd7b3608..9d5ae4c5f5 100644 --- a/src/ops/linear_attention/gated_delta_net/chunked/output.cu +++ b/src/ops/linear_attention/gated_delta_net/chunked/output.cu @@ -1,14 +1,13 @@ #include "ops/linear_attention/gated_delta_net/chunked/launch.h" #include "ops/linear_attention/gated_delta_net/chunked/output.cuh" +#include "ops/common/device_info.h" namespace ninfer::ops::detail::gated_delta_net::chunked { namespace { namespace kernel = output; -constexpr std::int64_t kRtx5090SmCount = 170; -constexpr std::int64_t kCtasPerSm = 4; -constexpr std::int64_t kTargetCtas = kRtx5090SmCount * kCtasPerSm; +constexpr std::int64_t kCtasPerSm = 4; template cudaError_t launch_fixed(const chunk_output_config& cfg, dim3 grid, head_map qk_map, int chunks) { @@ -40,10 +39,11 @@ cudaError_t launch_output(const chunk_output_config& cfg) { const auto qk_map = head_map::of((int)cfg.H_qk, (int)cfg.H_v); const std::int64_t NT = cfg.L / BT; - // Keep at most one resident RTX 5090 wave and distribute chunks evenly - // across it. Small grids retain one logical job per CTA. + // Keep at most one resident wave for THIS device and distribute chunks + // evenly across it. Small grids retain one logical job per CTA. + const std::int64_t target_ctas = static_cast(device_sm_count()) * kCtasPerSm; const std::int64_t logical_jobs = NT * cfg.H_v; - const std::int64_t jobs_per_block = (logical_jobs + kTargetCtas - 1) / kTargetCtas; + const std::int64_t jobs_per_block = (logical_jobs + target_ctas - 1) / target_ctas; const std::int64_t grid_chunks = (NT + jobs_per_block - 1) / jobs_per_block; NINFER_GATED_DELTA_NET_PROPAGATE(v.check_grid(grid_chunks, cfg.H_v)); diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu index 0127a8d1d5..c61bbf8837 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cu @@ -71,8 +71,10 @@ void launch_nvfp4_linear_swiglu_w4a4_tma(const std::uint8_t* activation_codes, activation_codes, activation_scales, weight_codes, weight_scales, tokens); constexpr int kPairN = M256N128S3::kBlockN / 2; const dim3 grid((Geometry::kOutputRows / 2) / kPairN, tokens / M256N128S3::kBlockM); + static_assert(sizeof(Nvfp4W4a4TmaDescriptors) == 512); + const std::uint64_t* descriptor_bytes = nvfp4_stage_tma_descriptor(descriptors, stream); nvfp4_linear_swiglu_w4a4_tma_kernel - <<>>(descriptors, alpha, output); + <<>>(descriptor_bytes, alpha, output); CUDA_CHECK(cudaGetLastError()); } diff --git a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh index a7664c7da7..32b7752663 100644 --- a/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh +++ b/src/ops/linear_swiglu/nvfp4/nvfp4_linear_swiglu_w4a4_tma.cuh @@ -46,9 +46,8 @@ template __global__ __launch_bounds__( Schedule::kThreads, Schedule:: - kMinBlocksPerSm) void nvfp4_linear_swiglu_w4a4_tma_kernel(const __grid_constant__ - Nvfp4W4a4TmaDescriptors - descriptors, + kMinBlocksPerSm) void nvfp4_linear_swiglu_w4a4_tma_kernel(const std::uint64_t + descriptors[64], float alpha, __nv_bfloat16* __restrict__ output) { static_assert(Geometry::kOutputRows == 34816); @@ -65,6 +64,13 @@ __global__ __launch_bounds__( extern __shared__ __align__(128) unsigned char shared_bytes[]; auto& shared = *reinterpret_cast*>(shared_bytes); + // MSVC cannot pass the 128-aligned TMA descriptor by value (C2719), so it is + // passed as a pointer to a device buffer in global memory. That buffer is + // written by the host (cudaMemcpyAsync H2D in nvfp4_stage_tma_descriptor), so + // it is already visible to the TMA unit's tensormap proxy — no in-kernel + // staging and no tensormap fence are required. (Staging the descriptor into + // local/shared memory inside the kernel makes it invisible to the TMA unit + // without a fence.proxy.tensormap, which surfaces as "Illegal instruction".) const int token_begin = static_cast(blockIdx.y) * Schedule::kBlockM; const int pair_begin = static_cast(blockIdx.x) * kPairN; @@ -77,6 +83,8 @@ __global__ __launch_bounds__( asm volatile("fence.mbarrier_init.release.cluster;" : : : "memory"); } __syncthreads(); + const Nvfp4W4a4TmaDescriptors* tma_desc = + reinterpret_cast(descriptors); constexpr int kKTiles = Geometry::kInputRows / Schedule::kBlockK; @@ -98,16 +106,16 @@ __global__ __launch_bounds__( nvfp4_mbarrier_arrive_expect_tx(&shared.full[stage], kTransactionBytes); auto& tensors = shared.scratch.tensors; - nvfp4_tma_load_2d(tensors.a_codes[stage], &descriptors.a_codes, + nvfp4_tma_load_2d(tensors.a_codes[stage], &tma_desc->a_codes, k_tile * Schedule::kCodeRowBytes, token_begin, &shared.full[stage]); - nvfp4_tma_load_2d(tensors.b_codes[stage], &descriptors.b_codes, + nvfp4_tma_load_2d(tensors.b_codes[stage], &tma_desc->b_codes, k_tile * Schedule::kCodeRowBytes, pair_begin, &shared.full[stage]); nvfp4_tma_load_2d(tensors.b_codes[stage] + kPairN * Schedule::kCodeRowBytes, - &descriptors.b_codes, k_tile * Schedule::kCodeRowBytes, + &tma_desc->b_codes, k_tile * Schedule::kCodeRowBytes, pair_begin + kIntermediate, &shared.full[stage]); - nvfp4_tma_load_2d(tensors.a_scale4[stage], &descriptors.a_scales, (k_tile / 2) * 16, + nvfp4_tma_load_2d(tensors.a_scale4[stage], &tma_desc->a_scales, (k_tile / 2) * 16, token_begin, &shared.full[stage]); const int gate_scale_row = ((pair_begin / 128) * Geometry::kScaleTilesPerRow + @@ -117,9 +125,9 @@ __global__ __launch_bounds__( (((pair_begin + kIntermediate) / 128) * Geometry::kScaleTilesPerRow + k_tile * Schedule::kK64PerStage) * 32; - nvfp4_tma_load_2d(tensors.b_scales[stage][0], &descriptors.b_scales, 0, + nvfp4_tma_load_2d(tensors.b_scales[stage][0], &tma_desc->b_scales, 0, gate_scale_row, &shared.full[stage]); - nvfp4_tma_load_2d(tensors.b_scales[stage][1], &descriptors.b_scales, 0, + nvfp4_tma_load_2d(tensors.b_scales[stage][1], &tma_desc->b_scales, 0, up_scale_row, &shared.full[stage]); } } diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode.h b/src/ops/sparse_moe/decode/sparse_moe_decode.h index 30da87bf71..c7b2e2b69b 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode.h +++ b/src/ops/sparse_moe/decode/sparse_moe_decode.h @@ -53,6 +53,7 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor cudaStream_t stream, const int* adaptive_route_jobs = nullptr); void sparse_moe_decode_launch(const Tensor& x, const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream); + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data = nullptr, std::size_t prefetch_bytes = 0); } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 12a68e0c18..5ccf460fd7 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -67,13 +67,26 @@ __device__ __forceinline__ float router_row_dot(const __nv_bfloat16* x, const __ return warp_reduce_sum(sum); } +// Ticket for the last-arriving D1 block. atomicInc wraps at gridDim.x - 1, so the counter returns +// to zero on its own and needs no host-side initialisation or workspace slot. +// Safety: one Engine per process, one compute stream. A second concurrent D1 grid on this +// device shares the ticket and silently corrupts routing. Today's executor is a single +// worker on device.stream; PDL dependents complete before the next decode step. +__device__ unsigned int g_sparse_moe_route_ticket = 0; + __global__ void sparse_moe_d1_kernel(const __nv_bfloat16* __restrict__ x, const __nv_bfloat16* __restrict__ router, - float* __restrict__ scores) { + float* __restrict__ scores, int* __restrict__ ids, + float* __restrict__ alpha, float* __restrict__ shared_scale, + const char* __restrict__ shared_down_payload, + unsigned long long shared_down_bytes) { __shared__ float partial[kD1Warps]; + __shared__ float selected_logits[kTopK]; + __shared__ bool is_last_block; const int row = static_cast(blockIdx.x); const int warp = static_cast(threadIdx.x) >> 5; const int lane = static_cast(threadIdx.x) & 31; + if (threadIdx.x == 0) { pdl::trigger_dependents(); } const float dot = router_row_dot(x, router + static_cast(row) * kHidden); if (lane == 0) { partial[warp] = dot; } __syncthreads(); @@ -82,14 +95,34 @@ __global__ void sparse_moe_d1_kernel(const __nv_bfloat16* __restrict__ x, value = warp_reduce_sum(value); if (lane == 0) { scores[row] = value; } } -} - -__global__ void sparse_moe_d2_warp_kernel(const float* __restrict__ scores, int* __restrict__ ids, - float* __restrict__ alpha, - float* __restrict__ shared_scale) { - __shared__ float selected_logits[kTopK]; - if (threadIdx.x == 0) { pdl::trigger_dependents(); } - sparse_moe_select_top8_warp(scores, ids, alpha, shared_scale, selected_logits); + if (shared_down_payload != nullptr) { + // Warm L2 for the shared-expert down payload while the bus idles behind the + // router: D4's shared warp streams these bytes last and otherwise sets the + // block's critical path. One 128B line per thread covers the payload in a + // single sweep. A pure cache hint. + const unsigned long long offset = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) * 128ull; + if (offset < shared_down_bytes) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(shared_down_payload + offset)); + } + } + // The top-8 selection used to be its own single-warp grid. Its cost was the price of the node + // rather than the work, so the block that arrives last runs the identical routine instead. The + // inputs, the comparison order and the softmax are unchanged, so the routing is identical. + if (threadIdx.x == 0) { + __threadfence(); + const unsigned int ticket = atomicInc(&g_sparse_moe_route_ticket, gridDim.x - 1u); + is_last_block = ticket == gridDim.x - 1u; + // Writer stores were fence-then-atomicInc. The winning block's later loads of all + // 257 scores still need a device-scope fence before those loads; without it this + // matches the canonical last-block pattern but is a PTX-model race (stale L1 on a + // future arch would change top-8 routing with no error). + if (is_last_block) { __threadfence(); } + } + __syncthreads(); + if (is_last_block && warp == 0) { + sparse_moe_select_top8_warp(scores, ids, alpha, shared_scale, selected_logits); + } } struct Q4Codec { @@ -378,7 +411,8 @@ __global__ void sparse_moe_d4_nine_warp_kernel( const float* __restrict__ shared_scale, const float* __restrict__ act, const std::uint8_t* __restrict__ routed_codes, const std::uint8_t* __restrict__ routed_high, const std::uint8_t* __restrict__ routed_scales, const std::uint8_t* __restrict__ shared_codes, - const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination) { + const std::uint8_t* __restrict__ shared_scales, __nv_bfloat16* __restrict__ destination, + const char* __restrict__ prefetch_data, unsigned long long prefetch_bytes) { __shared__ float paths[kTopK + 1][Rows]; pdl::wait_for_dependencies(); const int warp = static_cast(threadIdx.x) >> 5; @@ -412,6 +446,17 @@ __global__ void sparse_moe_d4_nine_warp_kernel( for (int path = 0; path < kTopK + 1; ++path) { value += paths[path][lane]; } destination[row_base + lane] = __float2bfloat16_rn(value); } + if (prefetch_data != nullptr) { + // Fire-and-forget L2 warmup of the next consumer's weight payload. D4 CTAs + // retire in waves across the tail of the MoE window while the bus is + // largely idle; one 128B line per thread covers the whole span in a single + // sweep. A pure cache hint: no value and no addition order is touched. + const unsigned long long offset = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) * 128ull; + if (offset < prefetch_bytes) { + asm volatile("prefetch.global.L2 [%0];" ::"l"(prefetch_data + offset)); + } + } } template @@ -480,12 +525,16 @@ __global__ void sparse_moe_d4_token_kernel( } } -void launch_d1(const Tensor& x, const Weight& router_shared_gate, +void launch_d1(const Tensor& x, const SparseMoeWeights& weights, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { sparse_moe_d1_kernel<<>>( static_cast(x.data), - static_cast(router_shared_gate.qdata), - static_cast(workspace.scratch.data)); + static_cast(weights.router_shared_gate.qdata), + static_cast(workspace.scratch.data), static_cast(workspace.ids.data), + static_cast(workspace.alpha.data), + static_cast(workspace.shared_scale.data), + static_cast(weights.shared_down.qdata), + static_cast(weights.shared_down.payload_bytes)); CUDA_CHECK(cudaGetLastError()); } @@ -507,13 +556,6 @@ void launch_d3_dependent_codec(const Tensor& x, const SparseMoeWeights& weights, void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { - const auto* scores = static_cast(workspace.scratch.data); - auto* ids = static_cast(workspace.ids.data); - auto* alpha = static_cast(workspace.alpha.data); - auto* shared_scale = static_cast(workspace.shared_scale.data); - sparse_moe_d2_warp_kernel<<<1, 32, 0, stream>>>(scores, ids, alpha, shared_scale); - CUDA_CHECK(cudaGetLastError()); - switch (weights.routed_gate_up.qtype) { case QType::Q4G64_F16S: launch_d3_dependent_codec(x, weights, workspace, stream); @@ -528,7 +570,8 @@ void launch_d2_d3(const Tensor& x, const SparseMoeWeights& weights, template void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data, std::size_t prefetch_bytes) { const auto* ids = static_cast(workspace.ids.data); const auto* alpha = static_cast(workspace.alpha.data); const auto* shared_scale = static_cast(workspace.shared_scale.data); @@ -542,20 +585,26 @@ void launch_d4_dependent_codec(const SparseMoeWeights& weights, Tensor& destinat CUDA_CHECK(pdl::launch_dependent({dim3(kHidden), dim3(9 * 32), 0, stream}, sparse_moe_d4_nine_warp_kernel, ids, alpha, shared_scale, act, routed_codes, routed_high, routed_scales, - shared_codes, shared_scales, output)); + shared_codes, shared_scales, output, + static_cast(prefetch_data), + static_cast(prefetch_bytes))); } void launch_d4_dependent(const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data, std::size_t prefetch_bytes) { switch (weights.routed_down.qtype) { case QType::Q5G64_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; case QType::Q6G64_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; case QType::W8G32_F16S: - launch_d4_dependent_codec(weights, destination, workspace, stream); + launch_d4_dependent_codec(weights, destination, workspace, stream, prefetch_data, + prefetch_bytes); return; default: throw std::invalid_argument("sparse_moe: unsupported D4 codec"); @@ -718,10 +767,11 @@ void sparse_moe_decode_launch_d4_small_t(const SparseMoeWeights& weights, Tensor } void sparse_moe_decode_launch(const Tensor& x, const SparseMoeWeights& weights, Tensor& destination, - const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream) { - launch_d1(x, weights.router_shared_gate, workspace, stream); + const SparseMoeDecodeWorkspace& workspace, cudaStream_t stream, + const void* prefetch_data, std::size_t prefetch_bytes) { + launch_d1(x, weights, workspace, stream); launch_d2_d3(x, weights, workspace, stream); - launch_d4_dependent(weights, destination, workspace, stream); + launch_d4_dependent(weights, destination, workspace, stream, prefetch_data, prefetch_bytes); } } // namespace ninfer::ops::detail diff --git a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu index 39653a2a57..5f5c6d70e7 100644 --- a/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu +++ b/src/ops/sparse_moe/prefill/sparse_moe_prefill_kernels.cu @@ -20,6 +20,8 @@ #include #include +#include "ops/common/device_info.h" + namespace ninfer::ops::detail { namespace { @@ -253,9 +255,7 @@ constexpr int kExpertBK = 64; constexpr int kExpertStages = 2; constexpr int kExpertWarps = 8; constexpr int kExpertThreads = 32 * kExpertWarps; -constexpr int kRtx5090SmCount = 170; -constexpr int kPrefillBlocksPerSm = 3; -constexpr int kPrefillPersistentBlocks = kPrefillBlocksPerSm * kRtx5090SmCount; +constexpr int kPrefillBlocksPerSm = 3; template __global__ __launch_bounds__(ExpertWarps * 32, 3) void sparse_moe_prefill_q4_gate_up_kernel( @@ -1087,6 +1087,9 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, throw std::invalid_argument("sparse_moe prefill: launch plan does not match tensors"); } + // One resident persistent wave sized for THIS device, not a reference part. + const int prefill_persistent_blocks = kPrefillBlocksPerSm * device_sm_count(); + const auto* router = static_cast(weights.router_shared_gate.qdata); const auto* routed_gate_codes = static_cast(weights.routed_gate_up.qdata); const auto* routed_gate_scales = @@ -1168,12 +1171,12 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, if (weights.routed_gate_up.qtype == QType::Q4G64_F16S) { if (wide_plan) { sparse_moe_prefill_q4_gate_up_kernel<8, 64> - <<>>( + <<>>( grouped_io, offsets, route_job_experts, route_job_columns, route_job_count, routed_gate_codes, routed_gate_scales, routed_activation); } else { sparse_moe_prefill_q4_gate_up_kernel<4, 32> - <<>>( + <<>>( grouped_io, offsets, route_job_experts, route_job_columns, route_job_count, routed_gate_codes, routed_gate_scales, routed_activation); } @@ -1207,13 +1210,13 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, case QType::Q5G64_F16S: if (wide_plan) { sparse_moe_prefill_qx_down_kernel - <<>>( + <<>>( routed_activation, offsets, route_job_experts, route_job_columns, route_job_count, routed_down_codes, routed_down_high, routed_down_scales, grouped_io); } else { sparse_moe_prefill_qx_down_kernel - <<>>( + <<>>( routed_activation, offsets, route_job_experts, route_job_columns, route_job_count, routed_down_codes, routed_down_high, routed_down_scales, grouped_io); @@ -1222,13 +1225,13 @@ void sparse_moe_prefill_launch(const Tensor& x, const SparseMoeWeights& weights, case QType::Q6G64_F16S: if (wide_plan) { sparse_moe_prefill_qx_down_kernel - <<>>( + <<>>( routed_activation, offsets, route_job_experts, route_job_columns, route_job_count, routed_down_codes, routed_down_high, routed_down_scales, grouped_io); } else { sparse_moe_prefill_qx_down_kernel - <<>>( + <<>>( routed_activation, offsets, route_job_experts, route_job_columns, route_job_count, routed_down_codes, routed_down_high, routed_down_scales, grouped_io); diff --git a/src/ops/wrapper/embedding.cpp b/src/ops/wrapper/embedding.cpp index c536663cdc..eb43709095 100644 --- a/src/ops/wrapper/embedding.cpp +++ b/src/ops/wrapper/embedding.cpp @@ -4,16 +4,82 @@ #include "ops/common/math.h" #include "ops/linear/fp8/fp8_format.h" #include "ops/launcher/embed_gather.h" // detail::embed_gather_*_launch +#include "core/verbose.h" #include "core/weight.h" +#include + #include +#include #include #include #include +#include namespace ninfer::ops { namespace { +// Verbose: true if the stream is in CUDA graph capture mode (blocking readbacks +// are illegal then). +bool verbose_stream_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + return cudaStreamIsCapturing(stream, &capture) == cudaSuccess && + capture != cudaStreamCaptureStatusNone; +} + +// Verbose probe: validate that each device pointer is a real device allocation, +// print the weight metadata, and dump the actual token ids (device->host) so an +// out-of-range row (the usual cause of an illegal address in a gather kernel) +// is visible. Runs before the launch, so the CUDA context is still clean. +void verbose_probe_pointers(const char* tag, const Tensor& ids, const Weight& table, + const Tensor& out, cudaStream_t stream) { + if (!verbose_enabled()) { return; } + auto describe = [](const char* name, const void* p) { + if (p == nullptr) { + std::fprintf(stderr, "[verbose] %-8s = (null)\n", name); + return; + } + cudaPointerAttributes attrs {}; + const cudaError_t err = cudaPointerGetAttributes(&attrs, p); + if (err != cudaSuccess) { + std::fprintf(stderr, "[verbose] %-8s = %p (cudaPointerGetAttributes FAILED: %s)\n", + name, p, cudaGetErrorString(err)); + return; + } + std::fprintf(stderr, "[verbose] %-8s = %p type=%d device=%d devptr=%p\n", name, p, + static_cast(attrs.type), attrs.device, attrs.devicePointer); + }; + const std::int32_t T = ids.ne[0]; + std::fprintf(stderr, + "[verbose] embedding(%s): T=%d vocab(n)=%d hidden(k)=%d out_d=%d " + "payload_bytes=%llu layout=%d scale_dtype=%d padded=[%d,%d,%d,%d]\n", + tag, T, table.n, table.k, out.ne[0], + static_cast(table.payload_bytes), + static_cast(table.layout), static_cast(table.scale_dtype), + table.padded_shape[0], table.padded_shape[1], table.padded_shape[2], + table.padded_shape[3]); + describe("ids", ids.data); + describe("qdata", table.qdata); + describe("scales", table.scales); + describe("out", out.data); + if (ids.data != nullptr && T > 0 && !verbose_stream_capturing(stream)) { + std::vector host_ids(static_cast(T)); + const cudaError_t err = cudaMemcpy(host_ids.data(), ids.data, + static_cast(T) * sizeof(std::int32_t), + cudaMemcpyDeviceToHost); + if (err != cudaSuccess) { + std::fprintf(stderr, "[verbose] ids dump FAILED: %s\n", cudaGetErrorString(err)); + } else { + std::fprintf(stderr, "[verbose] ids = ["); + for (std::int32_t i = 0; i < T; ++i) { + std::fprintf(stderr, "%s%d%s", i ? ", " : "", host_ids[i], + host_ids[i] >= table.n ? " OOB!" : ""); + } + std::fprintf(stderr, "]\n"); + } + } +} + std::int64_t numel_allow_zero(const Tensor& t, const char* label) { bool has_zero = false; for (int d = 0; d < 4; ++d) { @@ -230,6 +296,7 @@ void embedding(const Tensor& ids, const Weight& table, Tensor& out, cudaStream_t require_fp8_metadata(table, out); if (is_empty_T(ids, out)) { return; } require_non_empty_tensors(ids, out); + verbose_probe_pointers("fp8", ids, table, out, stream); detail::embed_gather_fp8_launch(ids, table, out, stream); break; default: diff --git a/src/ops/wrapper/gqa_attention.cpp b/src/ops/wrapper/gqa_attention.cpp index b85b6c17e9..3732df20fa 100644 --- a/src/ops/wrapper/gqa_attention.cpp +++ b/src/ops/wrapper/gqa_attention.cpp @@ -10,6 +10,7 @@ #include #include #include +#include "ninfer/ops/sigmoid_mul.h" namespace ninfer::ops { namespace { @@ -397,7 +398,8 @@ std::size_t gqa_attention_workspace_capacity_bytes(std::int32_t q_heads, DType c void gqa_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tensor& positions, const Tensor& valid_columns, const Tensor& kv_table_rows, float scale, PagedKVBatchLayerView cache, GqaExecutionEnvelope envelope, - WorkspaceArena& workspace, Tensor& out, cudaStream_t stream) { + WorkspaceArena& workspace, Tensor& out, cudaStream_t stream, + const Tensor* gate) { constexpr const char* op = "gqa_attention"; validate_batched_attention_tensors(q, positions, valid_columns, kv_table_rows, out, cache, envelope, scale, op); @@ -418,6 +420,7 @@ void gqa_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tens if (route == detail::GqaAttentionRoute::ChunkedSmallT) { launch_chunked_small_t(q, k, v, positions, valid_columns, kv_table_rows, scale, cache, envelope, workspace, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } return; } if (route == detail::GqaAttentionRoute::SmallT) { @@ -427,11 +430,13 @@ void gqa_attention(const Tensor& q, const Tensor& k, const Tensor& v, const Tens allocate_small_t_workspace(workspace, q.ne[1], width, splits, batch); detail::gqa_attention_small_t_launch(q, k, v, positions, valid_columns, kv_table_rows, scale, cache, envelope, 0, width, partial.acc, - partial.m, partial.l, out, stream); + partial.m, partial.l, out, stream, + gate == nullptr ? nullptr : gate->data); return; } detail::gqa_attention_prompt_launch(q, k, v, positions, valid_columns, kv_table_rows, scale, cache, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } } void gqa_kv_append(const Tensor& k, const Tensor& v, const Tensor& positions, @@ -462,7 +467,8 @@ void gqa_kv_append(const Tensor& k, const Tensor& v, const Tensor& positions, void gqa_attention_cached(const Tensor& q, const Tensor& positions, float scale, const PagedKVLayerView& cache, GqaExecutionEnvelope envelope, - WorkspaceArena& workspace, Tensor& out, cudaStream_t stream) { + WorkspaceArena& workspace, Tensor& out, cudaStream_t stream, + const Tensor* gate) { constexpr const char* op = "gqa_attention_cached"; validate_attention_tensors(q, positions, out, cache, envelope, scale, op); @@ -470,6 +476,7 @@ void gqa_attention_cached(const Tensor& q, const Tensor& positions, float scale, if (detail::gqa_attention_resolve_route(q.ne[1], q.ne[2], 1, envelope) == detail::GqaAttentionRoute::ChunkedSmallT) { launch_cached_chunked_small_t(q, positions, scale, cache, envelope, workspace, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } return; } if (detail::gqa_attention_uses_small_t(q.ne[2])) { @@ -477,10 +484,12 @@ void gqa_attention_cached(const Tensor& q, const Tensor& positions, float scale, detail::gqa_attention_split_capacity(q.ne[1], q.ne[2], cache.dtype, envelope); SmallTWorkspace partial = allocate_small_t_workspace(workspace, q.ne[1], q.ne[2], splits); detail::gqa_attention_cached_small_t_launch(q, positions, scale, cache, envelope, - partial.acc, partial.m, partial.l, out, stream); + partial.acc, partial.m, partial.l, out, stream, + gate == nullptr ? nullptr : gate->data); return; } detail::gqa_attention_prompt_attention_launch(q, positions, scale, cache, out, stream); + if (gate != nullptr) { sigmoid_mul(*gate, out, stream); } } } // namespace ninfer::ops diff --git a/src/ops/wrapper/rope.cpp b/src/ops/wrapper/rope.cpp index 17d574db88..61b21ac4aa 100644 --- a/src/ops/wrapper/rope.cpp +++ b/src/ops/wrapper/rope.cpp @@ -7,6 +7,7 @@ #include #include #include +#include "ninfer/ops/rmsnorm.h" namespace ninfer::ops { namespace { @@ -140,4 +141,22 @@ void rope(const Tensor& positions, int rotary_dim, float theta, Tensor& x, cudaS detail::rope_single_launch(positions, rotary_dim, theta, x, stream); } +void qk_norm_rope(const Tensor& positions, int rotary_dim, float theta, const Tensor& q_in, + const Tensor& q_weight, Tensor& q_out, const Tensor& k_in, + const Tensor& k_weight, Tensor& k_out, float eps, cudaStream_t stream) { + const bool text_heads = (q_in.ne[1] == 16 && k_in.ne[1] == 2) || + (q_in.ne[1] == 24 && k_in.ne[1] == 4); + const bool fused_shape = text_heads && q_in.ne[0] == 256 && k_in.ne[0] == 256 && + rotary_dim == 64 && positions.dtype == DType::I32; + if (fused_shape) { + detail::qk_norm_rope_text_launch(positions, q_in, q_weight, q_out, k_in, k_weight, k_out, + eps, stream); + rope(positions, rotary_dim, theta, q_out, k_out, stream); + return; + } + rmsnorm(q_in, q_weight, eps, true, q_out, stream); + rmsnorm(k_in, k_weight, eps, true, k_out, stream); + rope(positions, rotary_dim, theta, q_out, k_out, stream); +} + } // namespace ninfer::ops diff --git a/src/ops/wrapper/sparse_moe.cpp b/src/ops/wrapper/sparse_moe.cpp index 48e9c725fc..8b288880df 100644 --- a/src/ops/wrapper/sparse_moe.cpp +++ b/src/ops/wrapper/sparse_moe.cpp @@ -189,8 +189,18 @@ std::size_t sparse_moe_workspace_capacity_bytes(QType routed_gate_up, QType rout return required; } +namespace { +// Cap keeps the warmed span well inside L2 next to the layer's own streams. +constexpr std::size_t kNextWeightPrefetchLimit = std::size_t{8} << 20; +} // namespace + void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilogue epilogue, - Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream) { + Tensor& destination, WorkspaceArena& workspace, cudaStream_t stream, + WeightPrefetchSpan next_prefetch_request) { + const WeightPrefetchSpan next_prefetch{ + next_prefetch_request.data, + next_prefetch_request.bytes < kNextWeightPrefetchLimit ? next_prefetch_request.bytes + : kNextWeightPrefetchLimit}; if (epilogue != SparseMoeEpilogue::AddResidual) { throw std::invalid_argument("sparse_moe: unsupported epilogue"); } @@ -258,7 +268,8 @@ void sparse_moe(const Tensor& x, const SparseMoeWeights& weights, SparseMoeEpilo for (std::int32_t token = 0; token < tokens; ++token) { const Tensor x_column = x.slice(1, token, 1); Tensor destination_column = destination.slice(1, token, 1); - detail::sparse_moe_decode_launch(x_column, weights, destination_column, views, stream); + detail::sparse_moe_decode_launch(x_column, weights, destination_column, views, stream, + next_prefetch.data, next_prefetch.bytes); } } diff --git a/src/product/load_progress/load_progress.cpp b/src/product/load_progress/load_progress.cpp index 2617b825b9..9432641a92 100644 --- a/src/product/load_progress/load_progress.cpp +++ b/src/product/load_progress/load_progress.cpp @@ -1,6 +1,10 @@ #include "product/load_progress/load_progress.h" +#if defined(_WIN32) +#include +#else #include +#endif #include #include @@ -60,7 +64,14 @@ std::string format_line(std::string_view phase, std::uint64_t done, std::uint64_ } // namespace LoadProgressRendererOptions stderr_load_progress_options() noexcept { - if (::isatty(STDERR_FILENO) == 1) { +#if defined(_WIN32) + DWORD console_mode = 0; + const bool is_terminal = + ::GetConsoleMode(::GetStdHandle(STD_ERROR_HANDLE), &console_mode) != 0; +#else + const bool is_terminal = ::isatty(STDERR_FILENO) == 1; +#endif + if (is_terminal) { return LoadProgressRendererOptions{ .mode = LoadProgressOutputMode::Interactive, .min_refresh_interval = std::chrono::milliseconds(200), diff --git a/src/product/media_acquire/acquire.cpp b/src/product/media_acquire/acquire.cpp index 1f03ac9c64..40682324b5 100644 --- a/src/product/media_acquire/acquire.cpp +++ b/src/product/media_acquire/acquire.cpp @@ -2,15 +2,21 @@ #include +#if defined(_WIN32) +#include +#include +#else #include #include #include +#endif #include #include #include #include #include +#include #include #include #include @@ -203,6 +209,14 @@ std::vector fetch_url(std::string url, const Policy& policy) { if (!policy.allow_remote) { throw std::invalid_argument("remote media URLs are disabled"); } static std::once_flag init; std::call_once(init, [] { +#if defined(_WIN32) + WSADATA wsa_data {}; + if (::WSAStartup(MAKEWORD(2, 2), &wsa_data) != 0) { + throw std::runtime_error("failed to initialize Winsock"); + } + static const auto wsa_cleanup = [] { ::WSACleanup(); }; + std::atexit(wsa_cleanup); +#endif if (curl_global_init(CURL_GLOBAL_DEFAULT) != CURLE_OK) { throw std::runtime_error("failed to initialize libcurl"); } @@ -287,7 +301,11 @@ std::vector read_path(const Source& source, const Policy& policy) if (!policy.media_root.empty()) { const std::filesystem::path root = std::filesystem::weakly_canonical(policy.media_root, ec); const auto relative = std::filesystem::relative(path, root, ec); +#if defined(_WIN32) + if (ec || relative.empty() || relative.native().starts_with(L"..")) { +#else if (ec || relative.empty() || relative.native().starts_with("..")) { +#endif throw std::invalid_argument("media path is outside configured media root"); } } diff --git a/src/product/media_acquire/acquire_stub.cpp b/src/product/media_acquire/acquire_stub.cpp new file mode 100644 index 0000000000..935b8769a1 --- /dev/null +++ b/src/product/media_acquire/acquire_stub.cpp @@ -0,0 +1,25 @@ +// API-compatible stand-in for product/media_acquire/acquire.cpp in builds +// configured with NINFER_BUILD_MEDIA=OFF (no libcurl). The public API is +// preserved so the serve layer and prompt input compile unchanged; every +// entry point throws at runtime. Text-only servers never reach these calls: +// the generation service rejects media requests when started without +// --vision. + +#include "product/media_acquire/acquire.h" + +#include +#include + +namespace ninfer::product::media_acquire { + +[[noreturn]] static void unavailable() { + throw std::runtime_error( + "media acquisition is unavailable in this build; configure with " + "NINFER_BUILD_MEDIA=ON (requires libcurl) to serve vision models"); +} + +std::vector acquire_bytes(const Source&, const Policy&) { + unavailable(); +} + +} // namespace ninfer::product::media_acquire diff --git a/src/runtime/engine/concurrent_executor.h b/src/runtime/engine/concurrent_executor.h index 5e4729cab9..356f327f05 100644 --- a/src/runtime/engine/concurrent_executor.h +++ b/src/runtime/engine/concurrent_executor.h @@ -16,15 +16,19 @@ #include #include #include +#include #include #include +#include #include #include #include #include +#include #include #include #include +#include #include #include @@ -46,7 +50,13 @@ class ConcurrentExecutor { max_outstanding_(static_cast(options.max_concurrency) + options.max_pending_requests), pending_timeout_(std::chrono::milliseconds(options.pending_timeout_ms)), - admission_capacity_(instance.program->admission_capacity()) { + admission_capacity_(instance.program->admission_capacity()), + prefix_cache_min_bytes_(options.prefix_cache_min_bytes), + prefix_cache_max_bytes_(options.prefix_cache_max_bytes == 0 ? options.prefix_cache_bytes + : options.prefix_cache_max_bytes), + vram_floor_bytes_(options.vram_floor_bytes), + vram_observe_only_(options.vram_observe_only), + on_fatal_error_(options.on_fatal_error) { if (max_concurrency_ == 0 || max_concurrency_ > kMaximumConcurrency || options.max_pending_requests == 0 || pending_timeout_.count() <= 0) { throw std::invalid_argument("concurrent executor bounds are invalid"); @@ -190,6 +200,11 @@ class ConcurrentExecutor { return published_stats_; } + [[nodiscard]] bool is_healthy() const noexcept { + std::lock_guard lock(queue_mutex_); + return !stopping_ && !failed_; + } + void reset_memory_peaks() noexcept { try { std::scoped_lock lock(execution_mutex_); @@ -198,6 +213,112 @@ class ConcurrentExecutor { } catch (...) {} } + [[nodiscard]] VramControlState vram_control_state() const { + std::scoped_lock lock(execution_mutex_); + return make_vram_control_state(); + } + + void vram_release(const std::vector& tiers, std::size_t target_mib) { + std::scoped_lock lock(execution_mutex_); + bool seed = false; + for (const std::string& tier : tiers) { + if (tier == "seed") { + seed = true; + } else if (tier == "kv") { + throw std::invalid_argument( + "KV-tier VRAM release is not available in this phase; it needs quiesce and " + "CUDA Graph recapture"); + } else { + throw std::invalid_argument("unknown VRAM tier: " + tier); + } + } + if (!seed) { return; } + const std::size_t held = instance_.program->prefix_seed_held_bytes(); + const std::size_t min_bytes = + prefix_cache_min_bytes_ > prefix_cache_max_bytes_ ? 0 : prefix_cache_min_bytes_; + if (min_bytes > 0) { + throw std::invalid_argument("prefix-cache min is not zero; seed store cannot be released"); + } + (void)target_mib; + if (held == 0) { + last_transition_ = "noop"; + last_reason_ = "seed store already released"; + return; + } + if (vram_observe_only_) { + last_transition_ = "observe"; + last_reason_ = "would release seed store (" + std::to_string(held) + " bytes)"; + std::fprintf(stderr, "ninfer: vram observe-only: would release seed %zu MiB\n", + held >> 20); + return; + } + instance_.program->release_prefix_seeds(); + for (std::uint32_t lane = 0; lane < max_concurrency_; ++lane) { + invalidate_lane_plans(lane); + } + last_transition_ = "release"; + last_reason_ = "seed store released"; + std::fprintf(stderr, "ninfer: vram released seed store (%zu MiB)\n", held >> 20); + } + + void vram_reclaim() { + std::scoped_lock lock(execution_mutex_); + const std::size_t held = instance_.program->prefix_seed_held_bytes(); + if (held != 0) { + last_transition_ = "noop"; + last_reason_ = "seed store already held"; + return; + } + if (prefix_cache_max_bytes_ == 0) { + last_transition_ = "noop"; + last_reason_ = "seed store disabled at startup"; + return; + } + if (vram_observe_only_) { + last_transition_ = "observe"; + last_reason_ = "would reclaim seed store"; + std::fprintf(stderr, "ninfer: vram observe-only: would reclaim seed %zu MiB\n", + prefix_cache_max_bytes_ >> 20); + return; + } + const bool ok = instance_.program->reclaim_prefix_seeds(); + for (std::uint32_t lane = 0; lane < max_concurrency_; ++lane) { + invalidate_lane_plans(lane); + } + last_transition_ = ok ? "reclaim" : "reclaim-failed"; + last_reason_ = ok ? "seed store reclaimed" : "seed store reclaim failed; staying degraded"; + std::fprintf(stderr, "ninfer: vram %s\n", last_reason_.c_str()); + } + +private: + [[nodiscard]] VramControlState make_vram_control_state() const { + VramControlState out; + out.observe_only = vram_observe_only_; + out.floor_bytes = vram_floor_bytes_; + out.last_transition = last_transition_; + out.last_reason = last_reason_; + const std::size_t seed_held = instance_.program->prefix_seed_held_bytes(); + const std::size_t seed_max = prefix_cache_max_bytes_; + const std::size_t seed_min = prefix_cache_min_bytes_; + VramTierState seed; + seed.name = "seed"; + seed.held_bytes = seed_held; + seed.min_bytes = seed_min; + seed.max_bytes = seed_max; + seed.reclaimable_bytes = seed_held > seed_min ? seed_held - seed_min : 0; + seed.released = seed_max != 0 && seed_held == 0; + out.tiers.push_back(std::move(seed)); + VramTierState kv; + kv.name = "kv"; + kv.held_bytes = instance_.kv_capacity_resolution.runtime_reservation_bytes; + kv.min_bytes = 0; + kv.max_bytes = 0; + kv.reclaimable_bytes = 0; + kv.released = false; + out.tiers.push_back(std::move(kv)); + return out; + } + private: void publish_runtime_stats() { RuntimeStats snapshot = cumulative_stats_; @@ -823,6 +944,17 @@ class ConcurrentExecutor { const bool cancel_at_boundary = request->cancelled.load(std::memory_order_acquire); resolve_prefill_step(request, first, cancel_at_boundary); publish_runtime_stats(); + } catch (const RequestError&) { + if (target_started) { instance_.program->abort_lane(lane); } + if (prefill_lane_ && *prefill_lane_ == lane) { + instance_.request_memory.deactivate(); + prefill_lane_.reset(); + } + slots_[lane].reset(); + invalidate_lane_plans(lane); + complete_error(request, std::current_exception()); + publish_runtime_stats(); + return AdmissionProgress::ControlProgress; } catch (...) { const std::exception_ptr error = std::current_exception(); if (target_started) { instance_.program->abort_lane(lane); } @@ -869,7 +1001,7 @@ class ConcurrentExecutor { try { ensure_base_plan(head); - } catch (...) { + } catch (const RequestError&) { (void)remove_pending_error(head, std::current_exception()); control_progress = true; continue; @@ -887,7 +1019,7 @@ class ConcurrentExecutor { std::optional head_lane; try { head_lane = find_admission_lane(head); - } catch (...) { + } catch (const RequestError&) { (void)remove_pending_error(head, std::current_exception()); control_progress = true; continue; @@ -902,8 +1034,8 @@ class ConcurrentExecutor { } if (!protection_) { protection_.emplace(make_admission_protection(next_protection_epoch_++, head->id, - head_base.admission, active.span(), - admission_capacity_)); + head_base.admission, active.span(), + admission_capacity_)); } if (protected_head_safe_without_temporal(*protection_, active.span(), admission_capacity_)) { @@ -937,7 +1069,7 @@ class ConcurrentExecutor { try { ensure_base_plan(candidate); - } catch (...) { + } catch (const RequestError&) { (void)remove_pending_error(candidate, std::current_exception()); control_progress = true; continue; @@ -955,7 +1087,7 @@ class ConcurrentExecutor { std::optional candidate_lane; try { candidate_lane = find_admission_lane(candidate); - } catch (...) { + } catch (const RequestError&) { (void)remove_pending_error(candidate, std::current_exception()); control_progress = true; continue; @@ -1063,6 +1195,8 @@ class ConcurrentExecutor { publish_runtime_stats(); } + // Fail all active and pending requests. When followed by std::_Exit(1), notification + // of in-flight futures is best-effort before the process terminates. void fail_all(std::exception_ptr error) noexcept { std::vector> pending; { @@ -1144,17 +1278,55 @@ class ConcurrentExecutor { continue; } } catch (...) { - fail_all(std::current_exception()); + const std::exception_ptr error = std::current_exception(); + fail_all(error); + handle_fatal_error(error); return; } } } + void handle_fatal_error(std::exception_ptr error) noexcept { + std::string detail = "unknown fatal failure"; + if (error != nullptr) { + try { + std::rethrow_exception(error); + } catch (const std::exception& e) { + detail = std::string(typeid(e).name()) + ": " + e.what(); + } catch (...) { + detail = "non-std exception"; + } + } + const std::string message = + "fatal executor failure (" + detail + "); terminating process"; + if (on_fatal_error_) { + try { + on_fatal_error_(error, message); + } catch (...) {} + // When a custom callback is provided (e.g. GenerationService in production, or a + // test probe), the callback owns what action to take (logging and exiting hard + // in serve, or recording without exit in tests). + return; + } + // Default fallback when no callback is registered: log to stderr and exit hard. + std::cerr << message << '\n'; + std::cerr.flush(); + std::cout.flush(); + std::_Exit(1); + } + Instance& instance_; const std::uint32_t max_concurrency_; const std::size_t max_outstanding_; const std::chrono::milliseconds pending_timeout_; const AdmissionResources admission_capacity_; + const std::size_t prefix_cache_min_bytes_; + const std::size_t prefix_cache_max_bytes_; + const std::size_t vram_floor_bytes_; + const bool vram_observe_only_; + std::function on_fatal_error_; + std::string last_transition_; + std::string last_reason_; mutable std::mutex execution_mutex_; mutable std::mutex queue_mutex_; diff --git a/src/runtime/engine/engine.cpp b/src/runtime/engine/engine.cpp index 0f023dd575..5cf09c5950 100644 --- a/src/runtime/engine/engine.cpp +++ b/src/runtime/engine/engine.cpp @@ -346,6 +346,20 @@ RuntimeStats Engine::runtime_stats() const { impl_->executor); } +bool Engine::is_healthy() const noexcept { + if (impl_ == nullptr) { return false; } + return std::visit( + [](const auto& executor) -> bool { + using Executor = std::remove_cvref_t; + if constexpr (std::is_same_v) { + return false; + } else { + return executor != nullptr && executor->is_healthy(); + } + }, + impl_->executor); +} + void Engine::reset_memory_peaks() noexcept { if (impl_ == nullptr) { return; } std::visit( @@ -358,4 +372,46 @@ void Engine::reset_memory_peaks() noexcept { impl_->executor); } +VramControlState Engine::vram_control_state() const { + if (impl_ == nullptr) { throw std::logic_error("Engine is moved from"); } + return std::visit( + [](const auto& executor) -> VramControlState { + using Executor = std::remove_cvref_t; + if constexpr (std::is_same_v) { + throw std::logic_error("concurrent Engine executor is unavailable"); + } else { + return executor->vram_control_state(); + } + }, + impl_->executor); +} + +void Engine::vram_release(const std::vector& tiers, std::size_t target_mib) { + if (impl_ == nullptr) { throw std::logic_error("Engine is moved from"); } + std::visit( + [&](auto& executor) { + using Executor = std::remove_cvref_t; + if constexpr (std::is_same_v) { + throw std::logic_error("concurrent Engine executor is unavailable"); + } else { + executor->vram_release(tiers, target_mib); + } + }, + impl_->executor); +} + +void Engine::vram_reclaim() { + if (impl_ == nullptr) { throw std::logic_error("Engine is moved from"); } + std::visit( + [&](auto& executor) { + using Executor = std::remove_cvref_t; + if constexpr (std::is_same_v) { + throw std::logic_error("concurrent Engine executor is unavailable"); + } else { + executor->vram_reclaim(); + } + }, + impl_->executor); +} + } // namespace ninfer diff --git a/src/serve/anthropic_schema.cpp b/src/serve/anthropic_schema.cpp index ba3c6e26b6..763e2e10c6 100644 --- a/src/serve/anthropic_schema.cpp +++ b/src/serve/anthropic_schema.cpp @@ -8,6 +8,7 @@ #include #include #include +#include #include namespace ninfer::serve { @@ -141,6 +142,7 @@ void parse_tools(const Json& body, GenerationRequest& out) { const Json& tools = body.at("tools"); if (!tools.is_array()) { bad_request("tools must be an array", "tools"); } out.tools.reserve(tools.size()); + std::unordered_set names; for (const Json& item : tools) { if (!item.is_object()) { bad_request("tools entries must be objects", "tools"); } // Anthropic server/built-in tools carry a `type` and no `input_schema`; we @@ -155,6 +157,9 @@ void parse_tools(const Json& body, GenerationRequest& out) { } ToolDefinition tool; tool.name = require_function_name(item, "tools"); + if (!names.insert(tool.name).second) { + bad_request("duplicate tool name: " + tool.name, "tools"); + } Json function = Json{{"name", tool.name}}; if (item.contains("description") && !item.at("description").is_null()) { if (!item.at("description").is_string()) { diff --git a/src/serve/console_log.cpp b/src/serve/console_log.cpp index 7c58004793..24588fd1b5 100644 --- a/src/serve/console_log.cpp +++ b/src/serve/console_log.cpp @@ -38,7 +38,11 @@ std::string format_console_log_prefix(std::chrono::system_clock::time_point time const std::time_t wall_seconds = std::chrono::system_clock::to_time_t(std::chrono::system_clock::time_point(whole_seconds)); std::tm local{}; +#if defined(_WIN32) + ::localtime_s(&local, &wall_seconds); +#else localtime_r(&wall_seconds, &local); +#endif std::ostringstream out; out << '[' << std::put_time(&local, "%Y-%m-%d %H:%M:%S") << '.' << std::setfill('0') diff --git a/src/serve/generation_service.cpp b/src/serve/generation_service.cpp index b1c78d1c73..d84f65f7a3 100644 --- a/src/serve/generation_service.cpp +++ b/src/serve/generation_service.cpp @@ -8,9 +8,12 @@ #include #include #include +#include +#include #include #include #include +#include #include namespace ninfer::serve { @@ -189,13 +192,13 @@ void check_preparation_control(Clock::time_point deadline, class ServiceOutputSink final : public ninfer::OutputSink { public: - ServiceOutputSink(const StreamSink& sink, bool filter_tool_calls) - : sink_(&sink), filter_tool_calls_(filter_tool_calls) {} + ServiceOutputSink(const StreamSink& sink, bool filter_tool_calls, bool hold_reasoning) + : sink_(&sink), filter_tool_calls_(filter_tool_calls), hold_reasoning_(hold_reasoning) {} void publish(ninfer::OutputDelta delta) override { if (delta.text.empty()) { return; } if (delta.channel == ninfer::OutputChannel::Reasoning) { - if (sink_->on_reasoning) { sink_->on_reasoning(delta.text); } + if (!hold_reasoning_ && sink_->on_reasoning) { sink_->on_reasoning(delta.text); } } else { std::string visible = filter_tool_calls_ ? tool_filter_.feed(delta.text) : std::move(delta.text); @@ -217,6 +220,7 @@ class ServiceOutputSink final : public ninfer::OutputSink { const StreamSink* sink_ = nullptr; bool filter_tool_calls_ = false; + bool hold_reasoning_ = false; ToolCallStreamFilter tool_filter_; std::size_t content_bytes_ = 0; }; @@ -241,14 +245,31 @@ GenerationService::GenerationService(ServeOptions options, LoadProgress load_pro engine_options.media_cache_bytes = options_.media_cache_bytes; engine_options.media_live_bytes = options_.media_live_bytes; engine_options.media_preprocess_threads = options_.media_preprocess_threads; - engine_options.load_progress = std::move(load_progress); + engine_options.prefix_cache_bytes = options_.prefix_cache_bytes; + engine_options.prefix_cache_min_bytes = options_.prefix_cache_min_bytes; + engine_options.prefix_cache_max_bytes = options_.prefix_cache_max_bytes; + engine_options.kv_capacity_min_tokens = options_.kv_capacity_min_tokens; + engine_options.kv_capacity_max_tokens = options_.kv_capacity_max_tokens; + engine_options.vram_guarantee_context = options_.vram_guarantee_context; + engine_options.vram_guarantee_concurrency = options_.vram_guarantee_concurrency; + engine_options.vram_floor_bytes = options_.vram_floor_bytes; + engine_options.vram_idle_release_after_s = options_.vram_idle_release_after_s; + engine_options.vram_observe_only = options_.vram_observe_only; + engine_options.load_progress = std::move(load_progress); + engine_options.on_fatal_error = [](std::exception_ptr, const std::string& message) { + write_console_log(ConsoleLogLevel::Error, message); + std::cerr.flush(); + std::cout.flush(); + std::_Exit(1); + }; engine_ = std::make_unique(std::move(engine_options)); prompt_capabilities_ = engine_->prompt_capabilities(); request_capacity_ = std::make_shared( static_cast(options_.max_concurrency) + options_.max_pending_requests); } -std::shared_ptr GenerationService::acquire_request_lifetime() const { +std::shared_ptr GenerationService::acquire_request_lifetime( + std::optional timeout_override) const { const auto started = Clock::now(); { std::lock_guard lock(request_capacity_->mutex); @@ -258,10 +279,10 @@ std::shared_ptr GenerationService::acquire_request_lifetime() c } ++request_capacity_->active; } + const auto timeout = + timeout_override.value_or(std::chrono::milliseconds(options_.pending_timeout_ms)); try { - return std::make_shared( - request_capacity_, started, - started + std::chrono::milliseconds(options_.pending_timeout_ms)); + return std::make_shared(request_capacity_, started, started + timeout); } catch (...) { std::lock_guard lock(request_capacity_->mutex); --request_capacity_->active; @@ -269,13 +290,15 @@ std::shared_ptr GenerationService::acquire_request_lifetime() c } } -PreparedRequest GenerationService::prepare(const GenerationRequest& request, - std::function is_cancelled) const { +PreparedRequest GenerationService::prepare( + const GenerationRequest& request, std::function is_cancelled, + std::optional timeout_override) const { PreparedRequest prepared; ninfer::RequestOptions request_options = to_request_options(request, options_); prepared.include_usage = request.include_usage; prepared.tool_capable = request.uses_tools() || request.has_tool_history(); prepared.tool_name_max_length = request.tool_name_max_length; + prepared.param_types = build_tool_param_type_map(request.tools); const ResolvedPromptSemantics semantics = resolve_prompt_semantics(request, options_, prompt_capabilities_); prepared.enable_thinking = semantics.enable_thinking; @@ -286,7 +309,7 @@ PreparedRequest GenerationService::prepare(const GenerationRequest& request, const std::invalid_argument error("Vision is disabled for this server"); throw_invalid_input(error, "vision_disabled"); } - prepared.lifetime = acquire_request_lifetime(); + prepared.lifetime = acquire_request_lifetime(timeout_override); try { const auto acquisition_started = Clock::now(); @@ -355,7 +378,9 @@ GenerationOutcome GenerationService::run(PreparedRequest& prepared, const Stream std::function is_cancelled) { std::unique_ptr output_sink; if (sink != nullptr) { - output_sink = std::make_unique(*sink, prepared.tool_capable); + output_sink = std::make_unique( + *sink, prepared.tool_capable, + prepared.tool_capable && options_.tolerant_tool_calls); } ninfer::OutputSink* public_sink = output_sink.get(); ninfer::CancellationView cancellation; @@ -401,15 +426,35 @@ GenerationOutcome GenerationService::run(PreparedRequest& prepared, const Stream bool is_tool_call_response = false; if (prepared.tool_capable) { - ParsedToolCallOutput parsed = - parse_qwen_tool_call_output(outcome.text, prepared.tool_name_max_length); + ParsedToolCallOutput parsed = parse_qwen_tool_call_output( + outcome.text, prepared.tool_name_max_length, prepared.param_types, + options_.tolerant_tool_calls); outcome.text = std::move(parsed.content); is_tool_call_response = parsed.is_tool_call_response; - if (is_tool_call_response) { outcome.tool_calls = std::move(parsed.tool_calls); } + if (is_tool_call_response) { + outcome.tool_calls = std::move(parsed.tool_calls); + } else if (options_.tolerant_tool_calls && !outcome.reasoning.empty()) { + // A Qwen drift can emit the call before . In that case the + // frontend correctly classifies it as reasoning, so give the same + // tolerant recovery path a chance before returning raw XML. + ParsedToolCallOutput reasoning_parsed = parse_qwen_tool_call_output( + outcome.reasoning, prepared.tool_name_max_length, prepared.param_types, true); + if (reasoning_parsed.is_tool_call_response) { + outcome.reasoning = std::move(reasoning_parsed.content); + outcome.tool_calls = std::move(reasoning_parsed.tool_calls); + is_tool_call_response = true; + } + } } if (output_sink) { outcome.streamed_content_bytes = output_sink->finish(is_tool_call_response); } + // A tool turn should not show the model's pre-call chatter to the caller, but it + // can only be dropped when none of it has already gone out on the wire. Streaming + // recognises only once the whole marker lands in one chunk; when the + // marker straddles chunks the prefix has already been emitted, and shortening the + // terminal body below streamed_content_bytes would abort the request mid-stream. + if (is_tool_call_response && outcome.streamed_content_bytes == 0) { outcome.text.clear(); } return outcome; } @@ -426,11 +471,14 @@ void GenerationService::warmup() { request.messages.push_back(std::move(turn)); request.max_tokens = 4; request.max_tokens_set = true; - PreparedRequest prepared = prepare(request); + // Warmup is internal startup priming and must not inherit the client-facing + // request deadline (--pending-timeout-ms bounds incoming-request preparation + // and queue waiting, not warmup). + constexpr auto kWarmupTimeout = std::chrono::seconds(60); + PreparedRequest prepared = prepare(request, {}, kWarmupTimeout); run(prepared, nullptr); } catch (const std::exception& exception) { - write_console_log(ConsoleLogLevel::Warning, - std::string("warmup failed (continuing): ") + exception.what()); + throw std::runtime_error(std::string("warmup generation failed: ") + exception.what()); } } diff --git a/src/serve/generation_service.h b/src/serve/generation_service.h index a444146c20..8df43b9584 100644 --- a/src/serve/generation_service.h +++ b/src/serve/generation_service.h @@ -7,12 +7,14 @@ #include "ninfer/engine.h" #include "serve/request.h" #include "serve/serve_options.h" +#include "serve/tool_call_parser.h" #include #include #include #include #include +#include #include #include @@ -74,6 +76,7 @@ struct PreparedRequest { bool include_usage = false; bool tool_capable = false; std::size_t tool_name_max_length = 64; + ToolParamTypeMap param_types; bool enable_thinking = true; bool preserve_thinking = false; bool preserve_thinking_semantic_change = false; @@ -96,12 +99,24 @@ class GenerationService { return engine_->media_cache_summary(); } + [[nodiscard]] ninfer::VramControlState vram_control_state() const { + return engine_->vram_control_state(); + } + void vram_release(const std::vector& tiers, std::size_t target_mib = 0) { + engine_->vram_release(tiers, target_mib); + } + void vram_reclaim() { engine_->vram_reclaim(); } + [[nodiscard]] bool is_healthy() const noexcept { + return engine_ != nullptr && engine_->is_healthy(); + } + [[nodiscard]] ninfer::ModelSamplingDefaults sampling_defaults() const { return engine_->sampling_defaults(); } - [[nodiscard]] PreparedRequest prepare(const GenerationRequest& req, - std::function is_cancelled = {}) const; + [[nodiscard]] PreparedRequest prepare( + const GenerationRequest& req, std::function is_cancelled = {}, + std::optional timeout_override = std::nullopt) const; [[nodiscard]] int count_prompt_tokens(const GenerationRequest& req, std::function is_cancelled = {}) const; @@ -112,7 +127,8 @@ class GenerationService { void warmup(); private: - [[nodiscard]] std::shared_ptr acquire_request_lifetime() const; + [[nodiscard]] std::shared_ptr acquire_request_lifetime( + std::optional timeout_override = std::nullopt) const; ServeOptions options_; std::unique_ptr engine_; diff --git a/src/serve/http_server.cpp b/src/serve/http_server.cpp index f99e6d9021..6f553421ca 100644 --- a/src/serve/http_server.cpp +++ b/src/serve/http_server.cpp @@ -10,13 +10,16 @@ #include #include +#include #include #include #include #include +#include #include #include #include +#include namespace ninfer::serve { namespace { @@ -54,6 +57,15 @@ void write_error(httplib::Response& res, const ApiError& error) { res.set_content(make_error_body(error), "application/json"); } +CompletionUsage completion_usage(const GenerationOutcome& outcome) { + CompletionUsage usage; + usage.prompt_tokens = outcome.prompt_tokens; + usage.completion_tokens = outcome.completion_tokens; + usage.cached_prompt_tokens = static_cast(outcome.metrics.prefix_cache_hit_tokens); + usage.prefix_reuse_path = outcome.metrics.prefix_reuse_path; + return usage; +} + // Anthropic-shaped error body ({"type":"error","error":{...}}), used by the // /v1/messages endpoints so Claude clients see the error format they expect. void write_messages_error(httplib::Response& res, const ApiError& error) { @@ -172,19 +184,44 @@ void HttpServer::run_stats_reporter() { using Clock = std::chrono::steady_clock; ninfer::RuntimeStats previous = service_->runtime_stats(); Clock::time_point previous_time = Clock::now(); - const auto interval = std::chrono::milliseconds(options_.log_stats_interval_ms); + const auto stats_interval = std::chrono::milliseconds( + options_.log_stats_interval_ms == 0 ? 1000 : options_.log_stats_interval_ms); + std::optional idle_since; for (;;) { { std::unique_lock lock(stats_mutex_); - if (stats_cv_.wait_for(lock, interval, [this] { return stats_stopping_; })) { break; } + if (stats_cv_.wait_for(lock, stats_interval, [this] { return stats_stopping_; })) { + break; + } } const ninfer::RuntimeStats current = service_->runtime_stats(); const Clock::time_point now = Clock::now(); - const ThroughputReport report = make_throughput_report( - previous, current, std::chrono::duration(now - previous_time).count()); - if (report_has_activity(report)) { log_throughput(report); } + if (options_.log_stats_interval_ms != 0) { + const ThroughputReport report = make_throughput_report( + previous, current, std::chrono::duration(now - previous_time).count()); + if (report_has_activity(report)) { log_throughput(report); } + } + if (options_.vram_idle_release_after_s != 0) { + const bool gpu_idle = current.running_requests == 0 && + current.prefilling_requests == 0 && current.waiting_requests == 0; + if (!gpu_idle) { + idle_since.reset(); + } else { + if (!idle_since.has_value()) { idle_since = now; } + const auto idle_for = std::chrono::duration_cast( + now - *idle_since); + if (idle_for.count() >= static_cast(options_.vram_idle_release_after_s)) { + try { + service_->vram_release({"seed"}, 0); + } catch (const std::exception& error) { + log_line(std::string("vram idle release skipped: ") + error.what()); + } + idle_since = now; + } + } + } previous = current; previous_time = now; } @@ -225,7 +262,17 @@ void HttpServer::register_routes() { } server_.set_pre_routing_handler([this](const httplib::Request& req, httplib::Response& res) { - if (options_.api_key.empty() || req.path == "/health" || req.method == "OPTIONS") { + if (req.path.rfind("/admin", 0) == 0) { + if (options_.api_key.empty()) { + ApiError error; + error.status = 403; + error.type = "permission_error"; + error.code = "admin_disabled"; + error.message = "admin endpoints require --admin-vram and --api-key"; + write_error(res, error); + return httplib::Server::HandlerResponse::Handled; + } + } else if (options_.api_key.empty() || req.path == "/health" || req.method == "OPTIONS") { return httplib::Server::HandlerResponse::Unhandled; } // Accept both the OpenAI-style bearer token and the Anthropic-style @@ -266,9 +313,30 @@ void HttpServer::register_routes() { } }); - server_.Get("/health", [](const httplib::Request&, httplib::Response& res) { + server_.Get("/health", [this](const httplib::Request&, httplib::Response& res) { + if (service_ != nullptr && !service_->is_healthy()) { + res.status = 503; + res.set_content( + nlohmann::json{{"status", "unhealthy"}, + {"error", "inference engine is unavailable"}} + .dump(), + "application/json"); + return; + } res.set_content(nlohmann::json{{"status", "ok"}}.dump(), "application/json"); }); + if (options_.enable_admin_vram) { + server_.Get("/admin/vram", [this](const httplib::Request&, httplib::Response& res) { + handle_admin_vram(res); + }); + server_.Post("/admin/vram/release", + [this](const httplib::Request& req, httplib::Response& res) { + handle_admin_vram_release(req, res); + }); + server_.Post("/admin/vram/reclaim", [this](const httplib::Request&, httplib::Response& res) { + handle_admin_vram_reclaim(res); + }); + } server_.Get("/v1/models", [this](const httplib::Request& req, httplib::Response& res) { handle_models(req, res); }); @@ -316,7 +384,8 @@ void HttpServer::register_routes() { } void HttpServer::handle_models(const httplib::Request&, httplib::Response& res) const { - res.set_content(make_models_list(public_model_id_, unix_time_now()), "application/json"); + res.set_content(make_models_list(public_model_id_, unix_time_now(), options_.max_context), + "application/json"); } void HttpServer::handle_model(const httplib::Request& req, httplib::Response& res) const { @@ -330,7 +399,89 @@ void HttpServer::handle_model(const httplib::Request& req, httplib::Response& re write_error(res, error); return; } - res.set_content(make_model_object(public_model_id_, unix_time_now()), "application/json"); + res.set_content(make_model_object(public_model_id_, unix_time_now(), options_.max_context), + "application/json"); +} + +nlohmann::json vram_state_json(const ninfer::VramControlState& state) { + nlohmann::json tiers = nlohmann::json::array(); + for (const auto& tier : state.tiers) { + tiers.push_back({{"name", tier.name}, + {"held_bytes", tier.held_bytes}, + {"min_bytes", tier.min_bytes}, + {"max_bytes", tier.max_bytes}, + {"reclaimable_bytes", tier.reclaimable_bytes}, + {"released", tier.released}}); + } + return {{"tiers", std::move(tiers)}, + {"floor_bytes", state.floor_bytes}, + {"observe_only", state.observe_only}, + {"last_transition", state.last_transition}, + {"last_reason", state.last_reason}}; +} + +void HttpServer::handle_admin_vram(httplib::Response& res) const { + res.set_content(vram_state_json(service_->vram_control_state()).dump(), "application/json"); +} + +void HttpServer::handle_admin_vram_release(const httplib::Request& req, httplib::Response& res) { + std::vector tiers{"seed"}; + std::size_t target_mib = 0; + if (!req.body.empty()) { + nlohmann::json body; + try { + body = nlohmann::json::parse(req.body); + } catch (const std::exception&) { + ApiError error; + error.status = 400; + error.type = "invalid_request_error"; + error.code = "invalid_json"; + error.message = "request body is not valid JSON"; + write_error(res, error); + return; + } + if (body.contains("tiers")) { + if (!body.at("tiers").is_array()) { + ApiError error; + error.status = 400; + error.type = "invalid_request_error"; + error.message = "tiers must be an array"; + write_error(res, error); + return; + } + tiers.clear(); + for (const auto& item : body.at("tiers")) { + if (!item.is_string()) { + ApiError error; + error.status = 400; + error.type = "invalid_request_error"; + error.message = "tiers entries must be strings"; + write_error(res, error); + return; + } + tiers.push_back(item.get()); + } + } + if (body.contains("target_mib") && body.at("target_mib").is_number_unsigned()) { + target_mib = body.at("target_mib").get(); + } + } + try { + service_->vram_release(tiers, target_mib); + } catch (const std::invalid_argument& error) { + ApiError api; + api.status = 400; + api.type = "invalid_request_error"; + api.message = error.what(); + write_error(res, api); + return; + } + res.set_content(vram_state_json(service_->vram_control_state()).dump(), "application/json"); +} + +void HttpServer::handle_admin_vram_reclaim(httplib::Response& res) { + service_->vram_reclaim(); + res.set_content(vram_state_json(service_->vram_control_state()).dump(), "application/json"); } void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::Response& res) { @@ -398,7 +549,7 @@ void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::R return req.is_connection_alive && !req.is_connection_alive(); }); log_request_done(log_context, outcome); - const CompletionUsage usage{outcome.prompt_tokens, outcome.completion_tokens}; + const CompletionUsage usage = completion_usage(outcome); std::string response_body; if (!outcome.tool_calls.empty()) { response_body = make_chat_completion_tool_response( @@ -416,9 +567,10 @@ void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::R return; } - auto stream = std::make_shared(std::move(prepared)); - const bool include_usage = stream->prepared.include_usage; - const bool tool_capable = stream->prepared.tool_capable; + auto stream = std::make_shared(std::move(prepared)); + const bool include_usage = stream->prepared.include_usage; + const bool tool_capable = stream->prepared.tool_capable; + const bool buffer_reasoning = tool_capable && options_.tolerant_tool_calls; // SSE hints: disable client/proxy caching and reverse-proxy response buffering // so tokens flush immediately. Content-Type is set by the chunked provider. @@ -427,7 +579,7 @@ void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::R res.set_chunked_content_provider( "text/event-stream", - [this, stream, id, created, model, include_usage, tool_capable, + [this, stream, id, created, model, include_usage, tool_capable, buffer_reasoning, log_context](std::size_t, httplib::DataSink& sink) -> bool { if (stream->started) { sink.done(); @@ -455,6 +607,11 @@ void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::R const GenerationOutcome outcome = service_->run(stream->prepared, &output); log_request_done(log_context, outcome); + if (buffer_reasoning && !outcome.reasoning.empty()) { + write_stream_item(sink, *stream, + make_chat_chunk_reasoning(id, model, created, + outcome.reasoning, include_usage)); + } const std::string_view remaining = unstreamed_content(outcome); if (!outcome.tool_calls.empty()) { if (!remaining.empty()) { @@ -483,7 +640,7 @@ void HttpServer::handle_chat_completions(const httplib::Request& req, httplib::R include_usage)); } if (include_usage) { - const CompletionUsage usage{outcome.prompt_tokens, outcome.completion_tokens}; + const CompletionUsage usage = completion_usage(outcome); write_stream_item(sink, *stream, make_chat_chunk_usage(id, model, created, usage)); } @@ -611,7 +768,7 @@ void HttpServer::handle_messages(const httplib::Request& req, httplib::Response& return req.is_connection_alive && !req.is_connection_alive(); }); log_request_done(log_context, outcome); - const CompletionUsage usage{outcome.prompt_tokens, outcome.completion_tokens}; + const CompletionUsage usage = completion_usage(outcome); const char* stop_reason = messages_stop_reason(outcome.finish_reason, !outcome.tool_calls.empty()); set_owned_content(res, @@ -632,15 +789,16 @@ void HttpServer::handle_messages(const httplib::Request& req, httplib::Response& return; } - auto stream = std::make_shared(std::move(prepared)); - const bool tool_capable = stream->prepared.tool_capable; + auto stream = std::make_shared(std::move(prepared)); + const bool tool_capable = stream->prepared.tool_capable; + const bool buffer_reasoning = tool_capable && options_.tolerant_tool_calls; res.set_header("Cache-Control", "no-cache"); res.set_header("X-Accel-Buffering", "no"); res.set_chunked_content_provider( "text/event-stream", - [this, stream, id, model, input_tokens, tool_capable, + [this, stream, id, model, input_tokens, tool_capable, buffer_reasoning, log_context](std::size_t, httplib::DataSink& sink) -> bool { if (stream->started) { sink.done(); @@ -698,6 +856,14 @@ void HttpServer::handle_messages(const httplib::Request& req, httplib::Response& text_open = false; } + if (buffer_reasoning && !outcome.reasoning.empty()) { + const int idx = next_index++; + write_stream_item(sink, *stream, make_content_block_start_thinking(idx)); + write_stream_item(sink, *stream, + make_content_block_delta_thinking(idx, outcome.reasoning)); + write_stream_item(sink, *stream, make_content_block_stop(idx)); + } + if (tool_capable) { if (!remaining.empty()) { const int idx = next_index++; @@ -775,7 +941,7 @@ bool HttpServer::listen() { if (public_model_id_.empty()) { throw std::logic_error("HTTP public model id is not resolved"); } - if (options_.log_stats_interval_ms != 0) { + if (options_.log_stats_interval_ms != 0 || options_.vram_idle_release_after_s != 0) { stats_stopping_ = false; stats_thread_ = std::thread([this] { run_stats_reporter(); }); } diff --git a/src/serve/http_server.h b/src/serve/http_server.h index c0557e9d27..fa77bd0065 100644 --- a/src/serve/http_server.h +++ b/src/serve/http_server.h @@ -51,6 +51,9 @@ class HttpServer { void handle_response_compact(const httplib::Request& req, httplib::Response& res); void handle_models(const httplib::Request& req, httplib::Response& res) const; void handle_model(const httplib::Request& req, httplib::Response& res) const; + void handle_admin_vram(httplib::Response& res) const; + void handle_admin_vram_release(const httplib::Request& req, httplib::Response& res); + void handle_admin_vram_reclaim(httplib::Response& res); // The process-wide console logger serializes lines from request and reporter threads. void log_line(const std::string& line); diff --git a/src/serve/openai_schema.cpp b/src/serve/openai_schema.cpp index 6975ccf7d8..426c5f700c 100644 --- a/src/serve/openai_schema.cpp +++ b/src/serve/openai_schema.cpp @@ -1,5 +1,6 @@ #include "serve/openai_schema.h" +#include #include #include #include @@ -8,6 +9,7 @@ #include #include #include +#include namespace ninfer::serve { namespace { @@ -136,7 +138,10 @@ ninfer::product::media_acquire::Source parse_media_url(const Json& part, const c return source; } -void parse_content_parts(const Json& content, ChatTurn& turn, std::size_t index) { +} // namespace + +void parse_content_parts(const Json& content, ChatTurn& turn, std::size_t index, + std::vector allowed_types) { if (content.is_string()) { turn.content.push_back(ContentPart{ContentKind::Text, content.get(), "text"}); return; @@ -145,6 +150,11 @@ void parse_content_parts(const Json& content, ChatTurn& turn, std::size_t index) bad_request("message " + std::to_string(index) + " content must be a string or array", "messages"); } + std::string allowed_list; + for (const std::string& allowed : allowed_types) { + if (!allowed_list.empty()) { allowed_list += ", "; } + allowed_list += "'" + allowed + "'"; + } for (const Json& part : content) { if (!part.is_object() || !part.contains("type") || !part.at("type").is_string()) { bad_request("message " + std::to_string(index) + @@ -152,6 +162,12 @@ void parse_content_parts(const Json& content, ChatTurn& turn, std::size_t index) "messages"); } const std::string type = part.at("type").get(); + if (!allowed_types.empty() && + std::find(allowed_types.begin(), allowed_types.end(), type) == allowed_types.end()) { + bad_request("message " + std::to_string(index) + " content parts must have type " + + allowed_list, + "messages"); + } ContentPart out; out.type_raw = type; if (type == "text") { @@ -178,6 +194,8 @@ void parse_content_parts(const Json& content, ChatTurn& turn, std::size_t index) } } +namespace { + std::vector parse_assistant_tool_calls(const Json& item, std::size_t index) { std::vector calls; if (!item.contains("tool_calls") || item.at("tool_calls").is_null()) { return calls; } @@ -250,12 +268,15 @@ void parse_messages(const Json& body, GenerationRequest& out) { item.at("tool_call_id").get().empty()) { bad_request("tool messages must contain a string tool_call_id", "messages"); } - if (!item.contains("content") || !item.at("content").is_string()) { - bad_request("tool messages must contain string content", "messages"); + if (!item.contains("content") || item.at("content").is_null()) { + bad_request("tool messages must contain content", "messages"); } turn.tool_call_id = item.at("tool_call_id").get(); - turn.content.push_back( - ContentPart{ContentKind::Text, item.at("content").get(), "text"}); + // Tool results share the ordinary content grammar: a string, or an array of + // text/image/video parts. Agentic clients return screenshots and rendered pages as + // image parts inside the tool message (the template renders tool turns through the + // same media-placeholder path as user turns). + parse_content_parts(item.at("content"), turn, i); out.messages.push_back(std::move(turn)); continue; } @@ -298,6 +319,7 @@ void parse_tools(const Json& body, GenerationRequest& out) { const Json& tools = body.at("tools"); if (!tools.is_array()) { bad_request("tools must be an array", "tools"); } out.tools.reserve(tools.size()); + std::unordered_set names; for (std::size_t i = 0; i < tools.size(); ++i) { const Json& item = tools.at(i); if (!item.is_object()) { bad_request("tools entries must be objects", "tools"); } @@ -314,6 +336,9 @@ void parse_tools(const Json& body, GenerationRequest& out) { Json& fn = normalized["function"]; ToolDefinition tool; tool.name = require_function_name(fn, "tools"); + if (!names.insert(tool.name).second) { + bad_request("duplicate function tool name: " + tool.name, "tools"); + } if (fn.contains("description") && !fn.at("description").is_null()) { if (!fn.at("description").is_string()) { bad_request("function description must be a string", "tools"); @@ -479,6 +504,39 @@ std::string sse_event(const Json& payload) { return "data: " + payload.dump() + } // namespace +std::optional parse_openai_template_enable_thinking(const Json& body) { + if (!body.contains("chat_template_kwargs")) { return std::nullopt; } + const Json& kwargs = body.at("chat_template_kwargs"); + if (!kwargs.is_object()) { + bad_request("chat_template_kwargs must be an object", "chat_template_kwargs"); + } + if (!kwargs.contains("enable_thinking") || kwargs.at("enable_thinking").is_null()) { + return std::nullopt; + } + if (!kwargs.at("enable_thinking").is_boolean()) { + bad_request("chat_template_kwargs.enable_thinking must be a boolean or null", + "chat_template_kwargs"); + } + return kwargs.at("enable_thinking").get(); +} + +void apply_openai_enable_thinking(const Json& body, GenerationRequest& out) { + std::optional top_level; + if (body.contains("enable_thinking") && !body.at("enable_thinking").is_null()) { + top_level = get_bool(body, "enable_thinking", false); + } + const std::optional template_thinking = parse_openai_template_enable_thinking(body); + if (top_level && template_thinking && *top_level != *template_thinking) { + bad_request("conflicting enable_thinking values", "enable_thinking", + "conflicting_template_option"); + } + if (template_thinking) { + out.enable_thinking = *template_thinking; + } else if (top_level) { + out.enable_thinking = *top_level; + } +} + std::optional parse_openai_preserve_thinking(const Json& body) { std::optional top_level; if (body.contains("preserve_thinking") && !body.at("preserve_thinking").is_null()) { @@ -495,7 +553,8 @@ std::optional parse_openai_preserve_thinking(const Json& body) { bad_request("chat_template_kwargs must be an object", "chat_template_kwargs"); } for (auto it = kwargs.begin(); it != kwargs.end(); ++it) { - if (it.key() != "preserve_thinking" && !it.value().is_null()) { + if (it.key() != "preserve_thinking" && it.key() != "enable_thinking" && + !it.value().is_null()) { bad_request("chat_template_kwargs." + it.key() + " is not supported", "chat_template_kwargs", "chat_template_option_not_supported"); } @@ -553,11 +612,9 @@ GenerationRequest parse_chat_completion_request(const Json& body, const RequestL if (body.contains("stream_options") && body.at("stream_options").is_object()) { out.include_usage = get_bool(body.at("stream_options"), "include_usage", false); } - if (body.contains("enable_thinking") && !body.at("enable_thinking").is_null()) { - out.enable_thinking = get_bool(body, "enable_thinking", false); - } parse_openai_reasoning_effort(body, out); out.preserve_thinking = parse_openai_preserve_thinking(body); + apply_openai_enable_thinking(body, out); std::optional max_tokens = get_int(body, "max_completion_tokens"); if (!max_tokens) { max_tokens = get_int(body, "max_tokens"); } @@ -572,6 +629,21 @@ GenerationRequest parse_chat_completion_request(const Json& body, const RequestL return out; } +namespace { + +nlohmann::json usage_json(const CompletionUsage& usage) { + const int cached = std::clamp(usage.cached_prompt_tokens, 0, usage.prompt_tokens); + return nlohmann::json{ + {"prompt_tokens", usage.prompt_tokens}, + {"completion_tokens", usage.completion_tokens}, + {"total_tokens", usage.prompt_tokens + usage.completion_tokens}, + {"prompt_tokens_details", nlohmann::json{{"cached_tokens", cached}}}, + {"prefix_cache_hit_tokens", cached}, + {"prefix_reuse_path", prefix_reuse_path_name(usage.prefix_reuse_path)}}; +} + +} // namespace + std::string make_chat_completion_response(const std::string& id, const std::string& model, std::int64_t created, const std::string& content, const std::string& reasoning, const char* finish_reason, @@ -586,9 +658,7 @@ std::string make_chat_completion_response(const std::string& id, const std::stri {"choices", Json::array({Json{ {"index", 0}, {"message", std::move(message)}, {"finish_reason", finish_reason}}})}, - {"usage", Json{{"prompt_tokens", usage.prompt_tokens}, - {"completion_tokens", usage.completion_tokens}, - {"total_tokens", usage.prompt_tokens + usage.completion_tokens}}}}; + {"usage", usage_json(usage)}}; return payload.dump(); } @@ -609,9 +679,7 @@ std::string make_chat_completion_tool_response(const std::string& id, const std: {"choices", Json::array({Json{ {"index", 0}, {"message", std::move(message)}, {"finish_reason", "tool_calls"}}})}, - {"usage", Json{{"prompt_tokens", usage.prompt_tokens}, - {"completion_tokens", usage.completion_tokens}, - {"total_tokens", usage.prompt_tokens + usage.completion_tokens}}}}; + {"usage", usage_json(usage)}}; return payload.dump(); } @@ -673,26 +741,30 @@ std::string make_chat_chunk_usage(const std::string& id, const std::string& mode std::int64_t created, const CompletionUsage& usage) { Json payload = base_chunk(id, model, created); payload["choices"] = Json::array(); - payload["usage"] = Json{{"prompt_tokens", usage.prompt_tokens}, - {"completion_tokens", usage.completion_tokens}, - {"total_tokens", usage.prompt_tokens + usage.completion_tokens}}; + payload["usage"] = usage_json(usage); return sse_event(payload); } std::string sse_done() { return "data: [DONE]\n\n"; } -std::string make_models_list(const std::string& model_id, std::int64_t created) { +std::string make_models_list(const std::string& model_id, std::int64_t created, + std::uint32_t max_model_len) { const Json payload = {{"object", "list"}, {"data", Json::array({Json{{"id", model_id}, {"object", "model"}, {"created", created}, - {"owned_by", "ninfer"}}})}}; + {"owned_by", "ninfer"}, + {"max_model_len", max_model_len}}})}}; return payload.dump(); } -std::string make_model_object(const std::string& model_id, std::int64_t created) { - const Json payload = { - {"id", model_id}, {"object", "model"}, {"created", created}, {"owned_by", "ninfer"}}; +std::string make_model_object(const std::string& model_id, std::int64_t created, + std::uint32_t max_model_len) { + const Json payload = {{"id", model_id}, + {"object", "model"}, + {"created", created}, + {"owned_by", "ninfer"}, + {"max_model_len", max_model_len}}; return payload.dump(); } diff --git a/src/serve/openai_schema.h b/src/serve/openai_schema.h index 57f430abf2..d97665958f 100644 --- a/src/serve/openai_schema.h +++ b/src/serve/openai_schema.h @@ -23,6 +23,13 @@ namespace ninfer::serve { GenerationRequest parse_chat_completion_request(const nlohmann::json& body, const RequestLimits& limits); +// Parse a message's `content` field (string or content-part array) into `turn.content`. +// A non-empty `allowed_types` rejects parts whose `type` is not listed. +void parse_content_parts(const nlohmann::json& content, ChatTurn& turn, std::size_t index, + std::vector allowed_types = {}); + +std::optional parse_openai_template_enable_thinking(const nlohmann::json& body); +void apply_openai_enable_thinking(const nlohmann::json& body, GenerationRequest& out); std::optional parse_openai_preserve_thinking(const nlohmann::json& body); // Non-streaming chat completion response body (JSON string). When `reasoning` is @@ -65,9 +72,11 @@ std::string make_chat_chunk_usage(const std::string& id, const std::string& mode std::int64_t created, const CompletionUsage& usage); std::string sse_done(); -// /v1/models payloads. -std::string make_models_list(const std::string& model_id, std::int64_t created); -std::string make_model_object(const std::string& model_id, std::int64_t created); +// /v1/models payloads. max_model_len is the process's configured --max-context. +std::string make_models_list(const std::string& model_id, std::int64_t created, + std::uint32_t max_model_len); +std::string make_model_object(const std::string& model_id, std::int64_t created, + std::uint32_t max_model_len); // Error object body. std::string make_error_body(const ApiError& error); diff --git a/src/serve/request.h b/src/serve/request.h index 3088a6b849..89c8f5887a 100644 --- a/src/serve/request.h +++ b/src/serve/request.h @@ -48,9 +48,30 @@ struct RequestLimits { int default_max_tokens = 8192; }; +[[nodiscard]] constexpr const char* prefix_reuse_path_name(ninfer::PrefixReusePath path) noexcept { + switch (path) { + case ninfer::PrefixReusePath::FullReset: + return "full_reset"; + case ninfer::PrefixReusePath::AppendAtFrontier: + return "append_frontier"; + case ninfer::PrefixReusePath::RestoreTurnCheckpoint: + return "restore_turn_checkpoint"; + case ninfer::PrefixReusePath::RestoreResponseCheckpoint: + return "restore_response_checkpoint"; + case ninfer::PrefixReusePath::SeedPrefixCache: + return "seed_prefix"; + } + return "unknown"; +} + struct CompletionUsage { int prompt_tokens = 0; int completion_tokens = 0; + // Subset of prompt_tokens the prefix cache served (OpenAI prompt_tokens_details.cached_tokens). + // Not an addend: clients subtract it from prompt_tokens. Anthropic Messages does not emit this + // because that surface's input_tokens already excludes cache reads. + int cached_prompt_tokens = 0; + ninfer::PrefixReusePath prefix_reuse_path = ninfer::PrefixReusePath::FullReset; }; enum class ContentKind { diff --git a/src/serve/request_log.cpp b/src/serve/request_log.cpp index b2dd984b69..16aecc96e5 100644 --- a/src/serve/request_log.cpp +++ b/src/serve/request_log.cpp @@ -15,7 +15,11 @@ #include #include +#if defined(_WIN32) +#include +#else #include +#endif namespace ninfer::serve { namespace { @@ -31,8 +35,12 @@ std::uint64_t unix_time_ms() { std::string new_server_instance_id() { const auto now = std::chrono::system_clock::now().time_since_epoch(); const auto micros = std::chrono::duration_cast(now).count(); - return "serve-" + std::to_string(static_cast(::getpid())) + '-' + - std::to_string(micros); +#if defined(_WIN32) + const auto process_id = static_cast(::GetCurrentProcessId()); +#else + const auto process_id = static_cast(::getpid()); +#endif + return "serve-" + std::to_string(process_id) + '-' + std::to_string(micros); } std::filesystem::path normalized_absolute_path(const std::string& value) { @@ -103,20 +111,6 @@ const char* proposal_head_name(ninfer::ProposalHead proposal) { return proposal == ninfer::ProposalHead::Optimized ? "optimized" : "full"; } -const char* prefix_reuse_path_name(ninfer::PrefixReusePath path) { - switch (path) { - case ninfer::PrefixReusePath::FullReset: - return "full_reset"; - case ninfer::PrefixReusePath::AppendAtFrontier: - return "append_frontier"; - case ninfer::PrefixReusePath::RestoreTurnCheckpoint: - return "restore_turn_checkpoint"; - case ninfer::PrefixReusePath::RestoreResponseCheckpoint: - return "restore_response_checkpoint"; - } - return "unknown"; -} - Json event_base(const std::string& server_instance_id, std::uint64_t timestamp, const char* event) { return Json{{"artifact_type", kRequestLogArtifactType}, {"schema_version", kRequestLogSchemaVersion}, @@ -454,7 +448,8 @@ std::string format_server_start_json( {"request_log_jsonl", options.request_log_jsonl}, {"default_output_tokens", options.default_max_tokens}, {"default_thinking", options.enable_thinking}, - {"default_preserve_thinking", options.preserve_thinking}}; + {"default_preserve_thinking", options.preserve_thinking}, + {"tolerant_tool_calls", options.tolerant_tool_calls}}; record["artifact"] = Json{{"path", options.artifact_path}, {"size_bytes", std::move(artifact_size)}, {"target", load.target}, @@ -505,6 +500,7 @@ std::string format_server_start_json( {"planned_slack_bytes", memory.planned_slack_bytes}, {"cuda_graph_allowance_bytes", memory.cuda_graph_allowance_bytes}, {"cuda_graph_observed_bytes", memory.cuda_graph_observed_bytes}, + {"prefix_cache_bytes", memory.prefix_cache_bytes}, {"kv_payload_bytes", memory.kv_payload_bytes}}; record["environment"] = Json{{"device", environment.device}, diff --git a/src/serve/responses_schema.cpp b/src/serve/responses_schema.cpp index 3efdf81a2e..c8610922c5 100644 --- a/src/serve/responses_schema.cpp +++ b/src/serve/responses_schema.cpp @@ -575,6 +575,7 @@ void reject_unknown_top_level(const Json& body) { "chat_template_kwargs", "context_management", "conversation", + "enable_thinking", "include", "input", "instructions", @@ -742,6 +743,7 @@ ResponsesRequest parse_request_impl(const Json& body, const RequestLimits& limit parse_tool_choice(body, out); parse_reasoning(body, out); out.generation.preserve_thinking = parse_openai_preserve_thinking(body); + apply_openai_enable_thinking(body, out.generation); if (const std::optional temperature = optional_number(body, "temperature")) { if (*temperature < 0.0 || *temperature > 2.0) { @@ -913,7 +915,9 @@ BuiltResponse build_response(const std::string& id, std::int64_t created_at, {"input_tokens_details", Json{{"cached_tokens", cached_tokens}}}, {"output_tokens", outcome.completion_tokens}, {"output_tokens_details", Json{{"reasoning_tokens", outcome.reasoning_tokens}}}, - {"total_tokens", outcome.prompt_tokens + outcome.completion_tokens}}; + {"total_tokens", outcome.prompt_tokens + outcome.completion_tokens}, + {"prefix_cache_hit_tokens", cached_tokens}, + {"prefix_reuse_path", prefix_reuse_path_name(outcome.metrics.prefix_reuse_path)}}; built.body = std::move(response); return built; } @@ -946,7 +950,7 @@ ResponsesRequest parse_response_input_tokens_request(const Json& body, require_object(body); for (auto it = body.begin(); it != body.end(); ++it) { if (it.key() != "model" && it.key() != "input" && it.key() != "chat_template_kwargs" && - it.key() != "preserve_thinking") { + it.key() != "preserve_thinking" && it.key() != "enable_thinking") { bad_request("unknown parameter: " + it.key(), it.key(), "unknown_parameter"); } } diff --git a/src/serve/serve_options.cpp b/src/serve/serve_options.cpp index c991e2cc85..12b7c403b8 100644 --- a/src/serve/serve_options.cpp +++ b/src/serve/serve_options.cpp @@ -65,24 +65,38 @@ std::string serve_usage_text(const char* argv0) { return std::string("usage: ") + argv0 + " [--host H] [--port N] [--api-key KEY] " "[--model-id ID] [--max-context N] [--kv-capacity N|auto] [--max-concurrency N] " - "[--max-pending-requests N] [--pending-timeout-ms N] " + "[--max-pending-requests N] [--pending-timeout-ms N] [--boot-watchdog-timeout-s N] " "[--prefill-chunk N] [--log-stats-interval-ms N] [--device N] " "[--max-request-mib N] [--media-cache-mib N] [--media-live-mib N] " + "[--prefix-cache-mib N] [--prefix-cache-mib-min N] [--prefix-cache-mib-max N] " + "[--kv-capacity-min N] [--kv-capacity-max N] " + "[--vram-guarantee-context N] [--vram-guarantee-concurrency N] [--vram-floor-mib N] " + "[--vram-idle-release-after-s N] [--vram-observe-only] [--admin-vram] " "[--media-preprocess-threads N] " "[--request-log-jsonl FILE] " "[--response-store-max-records N] [--response-store-max-mib N] " "[--kv-dtype bf16|int8] [--spec mtp|dflash --draft-tokens N] " "[--default-max-tokens N] " "[--vision] [--no-cuda-graph] [--no-prefix-reuse] " - "[--lm-head-draft] [--no-thinking] [--preserve-thinking] [--cors] " + "[--lm-head-draft] [--no-thinking] [--preserve-thinking] [--tolerant-tool-calls] [--cors] " "[--temperature F] [--top-p F] [--top-k N] [--min-p F] [--presence-penalty F] " "[--frequency-penalty F] [--seed N] [--greedy]\n" " serves OpenAI Responses/Chat Completions and Anthropic Messages endpoints\n" " --default-max-tokens defaults to " + std::to_string(kDefaultMaxTokens) + " when omitted\n" + " --boot-watchdog-timeout-s defaults to 120; 0 disables the pre-listen boot watchdog\n" " --max-request-mib defaults to 384 and is enforced before JSON parsing\n" " --media-cache-mib defaults to 1024; 0 disables retained media reuse\n" + " --prefix-cache-mib reserves device memory for cross-request prefix seeds; 0 " + "(default) disables. --prefix-cache-mib-min/max set an elastic range; omitted " + "--prefix-cache-mib with a max boots at the max. min 0 is fully releasable\n" + " --kv-capacity-min/max set an elastic KV token range; omitted --kv-capacity " + "with a max boots at the max. --kv-capacity N still means min==max==N\n" + " --vram-idle-release-after-s N releases the seed store after N idle seconds " + "(0 disables, default). --vram-observe-only logs would-be releases without changing " + "allocations. --admin-vram exposes GET/POST /admin/vram and requires --api-key; " + "default off\n" " --media-live-mib defaults to 2048 and bounds all live BF16 patch payloads\n" " --media-preprocess-threads defaults to 0 (auto, at most 16 workers)\n" " --request-log-jsonl appends full-precision server/request records\n" @@ -93,9 +107,10 @@ std::string serve_usage_text(const char* argv0) { " --vision enables media and loads the fixed Vision GPU allocations\n" " --kv-capacity auto leaves " + std::to_string(kDefaultKvCapacityHeadroomBytes / (1024ULL * 1024ULL)) + - " MiB of sizing headroom\n" + " MiB of sizing headroom (bounded by max-context * max-concurrency)\n" " --no-prefix-reuse disables compatible-prefix caching (enabled by default)\n" " --preserve-thinking retains closed-turn assistant reasoning in later prompts\n" + " --tolerant-tool-calls recovers complete Qwen calls with malformed wrapper/suffix output\n" " sampler defaults come from the loaded model and resolved thinking mode; " "server flags and request fields override individual values.\n" " --greedy forces temperature 0 (exact argmax).\n"; @@ -116,6 +131,9 @@ ServeOptions parse_serve_options(int argc, char** argv) { } bool default_max_tokens_explicit = false; bool kv_capacity_explicit = false; + bool prefix_cache_min_explicit = false; + bool prefix_cache_max_explicit = false; + bool prefix_cache_bytes_explicit = false; if (argc >= 2 && (std::string(argv[1]) == "--help" || std::string(argv[1]) == "-h")) { options.help_requested = true; return options; @@ -154,6 +172,9 @@ ServeOptions parse_serve_options(int argc, char** argv) { } else if (arg == "--pending-timeout-ms") { options.pending_timeout_ms = static_cast( parse_nonnegative_int(require_value("--pending-timeout-ms"), "pending-timeout-ms")); + } else if (arg == "--boot-watchdog-timeout-s" || arg == "--boot-watchdog-s") { + options.boot_watchdog_timeout_s = static_cast( + parse_nonnegative_int(require_value(arg.c_str()), "boot-watchdog-timeout-s")); } else if (arg == "--prefill-chunk") { options.prefill_chunk = static_cast( parse_nonnegative_int(require_value("--prefill-chunk"), "prefill-chunk")); @@ -167,6 +188,68 @@ ServeOptions parse_serve_options(int argc, char** argv) { throw std::invalid_argument("--max-request-mib is out of range"); } options.max_request_bytes = static_cast(mib << 20); + } else if (arg == "--prefix-cache-mib") { + const std::uint64_t mib = + parse_u64(require_value("--prefix-cache-mib"), "prefix-cache-mib"); + if (mib > (1ULL << 20)) { + throw std::invalid_argument("--prefix-cache-mib is out of range"); + } + options.prefix_cache_bytes = static_cast(mib << 20); + prefix_cache_bytes_explicit = true; + } else if (arg == "--prefix-cache-mib-min") { + const std::uint64_t mib = + parse_u64(require_value("--prefix-cache-mib-min"), "prefix-cache-mib-min"); + if (mib > (1ULL << 20)) { + throw std::invalid_argument("--prefix-cache-mib-min is out of range"); + } + options.prefix_cache_min_bytes = static_cast(mib << 20); + prefix_cache_min_explicit = true; + } else if (arg == "--prefix-cache-mib-max") { + const std::uint64_t mib = + parse_u64(require_value("--prefix-cache-mib-max"), "prefix-cache-mib-max"); + if (mib > (1ULL << 20)) { + throw std::invalid_argument("--prefix-cache-mib-max is out of range"); + } + options.prefix_cache_max_bytes = static_cast(mib << 20); + prefix_cache_max_explicit = true; + } else if (arg == "--kv-capacity-min") { + const int value = + parse_nonnegative_int(require_value("--kv-capacity-min"), "kv-capacity-min"); + if (value == 0) { throw std::invalid_argument("--kv-capacity-min must be positive"); } + options.kv_capacity_min_tokens = static_cast(value); + } else if (arg == "--kv-capacity-max") { + const int value = + parse_nonnegative_int(require_value("--kv-capacity-max"), "kv-capacity-max"); + if (value == 0) { throw std::invalid_argument("--kv-capacity-max must be positive"); } + options.kv_capacity_max_tokens = static_cast(value); + } else if (arg == "--vram-guarantee-context") { + const int value = parse_nonnegative_int(require_value("--vram-guarantee-context"), + "vram-guarantee-context"); + if (value == 0) { + throw std::invalid_argument("--vram-guarantee-context must be positive"); + } + options.vram_guarantee_context = static_cast(value); + } else if (arg == "--vram-guarantee-concurrency") { + const int value = parse_nonnegative_int(require_value("--vram-guarantee-concurrency"), + "vram-guarantee-concurrency"); + if (value == 0) { + throw std::invalid_argument("--vram-guarantee-concurrency must be positive"); + } + options.vram_guarantee_concurrency = static_cast(value); + } else if (arg == "--vram-floor-mib") { + const std::uint64_t mib = + parse_u64(require_value("--vram-floor-mib"), "vram-floor-mib"); + if (mib > (1ULL << 20)) { + throw std::invalid_argument("--vram-floor-mib is out of range"); + } + options.vram_floor_bytes = static_cast(mib << 20); + } else if (arg == "--vram-idle-release-after-s") { + options.vram_idle_release_after_s = static_cast(parse_nonnegative_int( + require_value("--vram-idle-release-after-s"), "vram-idle-release-after-s")); + } else if (arg == "--vram-observe-only") { + options.vram_observe_only = true; + } else if (arg == "--admin-vram") { + options.enable_admin_vram = true; } else if (arg == "--media-cache-mib") { const std::uint64_t mib = parse_u64(require_value("--media-cache-mib"), "media-cache-mib"); @@ -233,6 +316,8 @@ ServeOptions parse_serve_options(int argc, char** argv) { options.enable_thinking = false; } else if (arg == "--preserve-thinking") { options.preserve_thinking = true; + } else if (arg == "--tolerant-tool-calls") { + options.tolerant_tool_calls = true; } else if (arg == "--cors") { options.enable_cors = true; } else if (arg == "--temperature") { @@ -261,9 +346,49 @@ ServeOptions parse_serve_options(int argc, char** argv) { throw std::invalid_argument("unknown argument: " + arg); } } + if (!prefix_cache_min_explicit && !prefix_cache_max_explicit) { + options.prefix_cache_min_bytes = options.prefix_cache_bytes; + options.prefix_cache_max_bytes = options.prefix_cache_bytes; + } else { + if (prefix_cache_max_explicit) { + options.prefix_cache_bytes = options.prefix_cache_max_bytes; + } else if (prefix_cache_bytes_explicit) { + options.prefix_cache_max_bytes = options.prefix_cache_bytes; + } + if (!prefix_cache_min_explicit) { + options.prefix_cache_min_bytes = + prefix_cache_bytes_explicit ? options.prefix_cache_bytes : 0; + } + } + if (options.prefix_cache_min_bytes > options.prefix_cache_max_bytes) { + throw std::invalid_argument("--prefix-cache-mib-min must not exceed --prefix-cache-mib-max"); + } + if (options.kv_capacity_max_tokens != 0 && !kv_capacity_explicit) { + options.kv_capacity = KvCapacityPolicy::explicit_capacity(options.kv_capacity_max_tokens); + kv_capacity_explicit = true; + } if (!kv_capacity_explicit) { options.kv_capacity = KvCapacityPolicy::explicit_capacity(options.max_context); } + if (options.kv_capacity.mode == KvCapacityMode::Explicit) { + if (options.kv_capacity_max_tokens == 0) { + options.kv_capacity_max_tokens = options.kv_capacity.explicit_tokens; + } + if (options.kv_capacity_min_tokens == 0) { + options.kv_capacity_min_tokens = options.kv_capacity.explicit_tokens; + } + if (options.kv_capacity_min_tokens > options.kv_capacity_max_tokens) { + throw std::invalid_argument("--kv-capacity-min must not exceed --kv-capacity-max"); + } + options.kv_capacity = + KvCapacityPolicy::explicit_capacity(options.kv_capacity_max_tokens); + } + if (options.vram_guarantee_context == 0) { + options.vram_guarantee_context = options.max_context; + } + if (options.enable_admin_vram && options.api_key.empty()) { + throw std::invalid_argument("--admin-vram requires --api-key"); + } if (options.port <= 0 || options.port > 65535) { throw std::invalid_argument("--port must be in [1,65535]"); } diff --git a/src/serve/serve_options.h b/src/serve/serve_options.h index b6db8e4fd5..c99297fed6 100644 --- a/src/serve/serve_options.h +++ b/src/serve/serve_options.h @@ -31,10 +31,22 @@ struct ServeOptions { std::uint32_t max_concurrency = 1; std::uint32_t max_pending_requests = 16; std::uint32_t pending_timeout_ms = 30000; + std::uint32_t boot_watchdog_timeout_s = 120; // 0 disables the pre-listen boot watchdog std::uint32_t prefill_chunk = 1024; std::uint32_t log_stats_interval_ms = 5000; // 0 disables periodic Engine throughput logs std::size_t max_request_bytes = kDefaultMaxRequestBytes; std::size_t media_cache_bytes = kDefaultMediaCacheBytes; + std::size_t prefix_cache_bytes = 0; + std::size_t prefix_cache_min_bytes = 0; + std::size_t prefix_cache_max_bytes = 0; + std::uint32_t kv_capacity_min_tokens = 0; + std::uint32_t kv_capacity_max_tokens = 0; + std::uint32_t vram_guarantee_context = 0; + std::uint32_t vram_guarantee_concurrency = 1; + std::size_t vram_floor_bytes = 0; + std::uint32_t vram_idle_release_after_s = 0; + bool vram_observe_only = false; + bool enable_admin_vram = false; std::size_t media_live_bytes = kDefaultMediaLiveBytes; std::uint32_t media_preprocess_threads = 0; std::size_t response_store_max_records = kDefaultResponseStoreRecords; @@ -47,8 +59,9 @@ struct ServeOptions { bool allow_prefix_reuse = true; bool enable_thinking = true; // default thinking mode for the generation prompt (--no-thinking opts out) - bool preserve_thinking = false; - int default_max_tokens = kDefaultMaxTokens; + bool preserve_thinking = false; + bool tolerant_tool_calls = false; // recover complete Qwen calls with malformed wrapper/suffix output + int default_max_tokens = kDefaultMaxTokens; bool enable_cors = false; // send permissive CORS headers for browser UIs // Process-level explicit overrides layered between registered model/mode defaults and request // fields. An omitted seed is replaced per request with a fresh random seed. diff --git a/src/serve/tool_call_parser.cpp b/src/serve/tool_call_parser.cpp index a7bd3ca205..74f1d3a8ec 100644 --- a/src/serve/tool_call_parser.cpp +++ b/src/serve/tool_call_parser.cpp @@ -4,11 +4,17 @@ #include #include +#include #include #include +#include #include #include +#include #include +#include +#include +#include namespace ninfer::serve { namespace { @@ -64,7 +70,120 @@ std::string new_tool_call_id() { return std::string(buf.data()); } -bool parse_parameter(std::string_view inner, std::size_t& pos, Json& args) { +const std::unordered_map>* tool_param_types( + const ToolParamTypeMap& map, const std::string& tool_name) { + const auto it = map.find(tool_name); + return it == map.end() ? nullptr : &it->second; +} + +// The full set of declared non-string types recorded for (tool_name, param), +// or nullptr when the parameter has no non-string schema permission. +const std::vector* param_declared_types(const ToolParamTypeMap& map, + const std::string& tool_name, + const std::string& param) { + const auto* params = tool_param_types(map, tool_name); + if (params == nullptr) { return nullptr; } + const auto it = params->find(param); + return it == params->end() ? nullptr : &it->second; +} + +// vLLM's qwen3coder coercion for boolean-declared parameters: the model may +// emit Python-style scalars (True/False, 1/0) that are not valid JSON. +// `value` is the lowercased raw text; "true"/"1" -> true, "false"/"0" -> +// false; anything else is not a boolean and stays raw text. +std::optional coerce_boolean(std::string_view value) { + if (value == "true" || value == "1") { return true; } + if (value == "false" || value == "0") { return false; } + return std::nullopt; +} + +} // namespace + +namespace { + +// Valid JSON Schema "type" values that are not string. A parameter is only +// allowed to deserialize when every type it declares is in this set. +const std::unordered_set& non_string_schema_types() { + static const std::unordered_set types = {"integer", "number", "boolean", + "array", "object", "null"}; + return types; +} + +// Classify a parameter's schema "type" (a string or an array of strings) into +// the set of declared types. Returns false if "type" is absent or not a +// string/array; in that case classification is uncertain and the caller +// preserves raw text. Returns true and fills `declared` otherwise. +bool classify_param_type(const Json& spec, std::vector& declared) { + const auto type_it = spec.find("type"); + if (type_it == spec.end()) { return false; } + if (type_it->is_string()) { + declared.push_back(type_it->get()); + return true; + } + if (type_it->is_array()) { + for (const Json& t : *type_it) { + if (!t.is_string()) { return false; } + declared.push_back(t.get()); + } + // An empty type array (e.g. "type":[]) is uncertain, not a positive + // declaration of a non-string type; returning false here preserves + // raw text and prevents all_non_string_types from succeeding vacuously + // on an empty set (which would record the parameter with no declared + // types). + if (declared.empty()) { return false; } + return true; + } + return false; +} + +// Whether every declared type is a valid non-string JSON Schema type. If any +// declared type is "string" or unknown/invalid, the schema permits (or may +// permit) a string value, so the parser must preserve raw text. +bool all_non_string_types(const std::vector& declared) { + const auto& valid = non_string_schema_types(); + for (const std::string& type : declared) { + if (valid.find(type) == valid.end()) { return false; } + } + return true; +} + +} // namespace + +ToolParamTypeMap build_tool_param_type_map(const std::vector& tools) { + ToolParamTypeMap map; + for (const ToolDefinition& tool : tools) { + // Replace any prior entry for this tool name first, before any early + // exit, so a redefinition with an empty/malformed/no-properties schema + // cannot leak stale non-string permissions from a previous definition. + map[tool.name] = {}; + if (tool.parameters_json.empty()) { continue; } + const Json schema = Json::parse(tool.parameters_json, nullptr, false); + if (!schema.is_object()) { continue; } + const auto props_it = schema.find("properties"); + if (props_it == schema.end() || !props_it->is_object()) { continue; } + // The entry for tool.name was already reset to empty at the top of + // the loop; populate it only from this definition's properties. + auto& inner = map[tool.name]; + for (const auto& [name, spec] : props_it->items()) { + if (!spec.is_object()) { continue; } + std::vector declared; + if (!classify_param_type(spec, declared)) { continue; } + // Record only when every declared type is a valid non-string type; + // string-allowed, unknown/invalid, and absent-type params are left + // out so the parser preserves raw text for them. Store the full + // declared set (not just the first element) so the parser can + // reason about nullable types (e.g. ["boolean","null"]) + // independently of the type-array order. + if (all_non_string_types(declared)) { inner[name] = declared; } + } + } + return map; +} + +namespace { + +bool parse_parameter(std::string_view inner, std::size_t& pos, Json& args, + const std::string& tool_name, const ToolParamTypeMap& param_types) { constexpr std::string_view kParamOpen = "* declared = param_declared_types(param_types, tool_name, key); + const bool is_boolean = + declared != nullptr && + std::find(declared->begin(), declared->end(), "boolean") != declared->end(); + if (is_boolean) { + // Boolean-declared params: the model may emit Python-style scalars + // (True/False, 1/0) that are not valid JSON, so coerce the raw text + // instead of adopting a parsed value. The literal null is JSON null: + // a valid value for a nullable boolean, and the faithful reading of + // the token otherwise (matching the fallthrough for other nullable + // types). Any other value stays raw text for the client to validate. + // Both comparisons are case-insensitive, matching the model's + // Python-style emissions (True/TRUE, Null/NULL). + std::string lower; + lower.reserve(raw_value.size()); + for (const char c : raw_value) { + lower.push_back(static_cast(std::tolower(static_cast(c)))); + } + if (std::optional coerced = coerce_boolean(lower)) { + args[key] = *coerced; + } else if (lower == "null") { + args[key] = nullptr; + } else { + args[key] = Json(raw_value); + } + pos = value_end + kParamClose.size(); + return true; + } + // Only adopt the deserialized JSON type when the schema explicitly + // permits a non-string type. For string-typed, unknown, or absent + // params, keep the raw text so the model's value reaches the client + // with its type intact; the client owns the schema and validates. + Json parsed = Json::parse(raw_value, nullptr, false); + const bool can_deserialize = declared != nullptr; + args[key] = (parsed.is_discarded() || !can_deserialize) ? Json(raw_value) : parsed; pos = value_end + kParamClose.size(); return true; } -bool parse_one_tool_call(std::string_view block, std::size_t max_name_length, ToolCall& out) { +Json typed_json_argument(const std::string& tool_name, const std::string& key, const Json& value, + const ToolParamTypeMap& param_types) { + const std::vector* declared = param_declared_types(param_types, tool_name, key); + const bool is_boolean = + declared != nullptr && + std::find(declared->begin(), declared->end(), "boolean") != declared->end(); + if (is_boolean) { + if (value.is_boolean() || value.is_null()) { return value; } + if (value.is_string()) { + const std::string raw = value.get(); + std::string lower; + lower.reserve(raw.size()); + for (const char c : raw) { + lower.push_back(static_cast(std::tolower(static_cast(c)))); + } + if (std::optional coerced = coerce_boolean(lower)) { return *coerced; } + if (lower == "null") { return nullptr; } + return value; + } + if (value.is_number_integer() && !value.is_number_float()) { + const auto n = value.get(); + if (n == 1) { return true; } + if (n == 0) { return false; } + } + return Json(value.dump()); + } + if (declared != nullptr) { return value; } + if (value.is_string()) { return value; } + return Json(value.dump()); +} + +std::size_t json_object_end(std::string_view text, std::size_t pos) { + if (pos >= text.size() || text[pos] != '{') { return std::string_view::npos; } + int depth = 0; + bool in_string = false; + bool escaped = false; + for (std::size_t i = pos; i < text.size(); ++i) { + const char c = text[i]; + if (in_string) { + if (escaped) { + escaped = false; + } else if (c == '\\') { + escaped = true; + } else if (c == '"') { + in_string = false; + } + continue; + } + if (c == '"') { + in_string = true; + } else if (c == '{') { + ++depth; + } else if (c == '}') { + --depth; + if (depth == 0) { return i + 1; } + } + } + return std::string_view::npos; +} + +bool parse_json_argument_object(std::string_view inner, std::size_t& pos, Json& args, + const std::string& tool_name, + const ToolParamTypeMap& param_types) { + const std::size_t end = json_object_end(inner, pos); + if (end == std::string_view::npos) { return false; } + Json payload = Json::parse(inner.substr(pos, end - pos), nullptr, false); + if (payload.is_discarded() || !payload.is_object()) { return false; } + for (auto it = payload.begin(); it != payload.end(); ++it) { + args[it.key()] = typed_json_argument(tool_name, it.key(), it.value(), param_types); + } + pos = end; + return true; +} + +bool parse_one_tool_call(std::string_view block, std::size_t max_name_length, + const ToolParamTypeMap& param_types, ToolCall& out, + std::size_t& consumed) { constexpr std::string_view kFunctionOpen = "', name_begin); if (name_end == std::string_view::npos || name_end == name_begin) { return false; } - const std::string name = std::string(block.substr(name_begin, name_end - name_begin)); + const std::string name = trim_ascii(block.substr(name_begin, name_end - name_begin)); if (!valid_function_name(name, max_name_length)) { return false; } pos = name_end + 1; - const std::size_t function_end = block.find(kFunctionClose, pos); - if (function_end == std::string_view::npos) { return false; } - const std::string_view params = block.substr(pos, function_end - pos); - Json args = Json::object(); - std::size_t param_pos = 0; - for (;;) { - skip_ws(params, param_pos); - if (param_pos >= params.size()) { break; } - if (!parse_parameter(params, param_pos, args)) { return false; } - } - - pos = function_end + kFunctionClose.size(); + Json args = Json::object(); skip_ws(block, pos); - if (pos != block.size()) { return false; } + if (pos < block.size() && block[pos] == '{') { + if (!parse_json_argument_object(block, pos, args, name, param_types)) { return false; } + skip_ws(block, pos); + if (!starts_with_at(block, pos, kFunctionClose)) { return false; } + pos += kFunctionClose.size(); + } else { + const std::size_t function_end = block.find(kFunctionClose, pos); + if (function_end == std::string_view::npos) { return false; } + const std::string_view params = block.substr(pos, function_end - pos); + std::size_t param_pos = 0; + for (;;) { + skip_ws(params, param_pos); + if (param_pos >= params.size()) { break; } + if (!parse_parameter(params, param_pos, args, name, param_types)) { return false; } + } + pos = function_end + kFunctionClose.size(); + } out.id = new_tool_call_id(); out.name = name; out.arguments_json = args.dump(); + consumed = pos; return true; } ParsedToolCallOutput fallback(const std::string& text) { ParsedToolCallOutput out; - out.content = text; + out.is_tool_call_response = false; + out.content = text; return out; } } // namespace ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, - std::size_t max_tool_name_length) { + std::size_t max_tool_name_length, bool tolerant) { + return parse_qwen_tool_call_output(text, max_tool_name_length, {}, tolerant); +} + +ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, + std::size_t max_tool_name_length, + const ToolParamTypeMap& param_types, + bool tolerant) { constexpr std::string_view kToolOpen = ""; constexpr std::string_view kToolClose = ""; @@ -133,23 +375,49 @@ ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, if (first == std::string::npos) { return fallback(text); } ParsedToolCallOutput out; + // Text before the first is ordinary content and is RETAINED here. + // Suppressing it in the parser makes the terminal body shorter than what a + // streaming response may already have emitted, which trips the + // streamed-vs-terminal invariant and kills the request mid-stream. The caller + // suppresses it instead, because only the caller knows how many bytes actually + // left the process. out.content = rtrim_ascii(std::string_view(text).substr(0, first)); std::size_t pos = first; while (pos < text.size()) { skip_ws(text, pos); if (pos >= text.size()) { break; } - if (!starts_with_at(text, pos, kToolOpen)) { return fallback(text); } + if (!starts_with_at(text, pos, kToolOpen)) { + const std::size_t next = text.find(kToolOpen, pos); + if (next == std::string::npos) { + if (!out.tool_calls.empty()) { break; } + return fallback(text); + } + pos = next; + } const std::size_t inner_begin = pos + kToolOpen.size(); - const std::size_t close = text.find(kToolClose, inner_begin); - if (close == std::string::npos) { return fallback(text); } ToolCall call; - if (!parse_one_tool_call(std::string_view(text).substr(inner_begin, close - inner_begin), - max_tool_name_length, call)) { + std::size_t consumed = 0; + if (!parse_one_tool_call(std::string_view(text).substr(inner_begin), max_tool_name_length, + param_types, call, consumed)) { + // Once one complete call has been recovered, do not discard it just + // because Qwen started a malformed second call, had conversational text, + // or was cut off by token budget. + if (!out.tool_calls.empty()) { break; } return fallback(text); } + pos = inner_begin + consumed; + skip_ws(text, pos); + if (starts_with_at(text, pos, kToolClose)) { + pos += kToolClose.size(); + } else if (!tolerant) { + if (!out.tool_calls.empty()) { break; } + return fallback(text); + } else { + out.tool_calls.push_back(std::move(call)); + break; + } out.tool_calls.push_back(std::move(call)); - pos = close + kToolClose.size(); } if (out.tool_calls.empty()) { return fallback(text); } @@ -174,12 +442,11 @@ std::string ToolCallStreamFilter::feed(std::string_view text) { std::isspace(static_cast(pending_[safe_end - 1])) != 0) { --safe_end; } - std::string visible = pending_.substr(0, safe_end); - tool_region_ = pending_.substr(safe_end); + held_prefix_ = pending_.substr(0, safe_end); + tool_region_ = pending_.substr(safe_end); pending_.clear(); saw_tool_marker_ = true; - emitted_bytes_ += visible.size(); - return visible; + return {}; } const std::size_t prefix = longest_suffix_prefix(pending_, kToolOpen); @@ -198,11 +465,14 @@ std::string ToolCallStreamFilter::finish(bool is_tool_call_response) { finished_ = true; if (is_tool_call_response) { pending_.clear(); + held_prefix_.clear(); tool_region_.clear(); return {}; } - std::string tail = std::move(pending_); + std::string tail = std::move(held_prefix_); + tail += pending_; tail += tool_region_; + pending_.clear(); tool_region_.clear(); emitted_bytes_ += tail.size(); return tail; diff --git a/src/serve/tool_call_parser.h b/src/serve/tool_call_parser.h index 1e98e2010f..0530efcf61 100644 --- a/src/serve/tool_call_parser.h +++ b/src/serve/tool_call_parser.h @@ -5,6 +5,7 @@ #include #include #include +#include #include namespace ninfer::serve { @@ -15,8 +16,46 @@ struct ParsedToolCallOutput { std::vector tool_calls; }; +// Per-tool parameter deserialization allow-list distilled from the request +// ToolDefinition list. Outer key: tool name. Inner key: parameter name; the +// inner value is the full set of declared non-string types, in schema order. +// Only parameters whose JSON Schema "type" is a valid non-string type (or an +// array of valid non-string types) are recorded. The parser deserializes +// recorded parameters and, for sets containing "boolean", coerces +// Python-style scalars (True/False, 1/0) to JSON booleans and the literal +// null to JSON null. A parameter absent from the inner map (and a tool +// absent from the outer map) has no schema permission to deserialize: the +// parser preserves raw text and the client owns type interpretation. +using ToolParamTypeMap = + std::unordered_map>>; + +// Distill the request ToolDefinition list into a per-tool parameter +// deserialization allow-list. Each ToolDefinition::parameters_json is a JSON +// Schema object; its "properties" object maps each parameter name to an +// object whose "type" field (a string or an array of strings) declares the +// schema type(s). A parameter is recorded only when every declared type is +// one of the valid non-string JSON Schema types {integer, number, boolean, +// array, object, null}; otherwise (string allowed, unknown/invalid type, or +// absent "type") it is omitted so the parser preserves raw text. A tool +// name seen again replaces its entry so a redefinition cannot leak stale +// non-string permissions from a prior definition. +ToolParamTypeMap build_tool_param_type_map(const std::vector& tools); + +// Parse Qwen's XML-like tool-call format. In tolerant mode, a complete function +// call is recovered even when the model adds wrapper garbage or suffix text. +// A successful parse RETAINS any preamble before as content. It is +// not user-visible assistant text for a tool turn (OpenAI sends content=null +// when tool_calls exist), but suppressing it here would make the terminal body +// shorter than what a streaming response may already have emitted, which aborts +// the request mid-stream. GenerationService drops it instead, once it knows +// streamed_content_bytes == 0. +ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, + std::size_t max_tool_name_length, + bool tolerant = false); ParsedToolCallOutput parse_qwen_tool_call_output(const std::string& text, - std::size_t max_tool_name_length); + std::size_t max_tool_name_length, + const ToolParamTypeMap& param_types, + bool tolerant = false); // Incrementally publishes text that is provably outside a possible Qwen // suffix. At terminal time, a valid tool response discards the @@ -30,6 +69,7 @@ class ToolCallStreamFilter { private: std::string pending_; + std::string held_prefix_; std::string tool_region_; std::size_t emitted_bytes_ = 0; bool saw_tool_marker_ = false; diff --git a/src/targets/qwen3_6/CMakeLists.txt b/src/targets/qwen3_6/CMakeLists.txt index a1979d6553..bccdb8513f 100644 --- a/src/targets/qwen3_6/CMakeLists.txt +++ b/src/targets/qwen3_6/CMakeLists.txt @@ -7,6 +7,7 @@ target_sources(ninfer_engine PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/impl/frontend/resources.cpp ${CMAKE_CURRENT_SOURCE_DIR}/impl/frontend/tokenizer.cpp ${CMAKE_CURRENT_SOURCE_DIR}/impl/runtime/prefix_identity.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/impl/runtime/prefix_seed_store.cpp ${CMAKE_CURRENT_SOURCE_DIR}/impl/state/decoder_state.cpp ${CMAKE_CURRENT_SOURCE_DIR}/impl/state/round_state.cpp ${CMAKE_CURRENT_SOURCE_DIR}/impl/runtime/visual_scatter.cpp diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/prepared_prompt.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/prepared_prompt.h index 0955c3ed34..b4f1523a82 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/prepared_prompt.h +++ b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/prepared_prompt.h @@ -73,6 +73,9 @@ struct RewriteCheckpointSpec { struct PromptIdentity { bool reusable = true; std::optional rewrite_checkpoint; + // Token frontier closing the prompt's leading stable span (the rendered system block). + // Message-boundary aligned; drives cross-request prefix-seed capture and matching. + std::optional prefix_seed_frontier; }; struct PrepareStats { diff --git a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/runtime.h b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/runtime.h index 3126f2a258..c6993928fe 100644 --- a/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/runtime.h +++ b/src/targets/qwen3_6/export/ninfer/targets/qwen3_6/runtime.h @@ -176,6 +176,9 @@ class Program { [[nodiscard]] MemorySummary memory_summary() const noexcept; void reset_memory_peaks() noexcept; + void release_prefix_seeds(); + bool reclaim_prefix_seeds(); + [[nodiscard]] std::size_t prefix_seed_held_bytes() const noexcept; private: explicit Program(std::unique_ptr> impl) noexcept; diff --git a/src/targets/qwen3_6/impl/frontend/chat_template.cpp b/src/targets/qwen3_6/impl/frontend/chat_template.cpp index d0316acbdc..e20052b022 100644 --- a/src/targets/qwen3_6/impl/frontend/chat_template.cpp +++ b/src/targets/qwen3_6/impl/frontend/chat_template.cpp @@ -349,6 +349,8 @@ RenderedChat CompiledChatTemplate::render(const std::vector& messag rendered += reasoning_instructions; rendered += "<|im_end|>\n"; } + const std::optional prefix_seed_offset = + rendered.empty() ? std::nullopt : std::optional(rendered.size()); const long last_query_index = last_real_user_query(messages); const bool preserve_thinking = options.preserve_thinking.value_or(effort_template); @@ -447,7 +449,9 @@ RenderedChat CompiledChatTemplate::render(const std::vector& messag .kind = RewriteCheckpointKind::ResponseReplay, .offset = rendered.size()}; } } - return RenderedChat{.text = std::move(rendered), .rewrite_checkpoint = rewrite_checkpoint}; + return RenderedChat{.text = std::move(rendered), + .rewrite_checkpoint = rewrite_checkpoint, + .prefix_seed_offset = prefix_seed_offset}; } } // namespace ninfer::targets::qwen3_6::frontend_internal diff --git a/src/targets/qwen3_6/impl/frontend/chat_template.h b/src/targets/qwen3_6/impl/frontend/chat_template.h index 292a0663d7..15bb654bf3 100644 --- a/src/targets/qwen3_6/impl/frontend/chat_template.h +++ b/src/targets/qwen3_6/impl/frontend/chat_template.h @@ -86,6 +86,8 @@ struct RewriteCheckpointByteSpec { struct RenderedChat { std::string text; std::optional rewrite_checkpoint; + // Byte offset just past the rendered leading system block, when one was emitted. + std::optional prefix_seed_offset; }; enum class ChatTemplateSemantics : std::uint8_t { diff --git a/src/targets/qwen3_6/impl/frontend/frontend.cpp b/src/targets/qwen3_6/impl/frontend/frontend.cpp index 478f63c2c6..86c758cdb1 100644 --- a/src/targets/qwen3_6/impl/frontend/frontend.cpp +++ b/src/targets/qwen3_6/impl/frontend/frontend.cpp @@ -924,8 +924,9 @@ PreparedPrompt Frontend::prepare(PromptInput input, const PreparationControl& co result.prepare.tokenize_seconds = std::chrono::duration(Clock::now() - tokenize_started).count(); fi::check_preparation_control(control, "tokenization"); - result.token_ids = std::move(encoded.input_ids); - result.identity.rewrite_checkpoint = encoded.rewrite_checkpoint; + result.token_ids = std::move(encoded.input_ids); + result.identity.rewrite_checkpoint = encoded.rewrite_checkpoint; + result.identity.prefix_seed_frontier = encoded.prefix_seed_frontier; assign_text_positions(result); } (void)checked_token_count(result.token_ids.size()); diff --git a/src/targets/qwen3_6/impl/frontend/processor.cpp b/src/targets/qwen3_6/impl/frontend/processor.cpp index c3d10d762d..9c3436a6ed 100644 --- a/src/targets/qwen3_6/impl/frontend/processor.cpp +++ b/src/targets/qwen3_6/impl/frontend/processor.cpp @@ -587,9 +587,34 @@ std::span ProcessedInput::position_axis(int axis) const { static_cast(axis) * input_ids.size(), input_ids.size()); } +namespace { + +std::optional exact_prefix_frontier(const Tokenizer& tokenizer, + const RenderedChat& rendered, + std::size_t offset, + const std::vector& input_ids) { + if (offset == 0 || offset >= rendered.text.size()) { return std::nullopt; } + const std::vector prefix = + tokenizer.encode(std::string_view(rendered.text).substr(0, offset)); + if (prefix.empty() || prefix.size() >= input_ids.size() || + !std::equal(prefix.begin(), prefix.end(), input_ids.begin())) { + return std::nullopt; + } + if (prefix.size() > std::numeric_limits::max()) { return std::nullopt; } + return static_cast(prefix.size()); +} + +} // namespace + EncodedChat encode_rendered_chat(const Tokenizer& tokenizer, const RenderedChat& rendered) { EncodedChat encoded; encoded.input_ids = tokenizer.encode(rendered.text); + if (rendered.prefix_seed_offset) { + // Best-effort: a system block whose byte boundary is not an exact token boundary simply + // yields no seed frontier; nothing downstream depends on one existing. + encoded.prefix_seed_frontier = exact_prefix_frontier( + tokenizer, rendered, *rendered.prefix_seed_offset, encoded.input_ids); + } if (!rendered.rewrite_checkpoint) { return encoded; } if (rendered.rewrite_checkpoint->offset > rendered.text.size()) { throw std::logic_error("rewrite checkpoint byte offset exceeds rendered chat"); diff --git a/src/targets/qwen3_6/impl/frontend/processor.h b/src/targets/qwen3_6/impl/frontend/processor.h index 28c1e74e0d..a55d1346a8 100644 --- a/src/targets/qwen3_6/impl/frontend/processor.h +++ b/src/targets/qwen3_6/impl/frontend/processor.h @@ -118,6 +118,7 @@ struct ProcessedInput { struct EncodedChat { std::vector input_ids; std::optional rewrite_checkpoint; + std::optional prefix_seed_frontier; }; EncodedChat encode_rendered_chat(const Tokenizer& tokenizer, const RenderedChat& rendered); diff --git a/src/targets/qwen3_6/impl/runtime/api_impl.h b/src/targets/qwen3_6/impl/runtime/api_impl.h index 0eadcde9bc..dae594a005 100644 --- a/src/targets/qwen3_6/impl/runtime/api_impl.h +++ b/src/targets/qwen3_6/impl/runtime/api_impl.h @@ -18,11 +18,15 @@ SequencePlan::SequencePlan( : impl_(std::move(impl)) {} template <> -SequencePlan::SequencePlan(SequencePlan&&) noexcept = default; +SequencePlan::SequencePlan(SequencePlan&& other) noexcept + : impl_(std::move(other.impl_)) {} template <> -SequencePlan& SequencePlan::operator=(SequencePlan&&) noexcept = default; +SequencePlan& SequencePlan::operator=(SequencePlan&& other) noexcept { + impl_ = std::move(other.impl_); + return *this; +} template <> -SequencePlan::~SequencePlan() = default; +SequencePlan::~SequencePlan() { impl_.reset(); } template <> std::uint32_t SequencePlan::capacity() const noexcept { @@ -60,11 +64,15 @@ SequencePlanner::SequencePlanner( : impl_(std::move(impl)) {} template <> -SequencePlanner::SequencePlanner(SequencePlanner&&) noexcept = default; +SequencePlanner::SequencePlanner(SequencePlanner&& other) noexcept + : impl_(std::move(other.impl_)) {} template <> -SequencePlanner& SequencePlanner::operator=(SequencePlanner&&) noexcept = default; +SequencePlanner& SequencePlanner::operator=(SequencePlanner&& other) noexcept { + impl_ = std::move(other.impl_); + return *this; +} template <> -SequencePlanner::~SequencePlanner() = default; +SequencePlanner::~SequencePlanner() { impl_.reset(); } template <> const runtime::SequenceCapacityCurve& SequencePlanner::capacity_curve() const noexcept { @@ -85,11 +93,15 @@ RequestBasePlan::RequestBasePlan( : impl_(std::move(impl)) {} template <> -RequestBasePlan::RequestBasePlan(RequestBasePlan&&) noexcept = default; +RequestBasePlan::RequestBasePlan(RequestBasePlan&& other) noexcept + : impl_(std::move(other.impl_)) {} template <> -RequestBasePlan& RequestBasePlan::operator=(RequestBasePlan&&) noexcept = default; +RequestBasePlan& RequestBasePlan::operator=(RequestBasePlan&& other) noexcept { + impl_ = std::move(other.impl_); + return *this; +} template <> -RequestBasePlan::~RequestBasePlan() = default; +RequestBasePlan::~RequestBasePlan() { impl_.reset(); } template <> const runtime::RequestPlanSummary& RequestBasePlan::summary() const noexcept { @@ -102,11 +114,15 @@ RequestPlan::RequestPlan(std::unique_ptr -RequestPlan::RequestPlan(RequestPlan&&) noexcept = default; +RequestPlan::RequestPlan(RequestPlan&& other) noexcept + : impl_(std::move(other.impl_)) {} template <> -RequestPlan& RequestPlan::operator=(RequestPlan&&) noexcept = default; +RequestPlan& RequestPlan::operator=(RequestPlan&& other) noexcept { + impl_ = std::move(other.impl_); + return *this; +} template <> -RequestPlan::~RequestPlan() = default; +RequestPlan::~RequestPlan() { impl_.reset(); } template <> const runtime::RequestPlanSummary& RequestPlan::summary() const noexcept { @@ -119,7 +135,7 @@ Program::Program(std::unique_ptr> impl) no : impl_(std::move(impl)) {} template <> -Program::~Program() noexcept = default; +Program::~Program() noexcept { impl_.reset(); } template <> RequestBasePlan @@ -221,6 +237,21 @@ void Program::reset_memory_peaks() noexcept { impl_->reset_memory_peaks(); } +template <> +void Program::release_prefix_seeds() { + impl_->release_prefix_seeds(); +} + +template <> +bool Program::reclaim_prefix_seeds() { + return impl_->reclaim_prefix_seeds(); +} + +template <> +std::size_t Program::prefix_seed_held_bytes() const noexcept { + return impl_->prefix_seed_held_bytes(); +} + template <> SequencePlanner make_sequence_planner(DeviceContext& device, const EngineOptions& options, diff --git a/src/targets/qwen3_6/impl/runtime/layouts.h b/src/targets/qwen3_6/impl/runtime/layouts.h index 737b700856..c116aabccf 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts.h +++ b/src/targets/qwen3_6/impl/runtime/layouts.h @@ -72,6 +72,7 @@ struct SequencePlanningInputs { StartupFeatures features; bool use_cuda_graph = true; int device = 0; + std::size_t prefix_cache_bytes = 0; }; } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS @@ -99,6 +100,7 @@ struct SequencePlanImpl { std::size_t request_transient_capacity_bytes = 0; std::size_t graph_allowance_bytes = 0; std::size_t device_reservation_bytes = 0; + std::size_t prefix_cache_bytes = 0; }; template <> diff --git a/src/targets/qwen3_6/impl/runtime/layouts_impl.h b/src/targets/qwen3_6/impl/runtime/layouts_impl.h index ca50405c34..4656f71aca 100644 --- a/src/targets/qwen3_6/impl/runtime/layouts_impl.h +++ b/src/targets/qwen3_6/impl/runtime/layouts_impl.h @@ -620,6 +620,7 @@ std::unique_ptr build_sequence_candidate(const SequencePlannin impl->proposal_head = inputs.proposal_head; impl->features = inputs.features; impl->use_cuda_graph = inputs.use_cuda_graph; + impl->prefix_cache_bytes = inputs.prefix_cache_bytes; impl->device = inputs.device; impl->kv_dtype = inputs.kv_dtype; impl->kv_quant_group = inputs.kv_quant_group; @@ -675,9 +676,12 @@ std::unique_ptr build_sequence_candidate(const SequencePlannin impl->device_reservation_bytes = checked_add( checked_add( - checked_add(impl->persistent.bytes, impl->workspace.capacity, "sequence memory plan"), - impl->request_transient_capacity_bytes, "request transient reservation"), - impl->graph_allowance_bytes, "sequence graph allowance"); + checked_add( + checked_add(impl->persistent.bytes, impl->workspace.capacity, + "sequence memory plan"), + impl->request_transient_capacity_bytes, "request transient reservation"), + impl->graph_allowance_bytes, "sequence graph allowance"), + impl->prefix_cache_bytes, "prefix-seed store reservation"); return impl; } @@ -701,6 +705,7 @@ make_sequence_planner_impl(DeviceContext& device, const EngineOptions& options, .features = qwen3_6::startup_features(options), .use_cuda_graph = options.use_cuda_graph, .device = options.device, + .prefix_cache_bytes = options.prefix_cache_bytes, }; const std::uint32_t logical_pages = page_count(inputs.capacity); const std::uint32_t minimum_pages = std::max(logical_pages, inputs.max_concurrency); diff --git a/src/targets/qwen3_6/impl/runtime/prefix_seed_store.cpp b/src/targets/qwen3_6/impl/runtime/prefix_seed_store.cpp new file mode 100644 index 0000000000..7a253a1927 --- /dev/null +++ b/src/targets/qwen3_6/impl/runtime/prefix_seed_store.cpp @@ -0,0 +1,302 @@ +#include "targets/qwen3_6/impl/runtime/prefix_seed_store.h" + +#include "core/device.h" + +#include +#include +#include + +namespace ninfer::targets::qwen3_6::detail { + +namespace { + +constexpr std::uint64_t kFnvOffset = 1469598103934665603ULL; +constexpr std::uint64_t kFnvPrime = 1099511628211ULL; + +void check_cuda(cudaError_t err, const char* what) { + if (err != cudaSuccess) { + throw std::runtime_error(std::string("PrefixSeedStore: ") + what + ": " + + cudaGetErrorString(err)); + } +} + +} // namespace + +std::uint64_t prefix_seed_hash(std::span tokens) { + std::uint64_t hash = kFnvOffset; + for (const TokenId token : tokens) { + std::uint64_t value = static_cast(token); + for (int i = 0; i < 4; ++i) { + hash ^= (value >> (8 * i)) & 0xFFULL; + hash *= kFnvPrime; + } + } + return hash; +} + +PrefixSeedStore::~PrefixSeedStore() noexcept { release(); } + +void PrefixSeedStore::release() noexcept { + entries_.clear(); + arena_used_ = 0; + arena_bytes_ = 0; + state_layers_ = 0; + conv_slot_bytes_ = 0; + recurrent_slot_bytes_ = 0; + hidden_bytes_ = 0; + if (arena_ != nullptr) { + (void)cudaFree(arena_); + arena_ = nullptr; + } +} + +void PrefixSeedStore::initialize(std::size_t budget_bytes, + const LinearAttentionStatePool& state_pool, + const PagedKVPool& text_pool, const PagedKVPool* backend_pool, + std::size_t hidden_bytes) { + (void)text_pool; + (void)backend_pool; + if (arena_ != nullptr) { throw std::logic_error("PrefixSeedStore is already initialized"); } + if (budget_bytes == 0) { return; } + state_layers_ = state_pool.layer_count(); + if (state_layers_ == 0) { throw std::invalid_argument("prefix seeds require GDN state"); } + conv_slot_bytes_ = state_pool.conv_slot(0, 0).bytes(); + recurrent_slot_bytes_ = state_pool.recurrent_slot(0, 0).bytes(); + if (hidden_bytes == 0) { + throw std::invalid_argument("prefix seeds require a hidden image size"); + } + hidden_bytes_ = hidden_bytes; + const std::size_t minimum = state_image_bytes() + hidden_bytes_ + (1ULL << 20); + if (budget_bytes < minimum) { + throw std::invalid_argument("prefix cache budget is below one seed entry"); + } + check_cuda(cudaMalloc(&arena_, budget_bytes), "arena allocation"); + arena_bytes_ = budget_bytes; + arena_used_ = 0; + std::fprintf(stderr, "ninfer: prefix-seed store enabled (%zu MiB)\n", budget_bytes >> 20); +} + +std::size_t PrefixSeedStore::state_image_bytes() const noexcept { + return static_cast(state_layers_) * (conv_slot_bytes_ + recurrent_slot_bytes_); +} + +std::size_t PrefixSeedStore::kv_page_bytes(const PagedKVPool& pool) const { + std::size_t bytes = 0; + for (std::size_t plane = 0; plane < pool.plane_count(); ++plane) { + bytes += pool.plane(plane).bytes() / pool.page_group_count(); + } + return bytes; +} + +std::int64_t PrefixSeedStore::find(const PreparedPromptData& prompt) const { + std::int64_t best = -1; + std::uint32_t best_len = 0; + for (std::size_t i = 0; i < entries_.size(); ++i) { + const Entry& entry = entries_[i]; + if (entry.frontier <= best_len || entry.frontier >= prompt.token_ids.size()) { continue; } + if (!entry_matches(static_cast(i), prompt)) { continue; } + best = static_cast(i); + best_len = entry.frontier; + } + return best; +} + +bool PrefixSeedStore::contains(const PreparedPromptData& prompt, std::uint32_t frontier) const { + for (std::size_t i = 0; i < entries_.size(); ++i) { + if (entries_[i].frontier == frontier && + entry_matches(static_cast(i), prompt)) { + return true; + } + } + return false; +} + +std::uint32_t PrefixSeedStore::entry_frontier(std::int64_t entry) const { + return entries_.at(static_cast(entry)).frontier; +} + +std::int32_t PrefixSeedStore::entry_rope_delta(std::int64_t entry) const { + return entries_.at(static_cast(entry)).rope_delta; +} + +std::span PrefixSeedStore::tokens(std::int64_t entry) const { + const Entry& e = entries_.at(static_cast(entry)); + return std::span(e.ledger.data(), e.ledger.size()); +} + +bool PrefixSeedStore::entry_matches(std::int64_t index, const PreparedPromptData& prompt) const { + const Entry& entry = entries_.at(static_cast(index)); + if (entry.frontier > prompt.token_ids.size()) { return false; } + const std::span head(prompt.token_ids.data(), entry.frontier); + if (prefix_seed_hash(head) != entry.hash) { return false; } + if (!std::equal(entry.ledger.begin(), entry.ledger.end(), prompt.token_ids.begin())) { + return false; + } + return prefix_matches(prompt, entry.ledger, entry.identity, entry.frontier); +} + +void PrefixSeedStore::copy_state_image(const LinearAttentionStatePool& state_pool, + std::int32_t slot, std::byte* arena_base, + std::size_t offset, bool to_arena, + cudaStream_t stream) const { + std::size_t cursor = offset; + for (std::uint32_t layer = 0; layer < state_layers_; ++layer) { + const Tensor conv = state_pool.conv_slot(layer, slot); + const Tensor recurrent = state_pool.recurrent_slot(layer, slot); + std::byte* arena_conv = arena_base + cursor; + std::byte* arena_recurrent = arena_base + cursor + conv_slot_bytes_; + if (to_arena) { + check_cuda(cudaMemcpyAsync(arena_conv, conv.data, conv_slot_bytes_, + cudaMemcpyDeviceToDevice, stream), + "conv state export"); + check_cuda(cudaMemcpyAsync(arena_recurrent, recurrent.data, recurrent_slot_bytes_, + cudaMemcpyDeviceToDevice, stream), + "recurrent state export"); + } else { + check_cuda(cudaMemcpyAsync(conv.data, arena_conv, conv_slot_bytes_, + cudaMemcpyDeviceToDevice, stream), + "conv state import"); + check_cuda(cudaMemcpyAsync(recurrent.data, arena_recurrent, recurrent_slot_bytes_, + cudaMemcpyDeviceToDevice, stream), + "recurrent state import"); + } + cursor += conv_slot_bytes_ + recurrent_slot_bytes_; + } +} + +void PrefixSeedStore::copy_kv_pages(const PagedKVPool& pool, const PagedKVAllocation& allocation, + std::uint32_t pages, std::byte* arena_base, std::size_t offset, + bool to_arena, cudaStream_t stream) const { + const std::span ids = allocation.page_ids(); + if (ids.size() < pages) { + throw std::logic_error("prefix seed KV span exceeds the mapped allocation"); + } + std::size_t cursor = offset; + for (std::size_t plane_index = 0; plane_index < pool.plane_count(); ++plane_index) { + const Tensor& plane = pool.plane(plane_index); + const std::size_t page_bytes = plane.bytes() / pool.page_group_count(); + auto* plane_base = static_cast(plane.data); + std::uint32_t logical = 0; + while (logical < pages) { + // Coalesce physically-consecutive pages into one transfer. + std::uint32_t run = 1; + while (logical + run < pages && ids[logical + run] == ids[logical + run - 1] + 1) { + ++run; + } + std::byte* pool_ptr = + plane_base + static_cast(ids[logical]) * page_bytes; + std::byte* arena_ptr = arena_base + cursor; + const std::size_t bytes = static_cast(run) * page_bytes; + if (to_arena) { + check_cuda(cudaMemcpyAsync(arena_ptr, pool_ptr, bytes, cudaMemcpyDeviceToDevice, + stream), + "KV page export"); + } else { + check_cuda(cudaMemcpyAsync(pool_ptr, arena_ptr, bytes, cudaMemcpyDeviceToDevice, + stream), + "KV page import"); + } + cursor += bytes; + logical += run; + } + } +} + +void PrefixSeedStore::capture(const PreparedPromptData& prompt, std::uint32_t frontier, + std::int32_t rope_delta, const LinearAttentionStatePool& state_pool, + std::int32_t state_slot, const Tensor& hidden, + const PagedKVPool& text_pool, const PagedKVAllocation& text_kv, + const PagedKVPool* backend_pool, + const PagedKVAllocation* backend_kv, cudaStream_t stream) { + if (!enabled()) { return; } + if (frontier == 0 || frontier > prompt.token_ids.size()) { + throw std::invalid_argument("prefix seed frontier does not lie inside the prompt"); + } + if (contains(prompt, frontier)) { return; } + if (hidden.data == nullptr || hidden.bytes() < hidden_bytes_) { + throw std::invalid_argument("prefix seed capture requires the captured hidden state"); + } + + const std::uint32_t text_pages = + 1U + (frontier - 1U) / static_cast(kPagedKVPageSize); + const std::uint32_t backend_pages = + (backend_pool != nullptr && backend_kv != nullptr) ? text_pages : 0U; + + Entry entry; + entry.hash = + prefix_seed_hash(std::span(prompt.token_ids.data(), frontier)); + entry.frontier = frontier; + entry.rope_delta = rope_delta; + entry.ledger.assign(prompt.token_ids.begin(), + prompt.token_ids.begin() + static_cast(frontier)); + entry.identity.reserve(frontier); + entry.identity.assign(prompt); + entry.identity.truncate(frontier); + + const std::size_t text_bytes = static_cast(text_pages) * kv_page_bytes(text_pool); + const std::size_t backend_bytes = + backend_pages != 0 ? static_cast(backend_pages) * kv_page_bytes(*backend_pool) + : 0ULL; + const std::size_t total = + state_image_bytes() + hidden_bytes_ + text_bytes + backend_bytes; + if (total > arena_bytes_) { return; } // cannot ever fit; skip silently + if (arena_used_ + total > arena_bytes_) { + // Generation flush: the bump arena reclaims space only wholesale. Captures are cheap and + // repopulate on demand, so correctness never depends on retained entries. + entries_.clear(); + arena_used_ = 0; + } + + entry.arena_offset = arena_used_; + entry.arena_bytes = total; + entry.state_offset = entry.arena_offset; + entry.hidden_offset = entry.state_offset + state_image_bytes(); + entry.text_kv_offset = entry.hidden_offset + hidden_bytes_; + entry.text_pages = text_pages; + entry.backend_kv_offset = entry.text_kv_offset + text_bytes; + entry.backend_pages = backend_pages; + + auto* base = static_cast(arena_); + copy_state_image(state_pool, state_slot, base, entry.state_offset, /*to_arena=*/true, stream); + check_cuda(cudaMemcpyAsync(base + entry.hidden_offset, hidden.data, hidden_bytes_, + cudaMemcpyDeviceToDevice, stream), + "hidden export"); + copy_kv_pages(text_pool, text_kv, text_pages, base, entry.text_kv_offset, /*to_arena=*/true, + stream); + if (backend_pages != 0) { + copy_kv_pages(*backend_pool, *backend_kv, backend_pages, base, entry.backend_kv_offset, + /*to_arena=*/true, stream); + } + + arena_used_ += total; + entries_.push_back(std::move(entry)); + std::fprintf(stderr, "ninfer: prefix seed captured frontier=%u bytes=%zu entries=%zu\n", + frontier, total, entries_.size()); +} + +void PrefixSeedStore::restore(std::int64_t index, const LinearAttentionStatePool& state_pool, + std::int32_t state_slot, Tensor& tail_hidden, + const PagedKVPool& text_pool, const PagedKVAllocation& text_kv, + const PagedKVPool* backend_pool, + const PagedKVAllocation* backend_kv, cudaStream_t stream) const { + const Entry& entry = entries_.at(static_cast(index)); + if (tail_hidden.data == nullptr || tail_hidden.bytes() < hidden_bytes_) { + throw std::invalid_argument("prefix seed restore requires the lane tail-hidden tensor"); + } + if (entry.backend_pages != 0 && (backend_pool == nullptr || backend_kv == nullptr)) { + throw std::logic_error("prefix seed entry carries backend KV the engine no longer has"); + } + auto* base = static_cast(arena_); + copy_state_image(state_pool, state_slot, base, entry.state_offset, /*to_arena=*/false, stream); + check_cuda(cudaMemcpyAsync(tail_hidden.data, base + entry.hidden_offset, hidden_bytes_, + cudaMemcpyDeviceToDevice, stream), + "hidden import"); + copy_kv_pages(text_pool, text_kv, entry.text_pages, base, entry.text_kv_offset, + /*to_arena=*/false, stream); + if (entry.backend_pages != 0) { + copy_kv_pages(*backend_pool, *backend_kv, entry.backend_pages, base, + entry.backend_kv_offset, /*to_arena=*/false, stream); + } +} + +} // namespace ninfer::targets::qwen3_6::detail diff --git a/src/targets/qwen3_6/impl/runtime/prefix_seed_store.h b/src/targets/qwen3_6/impl/runtime/prefix_seed_store.h new file mode 100644 index 0000000000..d403fc5a5d --- /dev/null +++ b/src/targets/qwen3_6/impl/runtime/prefix_seed_store.h @@ -0,0 +1,138 @@ +#pragma once + +// Content-addressed store of immutable prompt-prefix state snapshots ("seeds"). +// +// A seed captures the complete sequence state at a message-boundary frontier F produced by a +// real prefill: every Linear Attention layer's convolution and recurrent state, the Text (and, +// when MTP is active, backend) KV page payloads for tokens [0,F), the hidden state at F-1, and +// the host token ledger with its prefix identity. A later request whose prompt begins with the +// identical F tokens seeds a fresh lane by copying the entry in, then prefills only its suffix. +// +// Entries are immutable and restore-by-copy, so any number of concurrent requests can seed from +// the same entry; nothing is claimed or consumed. The store owns one fixed device arena sized at +// startup (GPU residency stays process-fixed) and holds no KV-pool pages, so admission +// accounting for the shared pools is unchanged. All device transfers are ordered on the caller's +// stream; the store is mutated only from the GPU executor lane. + +#include "core/linear_attention_state.h" +#include "core/paged_kv_cache.h" +#include "core/tensor.h" + +#include "targets/qwen3_6/impl/runtime/prefix_identity.h" + +#include + +#include + +#include +#include +#include +#include +#include + +namespace ninfer::targets::qwen3_6::detail { + +class PrefixSeedStore { +public: + PrefixSeedStore() = default; + ~PrefixSeedStore() noexcept; + + PrefixSeedStore(const PrefixSeedStore&) = delete; + PrefixSeedStore& operator=(const PrefixSeedStore&) = delete; + + /** + * Allocates the device arena. budget_bytes==0 leaves the store disabled. The layouts + * fix every entry's device image sizes except the per-entry KV span, which scales with the + * entry frontier. Safe to call again after release(); throws if an arena is already live. + */ + void initialize(std::size_t budget_bytes, const LinearAttentionStatePool& state_pool, + const PagedKVPool& text_pool, const PagedKVPool* backend_pool, + std::size_t hidden_bytes); + + // Frees the arena and drops every entry. enabled() becomes false. Must run on the GPU + // executor with the device stream idle. + void release() noexcept; + + [[nodiscard]] bool enabled() const noexcept { return arena_ != nullptr; } + [[nodiscard]] std::size_t arena_bytes() const noexcept { return arena_bytes_; } + [[nodiscard]] std::size_t entry_count() const noexcept { return entries_.size(); } + + /** Exact-token-prefix probe. Returns the entry index or -1. */ + [[nodiscard]] std::int64_t find(const PreparedPromptData& prompt) const; + + /** True when an entry already covers exactly this prompt's first `frontier` tokens. */ + [[nodiscard]] bool contains(const PreparedPromptData& prompt, std::uint32_t frontier) const; + + [[nodiscard]] std::uint32_t entry_frontier(std::int64_t entry) const; + [[nodiscard]] std::int32_t entry_rope_delta(std::int64_t entry) const; + [[nodiscard]] bool entry_matches(std::int64_t entry, const PreparedPromptData& prompt) const; + + /** + * Copies the state at `frontier` into a new entry. `state_slot` names the Linear Attention + * pool slot holding the captured image (the lane's rewrite-checkpoint slot immediately after + * the in-graph capture), `hidden` the captured hidden state at frontier-1, and the + * allocations the sequence's live KV whose leading pages cover [0,frontier). When the + * remaining arena cannot hold the new entry, the store flushes every resident entry and + * resets the bump offset, then captures if the empty arena can hold it. Silently skips + * capture when the entry cannot fit even in an empty arena. + */ + void capture(const PreparedPromptData& prompt, std::uint32_t frontier, std::int32_t rope_delta, + const LinearAttentionStatePool& state_pool, std::int32_t state_slot, + const Tensor& hidden, const PagedKVPool& text_pool, + const PagedKVAllocation& text_kv, const PagedKVPool* backend_pool, + const PagedKVAllocation* backend_kv, cudaStream_t stream); + + /** + * Copies entry state into a lane: Linear Attention image into `state_slot`, KV payloads into + * the leading pages of the destination allocations, and the entry hidden into `tail_hidden`. + * Host-side sequence fields (ledger, identity, frontiers) are the caller's responsibility, + * fed from tokens()/entry_rope_delta(). + */ + void restore(std::int64_t entry, const LinearAttentionStatePool& state_pool, + std::int32_t state_slot, Tensor& tail_hidden, const PagedKVPool& text_pool, + const PagedKVAllocation& text_kv, const PagedKVPool* backend_pool, + const PagedKVAllocation* backend_kv, cudaStream_t stream) const; + + [[nodiscard]] std::span tokens(std::int64_t entry) const; + +private: + struct Entry { + std::uint64_t hash = 0; + std::uint32_t frontier = 0; + std::int32_t rope_delta = 0; + std::vector ledger; + ResidentPrefixIdentity identity; + std::size_t arena_offset = 0; + std::size_t arena_bytes = 0; + std::size_t state_offset = 0; // conv+recurrent images, layer-major + std::size_t hidden_offset = 0; + std::size_t text_kv_offset = 0; + std::uint32_t text_pages = 0; + std::size_t backend_kv_offset = 0; + std::uint32_t backend_pages = 0; + }; + + [[nodiscard]] std::size_t state_image_bytes() const noexcept; + [[nodiscard]] std::size_t kv_page_bytes(const PagedKVPool& pool) const; + void copy_kv_pages(const PagedKVPool& pool, const PagedKVAllocation& allocation, + std::uint32_t pages, std::byte* arena_base, std::size_t offset, + bool to_arena, cudaStream_t stream) const; + void copy_state_image(const LinearAttentionStatePool& state_pool, std::int32_t slot, + std::byte* arena_base, std::size_t offset, bool to_arena, + cudaStream_t stream) const; + + void* arena_ = nullptr; + std::size_t arena_bytes_ = 0; + std::size_t arena_used_ = 0; // bump offset; a full arena is reclaimed by wholesale flush + std::deque entries_; + + // Fixed per-entry geometry captured at initialize(). + std::uint32_t state_layers_ = 0; + std::size_t conv_slot_bytes_ = 0; // one layer's conv image for one slot + std::size_t recurrent_slot_bytes_ = 0; // one layer's recurrent image for one slot + std::size_t hidden_bytes_ = 0; +}; + +[[nodiscard]] std::uint64_t prefix_seed_hash(std::span tokens); + +} // namespace ninfer::targets::qwen3_6::detail diff --git a/src/targets/qwen3_6/impl/runtime/program.h b/src/targets/qwen3_6/impl/runtime/program.h index c8a129c355..2c81f10110 100644 --- a/src/targets/qwen3_6/impl/runtime/program.h +++ b/src/targets/qwen3_6/impl/runtime/program.h @@ -12,6 +12,7 @@ #include "targets/qwen3_6/impl/runtime/dflash_context.h" #include "targets/qwen3_6/impl/runtime/linear_state_slots.h" #include "targets/qwen3_6/impl/runtime/prefix_identity.h" +#include "targets/qwen3_6/impl/runtime/prefix_seed_store.h" #include "targets/qwen3_6/impl/runtime/text_context.h" #include "targets/qwen3_6/impl/runtime/vision_context.h" #include "targets/qwen3_6/impl/runtime/vision_prefill.h" @@ -27,6 +28,10 @@ namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS { using PreparedPromptData = qwen3_6::PreparedPromptData; using RewriteCheckpointKind = qwen3_6::RewriteCheckpointKind; +// (prefix-seed store: see targets/qwen3_6/impl/runtime/prefix_seed_store.h) +// A seed entry costs a fixed Linear Attention state image regardless of its frontier, so tiny +// prompts (warmup probes, smoke requests) are not worth a capture. +inline constexpr std::uint32_t kMinimumSeedFrontierTokens = 256; using RewriteCheckpointSpec = qwen3_6::RewriteCheckpointSpec; using ReusePath = ninfer::PrefixReusePath; @@ -67,6 +72,7 @@ struct RequestBasePlanImpl { std::shared_ptr vision_control; std::size_t vision_transient_bytes = 0; std::optional rewrite_checkpoint; + std::optional prefix_seed_frontier; bool allow_prefix_reuse = false; }; @@ -82,6 +88,8 @@ struct RequestPlanImpl { NINFER_QWEN36_RUNTIME_NS::RewriteCheckpointAction rewrite_checkpoint_action = NINFER_QWEN36_RUNTIME_NS::RewriteCheckpointAction::Drop; std::optional rewrite_checkpoint_capture; + std::int64_t seed_entry = -1; + std::optional seed_capture; ops::SamplingConfig sampling; std::uint32_t text_kv_page_entitlement = 0; std::uint32_t backend_kv_page_entitlement = 0; @@ -194,6 +202,8 @@ struct RequestControl { bool prepare_mtp = false; ReusePath reuse = ReusePath::FullReset; MtpBridgeMode mtp_bridge = MtpBridgeMode::None; + std::optional seed_capture; + bool seed_captured = false; }; std::optional prefill; @@ -239,6 +249,10 @@ class ProgramImplCore { void reset_memory_peaks() noexcept; + void release_prefix_seeds(); + bool reclaim_prefix_seeds(); + [[nodiscard]] std::size_t prefix_seed_held_bytes() const noexcept; + const LoadedModelData& model; DeviceContext& device; const std::uint32_t capacity; @@ -254,6 +268,7 @@ class ProgramImplCore { const bool use_cuda_graph; const std::size_t kv_payload_bytes; const std::size_t graph_allowance_bytes; + const std::size_t prefix_cache_bytes; std::size_t graph_observed_bytes = 0; const WorkspacePlan workspace_plan; @@ -269,6 +284,7 @@ class ProgramImplCore { Tensor token_counts; Tensor tail_hidden_store; Tensor rewrite_checkpoint_hidden_store; + qwen3_6::detail::PrefixSeedStore prefix_seeds; std::array sequences; std::array requests; diff --git a/src/targets/qwen3_6/impl/runtime/program_impl.h b/src/targets/qwen3_6/impl/runtime/program_impl.h index b6b23f3ebf..d31ac6b265 100644 --- a/src/targets/qwen3_6/impl/runtime/program_impl.h +++ b/src/targets/qwen3_6/impl/runtime/program_impl.h @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -186,7 +187,8 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence kv_dtype(plan.kv_dtype), kv_quant_group(plan.kv_quant_group), proposal_head(plan.proposal_head), vision_enabled(plan.features.vision), use_cuda_graph(plan.use_cuda_graph), kv_payload_bytes(plan.persistent.kv_payload_bytes), - graph_allowance_bytes(plan.graph_allowance_bytes), workspace_plan(plan.workspace), + graph_allowance_bytes(plan.graph_allowance_bytes), prefix_cache_bytes(plan.prefix_cache_bytes), + workspace_plan(plan.workspace), persistent(plan.persistent.bytes), workspace_storage(plan.workspace.capacity), work(DeviceSpan{workspace_storage.base(), workspace_storage.capacity()}), round_host(sizeof(TokenId)), @@ -263,6 +265,11 @@ ProgramImplCore::ProgramImplCore(const LoadedModelData& model_in, const Sequence sequence.prefix_identity.reserve(static_cast(capacity) + 1ULL); } + prefix_seeds.initialize( + plan.prefix_cache_bytes, decoder->linear_attention, decoder->text_kv.pool(), + decoder->mtp_cache() != nullptr ? &decoder->mtp_cache()->pool() : nullptr, + sequences[0].rewrite_checkpoint_hidden.bytes()); + set_device_i32(io.text_kv_table_row, 0); set_device_i32(io.backend_kv_table_row, 0); @@ -418,24 +425,41 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan prompt.token_types.begin() + static_cast(request_plan.reuse_base), prompt.token_types.end(), [](std::uint8_t type) { return type != 0; }); if (suffix_has_visual != request_plan.vision.has_value()) { - throw std::invalid_argument("request plan does not describe the prompt suffix modality"); + throw std::logic_error("request plan does not describe the prompt suffix modality"); } if (request_plan.summary.transient_bytes != 0 && (transient.data == nullptr || transient.size < request_plan.summary.transient_bytes || transient.alignment < request_plan.summary.transient_alignment)) { - throw std::invalid_argument("request transient region does not satisfy the plan"); + throw std::logic_error("request transient region does not satisfy the plan"); } if (request_plan.reuse != ReusePath::FullReset && + request_plan.reuse != ReusePath::SeedPrefixCache && (!sequence.retained || !qwen3_6::detail::prefix_matches(prompt, sequence.ledger, sequence.prefix_identity, request_plan.reuse_base))) { - throw std::logic_error("planned resident prefix is no longer reusable"); + throw RequestError(RequestErrorKind::Unavailable, + "planned resident prefix is no longer reusable"); + } + if (request_plan.reuse == ReusePath::SeedPrefixCache) { + const bool seed_ok = + prefix_seeds.enabled() && request_plan.seed_entry >= 0 && + static_cast(request_plan.seed_entry) < prefix_seeds.entry_count() && + prefix_seeds.entry_matches(request_plan.seed_entry, prompt) && + prefix_seeds.entry_frontier(request_plan.seed_entry) == request_plan.reuse_base; + if (!seed_ok) { + request_plan.reuse = ReusePath::FullReset; + request_plan.reuse_base = 0; + request_plan.seed_entry = -1; + request_plan.mtp_bridge = MtpBridgeMode::None; + if (!prefix_seeds.enabled()) { request_plan.seed_capture.reset(); } + } } if (is_rewrite_checkpoint_restore(request_plan.reuse) && (!sequence.rewrite_checkpoint.valid || sequence.rewrite_checkpoint.frontier != request_plan.reuse_base || request_plan.reuse != restore_path(sequence.rewrite_checkpoint.kind))) { - throw std::logic_error("planned rewrite checkpoint is unavailable"); + throw RequestError(RequestErrorKind::Unavailable, + "planned rewrite checkpoint is unavailable"); } if (request_plan.rewrite_checkpoint_action == RewriteCheckpointAction::KeepExisting && (!prompt.identity.rewrite_checkpoint || !sequence.rewrite_checkpoint.valid || @@ -444,7 +468,8 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan request_plan.reuse == ReusePath::FullReset || !qwen3_6::detail::prefix_matches(prompt, sequence.ledger, sequence.prefix_identity, sequence.rewrite_checkpoint.frontier))) { - throw std::logic_error("planned rewrite checkpoint retention is unavailable"); + throw RequestError(RequestErrorKind::Unavailable, + "planned rewrite checkpoint retention is unavailable"); } if (request_plan.rewrite_checkpoint_action == RewriteCheckpointAction::ReclassifyExisting && (!prompt.identity.rewrite_checkpoint || !sequence.rewrite_checkpoint.valid || @@ -453,7 +478,8 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan request_plan.reuse == ReusePath::FullReset || !qwen3_6::detail::prefix_matches(prompt, sequence.ledger, sequence.prefix_identity, sequence.rewrite_checkpoint.frontier))) { - throw std::logic_error("planned rewrite checkpoint reclassification is unavailable"); + throw RequestError(RequestErrorKind::Unavailable, + "planned rewrite checkpoint reclassification is unavailable"); } if (request_plan.rewrite_checkpoint_action == RewriteCheckpointAction::CaptureNew && (!request_plan.rewrite_checkpoint_capture || !prompt.identity.rewrite_checkpoint || @@ -547,6 +573,35 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan device.stream); if (base == prompt_tokens) { copy_tail(sequence, sequence.rewrite_checkpoint_hidden); } sequence.ledger.resize(base); + } else if (request_plan.reuse == ReusePath::SeedPrefixCache) { + sequence.kv.reset(); + ordered_reset(sequence); + sequence.ledger.clear(); + sequence.text_kv_valid = 0; + sequence.mtp_kv_valid = 0; + reserve_sequence_kv(sequence, request_plan.text_kv_page_entitlement, + request_plan.backend_kv_page_entitlement); + sequence.kv->text.materialize_tokens(base, device.stream); + if (sequence.kv->backend) { + sequence.kv->backend->materialize_tokens(base, device.stream); + } + prefix_seeds.restore( + request_plan.seed_entry, decoder->linear_attention, + LinearStateSlots::current_state_slot(sequence.lane, max_concurrency), + sequence.tail_hidden, decoder->text_kv.pool(), sequence.kv->text, + decoder->mtp_cache() != nullptr ? &decoder->mtp_cache()->pool() : nullptr, + sequence.kv->backend ? &*sequence.kv->backend : nullptr, device.stream); + const std::span seed_tokens = + prefix_seeds.tokens(request_plan.seed_entry); + sequence.ledger.assign(seed_tokens.begin(), seed_tokens.end()); + sequence.prefix_identity.assign(prompt); + sequence.prefix_identity.truncate(base); + sequence.rope_delta = prefix_seeds.entry_rope_delta(request_plan.seed_entry); + sequence.tail_hidden_valid = true; + sequence.text_kv_valid = base; + if (speculative_backend == SpeculativeBackend::Mtp) { + sequence.mtp_kv_valid = base == 0 ? 0 : base - 1; + } } else { throw std::logic_error("request plan has an invalid prefix reuse path"); } @@ -594,7 +649,8 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan std::vector used(prompt.media_payloads.size(), false); for (const VisionUseSpan& use : request_plan.vision->uses) { if (use.item_index >= used.size()) { - throw std::logic_error("Vision plan references a missing media payload"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "Vision plan references a missing media payload"); } used[use.item_index] = true; } @@ -618,6 +674,8 @@ runtime::PrefillStepResult ProgramImplCore::start_prefill_lane(std::uint32_t lan .prepare_mtp = request_plan.prepare_mtp, .reuse = request_plan.reuse, .mtp_bridge = request_plan.mtp_bridge, + .seed_capture = request_plan.seed_capture, + .seed_captured = false, }; request.prefill.emplace(std::move(prefill)); auto& staged = *request.prefill; @@ -1573,8 +1631,24 @@ runtime::PrefillStepResult ProgramImplCore::advance_prefill(SequenceState& seque } if (staged.cursor < staged.prompt_tokens) { - const std::uint32_t nominal = - std::min(prefill_chunk, staged.prompt_tokens - staged.cursor); + std::uint32_t nominal = std::min(prefill_chunk, staged.prompt_tokens - staged.cursor); + const std::optional rewrite_frontier = + staged.rewrite_checkpoint_capture + ? std::optional(staged.rewrite_checkpoint_capture->frontier) + : std::nullopt; + // The lane checkpoint slot accepts one in-graph capture per chunk. When a pending + // seed capture and the request rewrite capture fall inside the same chunk at + // different frontiers, split the chunk at the seed frontier so each capture gets its + // own chunk (a large stable system head plus a short live turn makes this collision + // the common shape, not the exception). + if (staged.seed_capture && !staged.seed_captured && + *staged.seed_capture > staged.cursor && + *staged.seed_capture <= staged.cursor + nominal && rewrite_frontier && + *rewrite_frontier > staged.cursor && + *rewrite_frontier <= staged.cursor + nominal && + *rewrite_frontier != *staged.seed_capture) { + nominal = *staged.seed_capture - staged.cursor; + } const bool final_candidate = staged.cursor + nominal == staged.prompt_tokens; mark_workspace_usage(staged.prepare_mtp ? workspace_plan.mtp_prefill : workspace_plan.text_prefill); @@ -1582,10 +1656,16 @@ runtime::PrefillStepResult ProgramImplCore::advance_prefill(SequenceState& seque mark_workspace_usage(workspace_plan.dflash_context); } schedule::PrefillChunkResult result; + const bool rewrite_in_chunk = rewrite_frontier && + *rewrite_frontier > staged.cursor && + *rewrite_frontier <= staged.cursor + nominal; + const bool seed_in_chunk = + staged.seed_capture && !staged.seed_captured && + *staged.seed_capture > staged.cursor && + *staged.seed_capture <= staged.cursor + nominal && + (!rewrite_in_chunk || *rewrite_frontier == *staged.seed_capture); const std::optional rewrite_checkpoint_capture_frontier = - staged.rewrite_checkpoint_capture - ? std::optional(staged.rewrite_checkpoint_capture->frontier) - : std::nullopt; + seed_in_chunk ? staged.seed_capture : rewrite_frontier; if (staged.vision) { mark_workspace_usage(workspace_plan.vision_encode); result = schedule::prefill_multimodal_chunk( @@ -1607,6 +1687,21 @@ runtime::PrefillStepResult ProgramImplCore::advance_prefill(SequenceState& seque if (speculative_backend == SpeculativeBackend::DFlash) { sequence.dflash_context_frontier = staged.cursor; } + if (staged.seed_capture && !staged.seed_captured && + staged.cursor >= *staged.seed_capture) { + if (seed_in_chunk) { + prefix_seeds.capture( + staged.prompt, *staged.seed_capture, sequence.rope_delta, + decoder->linear_attention, + LinearStateSlots::rewrite_checkpoint_state_slot(sequence.lane, + max_concurrency), + sequence.rewrite_checkpoint_hidden, decoder->text_kv.pool(), + sequence.kv->text, + decoder->mtp_cache() != nullptr ? &decoder->mtp_cache()->pool() : nullptr, + sequence.kv->backend ? &*sequence.kv->backend : nullptr, device.stream); + } + staged.seed_captured = true; + } if (!result.finalized) { if (staged.cursor == staged.prompt_tokens) { @@ -2218,6 +2313,8 @@ MemorySummary ProgramImplCore::memory_summary() const noexcept { out.workspace_logical_peak_bytes = workspace_logical_peak_bytes; out.cuda_graph_allowance_bytes = graph_allowance_bytes; out.cuda_graph_observed_bytes = graph_observed_bytes; + out.prefix_cache_bytes = prefix_cache_bytes; + out.prefix_cache_held_bytes = prefix_seeds.arena_bytes(); out.kv_payload_bytes = kv_payload_bytes; return out; } @@ -2229,4 +2326,28 @@ void ProgramImplCore::reset_memory_peaks() noexcept { workspace_logical_peak_bytes = 0; } +void ProgramImplCore::release_prefix_seeds() { + device.synchronize(); + prefix_seeds.release(); +} + +bool ProgramImplCore::reclaim_prefix_seeds() { + if (prefix_seeds.enabled() || prefix_cache_bytes == 0) { return prefix_seeds.enabled(); } + device.synchronize(); + try { + prefix_seeds.initialize( + prefix_cache_bytes, decoder->linear_attention, decoder->text_kv.pool(), + decoder->mtp_cache() != nullptr ? &decoder->mtp_cache()->pool() : nullptr, + sequences[0].rewrite_checkpoint_hidden.bytes()); + return prefix_seeds.enabled(); + } catch (const std::exception& error) { + std::fprintf(stderr, "ninfer: prefix-seed reclaim failed: %s\n", error.what()); + return false; + } +} + +std::size_t ProgramImplCore::prefix_seed_held_bytes() const noexcept { + return prefix_seeds.arena_bytes(); +} + } // namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS diff --git a/src/targets/qwen3_6/impl/runtime/request_plan_impl.h b/src/targets/qwen3_6/impl/runtime/request_plan_impl.h index de125fd24b..bf78d15cd2 100644 --- a/src/targets/qwen3_6/impl/runtime/request_plan_impl.h +++ b/src/targets/qwen3_6/impl/runtime/request_plan_impl.h @@ -4,6 +4,7 @@ #include "targets/qwen3_6/impl/runtime/schedule.h" #include +#include #include #include #include @@ -15,13 +16,13 @@ void validate_sampling(const ResolvedSamplingParameters& sampling) { if (!std::isfinite(sampling.temperature) || !std::isfinite(sampling.top_p) || !std::isfinite(sampling.min_p) || !std::isfinite(sampling.presence_penalty) || !std::isfinite(sampling.frequency_penalty)) { - throw std::invalid_argument("sampling parameters must be finite"); + throw RequestError(RequestErrorKind::Unavailable, "sampling parameters must be finite"); } if (sampling.top_p < 0.0F || sampling.top_p > 1.0F) { - throw std::invalid_argument("top_p must be in [0,1]"); + throw RequestError(RequestErrorKind::Unavailable, "top_p must be in [0,1]"); } if (sampling.min_p < 0.0F || sampling.min_p > 1.0F) { - throw std::invalid_argument("min_p must be in [0,1]"); + throw RequestError(RequestErrorKind::Unavailable, "min_p must be in [0,1]"); } } @@ -60,35 +61,43 @@ std::uint64_t projected_service_work(const runtime::RequestPlanSummary& summary, RequestBasePlan ProgramImplCore::plan_request_base(const PreparedPromptData& prompt, const runtime::ResolvedExecutionOptions& options) { - if (prompt.token_ids.empty()) { throw std::invalid_argument("prompt must contain tokens"); } + if (prompt.token_ids.empty()) { + throw RequestError(RequestErrorKind::ContextLengthExceeded, "prompt must contain tokens"); + } if (prompt.token_ids.size() > capacity) { - throw std::invalid_argument("prompt exceeds configured context capacity"); + throw RequestError(RequestErrorKind::ContextLengthExceeded, + "prompt exceeds configured context capacity"); } if (prompt.token_ids.size() > std::numeric_limits::max()) { - throw std::overflow_error("prompt token count exceeds uint32"); + throw RequestError(RequestErrorKind::ContextLengthExceeded, + "prompt token count exceeds uint32"); } for (const TokenId id : prompt.token_ids) { if (id < 0 || id >= TextConfig::token_domain) { - throw std::invalid_argument("prompt contains token outside the 248077-token domain"); + throw RequestError(RequestErrorKind::ContextLengthExceeded, + "prompt contains token outside the 248077-token domain"); } } if (prompt.token_types.size() != prompt.token_ids.size() || prompt.positions.size() != 3ULL * prompt.token_ids.size()) { - throw std::invalid_argument("prepared prompt token metadata has an invalid shape"); + throw RequestError(RequestErrorKind::ContextLengthExceeded, + "prepared prompt token metadata has an invalid shape"); } if (prompt.has_media() != !prompt.media_payloads.empty() || prompt.media_payloads.size() != prompt.vision_items.size()) { - throw std::invalid_argument("prepared prompt media payload is incomplete"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "prepared prompt media payload is incomplete"); } for (std::size_t i = 0; i < prompt.media_payloads.size(); ++i) { if (!prompt.media_payloads[i] || prompt.media_payloads[i]->patch_elements != prompt.vision_items[i].patch_count * kPreparedVisionPatchFeatures) { - throw std::invalid_argument("prepared prompt media item payload has an invalid shape"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "prepared prompt media item payload has an invalid shape"); } } if (prompt.has_media() && !vision_enabled) { - throw std::invalid_argument("Vision is disabled for this Engine"); + throw RequestError(RequestErrorKind::Unavailable, "Vision is disabled for this Engine"); } validate_sampling(options.sampling); @@ -130,7 +139,8 @@ ProgramImplCore::plan_request_base(const PreparedPromptData& prompt, std::uint32_t previous_end = 0; for (const qwen3_6::VisionItemControl& item : control->items) { if (item.scatter_indices.empty()) { - throw std::invalid_argument("vision item has no Text consumer columns"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "vision item has no Text consumer columns"); } const auto first = static_cast(item.scatter_indices.front()); const auto last = static_cast(item.scatter_indices.back()); @@ -138,13 +148,16 @@ ProgramImplCore::plan_request_base(const PreparedPromptData& prompt, speculative_backend == SpeculativeBackend::Mtp && first != 0 ? first - 1 : first; const std::uint32_t end = last + 1; if (begin < previous_end) { - throw std::invalid_argument("vision item consumer spans overlap"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "vision item consumer spans overlap"); } if (end > base->summary.prompt_tokens) { - throw std::invalid_argument("vision item consumer span exceeds prompt"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "vision item consumer span exceeds prompt"); } if (schedule::VisionContext::workspace_bytes(item) > work.capacity()) { - throw std::invalid_argument("vision item exceeds the Program workspace envelope"); + throw RequestError(RequestErrorKind::MediaBudgetExceeded, + "vision item exceeds the Program workspace envelope"); } previous_end = end; max_merged = std::max(max_merged, item.merged_count); @@ -156,17 +169,26 @@ ProgramImplCore::plan_request_base(const PreparedPromptData& prompt, if (prompt.identity.rewrite_checkpoint) { const RewriteCheckpointSpec candidate = *prompt.identity.rewrite_checkpoint; if (candidate.frontier == 0 || candidate.frontier > base->summary.prompt_tokens) { - throw std::invalid_argument( + throw RequestError( + RequestErrorKind::ContextLengthExceeded, "rewrite checkpoint frontier must lie at or inside the prompt frontier"); } base->rewrite_checkpoint = candidate; } + if (prompt.identity.prefix_seed_frontier) { + const std::uint32_t frontier = *prompt.identity.prefix_seed_frontier; + if (frontier != 0 && frontier < base->summary.prompt_tokens) { + base->prefix_seed_frontier = frontier; + } + } const std::size_t cold_prefill_splits = (base->vision_control != nullptr ? base->vision_control->items.size() : 0ULL) + (base->rewrite_checkpoint && base->rewrite_checkpoint->frontier < base->summary.prompt_tokens ? 1ULL - : 0ULL); + : 0ULL) + + // A planned seed capture may split its chunk when the rewrite capture shares it. + (base->prefix_seed_frontier ? 1ULL : 0ULL); base->summary.service_work_quanta = projected_service_work(base->summary, 0, prefill_chunk, cold_prefill_splits); return RequestBasePlan(std::move(base)); @@ -210,6 +232,19 @@ RequestPlan ProgramImplCore::plan_request_for_lane(std::uint32_t lane, } } + // Cross-request prefix seeding: only when no resident sequence state is reusable, on a + // text-only prompt, outside the DFlash backend (whose context frontier a seed cannot feed). + if (plan->reuse == ReusePath::FullReset && prefix_seeds.enabled() && base.allow_prefix_reuse && + prompt.identity.reusable && !prompt.has_media() && + speculative_backend != SpeculativeBackend::DFlash) { + const std::int64_t entry = prefix_seeds.find(prompt); + if (entry >= 0) { + plan->reuse = ReusePath::SeedPrefixCache; + plan->reuse_base = prefix_seeds.entry_frontier(entry); + plan->seed_entry = entry; + } + } + if (speculative_backend == SpeculativeBackend::Mtp) { const bool append_ready = plan->reuse == ReusePath::AppendAtFrontier && sequence.tail_hidden_valid && @@ -218,9 +253,15 @@ RequestPlan ProgramImplCore::plan_request_for_lane(std::uint32_t lane, const bool checkpoint_ready = is_rewrite_checkpoint_restore(plan->reuse) && decoder->mtp_cache() != nullptr && plan->reuse_base != 0 && sequence.mtp_kv_valid >= plan->reuse_base - 1; - if (plan->reuse != ReusePath::FullReset && !append_ready && !checkpoint_ready) { + // A seed entry carries its own tail hidden and backend KV span, so its MTP readiness + // does not depend on resident sequence state. + const bool seed_ready = plan->reuse == ReusePath::SeedPrefixCache && + decoder->mtp_cache() != nullptr && plan->reuse_base != 0; + if (plan->reuse != ReusePath::FullReset && !append_ready && !checkpoint_ready && + !seed_ready) { plan->reuse = ReusePath::FullReset; plan->reuse_base = 0; + plan->seed_entry = -1; } } @@ -254,6 +295,20 @@ RequestPlan ProgramImplCore::plan_request_for_lane(std::uint32_t lane, plan->rewrite_checkpoint_action = RewriteCheckpointAction::DeferCapture; } + // Cross-request seed capture rides the in-graph checkpoint mechanism, so it is admissible + // only while the lane checkpoint slot holds no live state this request must preserve: + // CaptureNew rewrites the slot later at a farther frontier and Drop leaves it dead, while + // KeepExisting/ReclassifyExisting/DeferCapture all retain live checkpoint state. + if (prefix_seeds.enabled() && base.prefix_seed_frontier && !prompt.has_media() && + speculative_backend != SpeculativeBackend::DFlash && + (plan->rewrite_checkpoint_action == RewriteCheckpointAction::CaptureNew || + plan->rewrite_checkpoint_action == RewriteCheckpointAction::Drop) && + *base.prefix_seed_frontier >= kMinimumSeedFrontierTokens && + *base.prefix_seed_frontier > plan->reuse_base && + !prefix_seeds.contains(prompt, *base.prefix_seed_frontier)) { + plan->seed_capture = base.prefix_seed_frontier; + } + plan->summary.reusable_prompt_tokens = plan->reuse_base; if (speculative_backend == SpeculativeBackend::Mtp) { if (plan->reuse == ReusePath::FullReset) { @@ -268,6 +323,16 @@ RequestPlan ProgramImplCore::plan_request_for_lane(std::uint32_t lane, plan->mtp_bridge = plan->reuse_base < plan->summary.prompt_tokens ? MtpBridgeMode::BeforeSuffix : MtpBridgeMode::AfterExactHit; + } else if (plan->reuse == ReusePath::SeedPrefixCache) { + if (decoder->mtp_cache() == nullptr) { + plan->reuse = ReusePath::FullReset; + plan->reuse_base = 0; + plan->seed_entry = -1; + plan->prepare_mtp = true; + } else { + plan->prepare_mtp = true; + plan->mtp_bridge = MtpBridgeMode::BeforeSuffix; // frontier < prompt_tokens + } } } @@ -296,7 +361,9 @@ RequestPlan ProgramImplCore::plan_request_for_lane(std::uint32_t lane, (plan->rewrite_checkpoint_capture && plan->rewrite_checkpoint_capture->frontier < plan->summary.prompt_tokens ? 1ULL - : 0ULL); + : 0ULL) + + // The seed capture splits its chunk when the rewrite capture lands in the same one. + (plan->seed_capture ? 1ULL : 0ULL); plan->summary.service_work_quanta = projected_service_work(plan->summary, plan->reuse_base, prefill_chunk, prefill_splits); return RequestPlan(std::move(plan)); diff --git a/src/targets/qwen3_6/impl/runtime/text_context.h b/src/targets/qwen3_6/impl/runtime/text_context.h index 4e59e9bd82..68c2a55741 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context.h +++ b/src/targets/qwen3_6/impl/runtime/text_context.h @@ -11,6 +11,7 @@ #include "core/weight.h" #include "ninfer/ops/sampling.h" #include "ninfer/ops/gqa_attention.h" +#include "ninfer/ops/sparse_moe.h" #include #include #include @@ -240,7 +241,9 @@ class TextContext { [[nodiscard]] const MtpW& mtp_weights() const; void attn_mix(const FullLayerW& weights, Tensor& x, int index, Phase phase); void gdn_mix(const GdnLayerW& weights, Tensor& x, int index, Phase phase); - void mlp_tail(const Tensor* post_norm, const MlpW& weights, Tensor& x, Phase phase); + void mlp_tail(const Tensor* post_norm, const MlpW& weights, Tensor& x, Phase phase, + ops::WeightPrefetchSpan next_prefetch); + [[nodiscard]] ops::WeightPrefetchSpan next_projection_prefetch(int layer) const; void run_layers(Tensor& x, Phase phase); template void run_layers(Tensor& x, Phase phase, Tap& tap); diff --git a/src/targets/qwen3_6/impl/runtime/text_context_impl.h b/src/targets/qwen3_6/impl/runtime/text_context_impl.h index 5d7082996b..e7a42b6528 100644 --- a/src/targets/qwen3_6/impl/runtime/text_context_impl.h +++ b/src/targets/qwen3_6/impl/runtime/text_context_impl.h @@ -3,6 +3,7 @@ #include "targets/qwen3_6/impl/runtime/workspace_recipe.h" #include "core/nvtx.h" +#include "core/verbose.h" #include "targets/qwen3_6/impl/runtime/visual_scatter.h" #include "targets/qwen3_6/impl/runtime/vision_context.h" #include @@ -25,6 +26,7 @@ #include "ninfer/ops/residual_add.h" #include "ninfer/ops/rmsnorm.h" #include "ninfer/ops/rope.h" +#include "ninfer/ops/sparse_moe.h" #include "ninfer/ops/scatter.h" #include "ninfer/ops/scalar.h" #include "ninfer/ops/sigmoid_mul.h" @@ -34,6 +36,7 @@ #include #include +#include #include #include #include @@ -44,15 +47,66 @@ namespace ninfer::targets::qwen3_6::detail::NINFER_QWEN36_RUNTIME_NS::schedule { namespace { +// Verbose: dump the first few host int32 values about to be copied to device. +// Lets us tell whether a device buffer holds garbage because the *host* source +// was already garbage (upstream bug) or because the copy/device side is at fault. +void verbose_dump_i32(const char* label, const std::int32_t* source, std::size_t count) { + if (!verbose_enabled() || source == nullptr) { return; } + const std::size_t shown = std::min(count, 16); + std::fprintf(stderr, "[verbose] %s: n=%zu values=[", label, count); + for (std::size_t i = 0; i < shown; ++i) { + std::fprintf(stderr, "%s%d", i ? ", " : "", source[i]); + } + if (count > shown) { std::fprintf(stderr, ", ..."); } + std::fprintf(stderr, "]\n"); +} + void copy_i32(const std::int32_t* source, Tensor& destination, cudaStream_t stream) { if (source == nullptr || destination.dtype != DType::I32 || !destination.is_contiguous() || destination.data == nullptr) { throw std::invalid_argument("copy_i32: invalid host source or I32 destination"); } + verbose_dump_i32("copy_i32 h2d", source, + static_cast(destination.bytes() / sizeof(std::int32_t))); CUDA_CHECK(cudaMemcpyAsync(destination.data, source, destination.bytes(), cudaMemcpyHostToDevice, stream)); } +// Verbose: true if the stream is currently in CUDA graph capture mode. Synchronizing +// or doing a blocking device readback is illegal during capture, so probes must +// bail out when this is true. +bool verbose_stream_capturing(cudaStream_t stream) { + cudaStreamCaptureStatus capture = cudaStreamCaptureStatusNone; + return cudaStreamIsCapturing(stream, &capture) == cudaSuccess && + capture != cudaStreamCaptureStatusNone; +} + +// Verbose: read back a device int32 tensor (synchronously) and dump the first +// few values. Used to inspect the argmax index and the remap table at the point +// where the MTP draft token is produced. +void verbose_dump_device_i32(const char* label, const void* device_ptr, std::size_t count, + cudaStream_t stream) { + if (!verbose_enabled() || device_ptr == nullptr || count == 0 || + verbose_stream_capturing(stream)) { + return; + } + const std::size_t shown = std::min(count, 16); + std::vector host(shown); + const cudaError_t err = cudaMemcpy(host.data(), device_ptr, shown * sizeof(std::int32_t), + cudaMemcpyDeviceToHost); + if (err != cudaSuccess) { + std::fprintf(stderr, "[verbose] %s: readback FAILED: %s\n", label, + cudaGetErrorString(err)); + return; + } + std::fprintf(stderr, "[verbose] %s: n=%zu values=[", label, count); + for (std::size_t i = 0; i < shown; ++i) { + std::fprintf(stderr, "%s%d", i ? ", " : "", host[i]); + } + if (count > shown) { std::fprintf(stderr, ", ..."); } + std::fprintf(stderr, "]\n"); +} + void require_tensor_shape(const Tensor& t, DType dtype, std::initializer_list shape, const char* label) { if (t.dtype != dtype) { throw std::invalid_argument(std::string(label) + " dtype mismatch"); } @@ -379,10 +433,9 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po const auto results = workspace_recipe::mtp_attention_results(work_, T); Tensor qn = results.normalized_query.view({kCfg.head_dim, kCfg.n_q, T}); Tensor kn = results.normalized_key.view({kCfg.head_dim, kCfg.n_kv, T}); - ops::rmsnorm(q, *mtp_.q_norm, kCfg.rms_eps, true, qn, s); - ops::rmsnorm(k, *mtp_.k_norm, kCfg.rms_eps, true, kn, s); Tensor rope_for_op = active_sequence_batch_ != 0 ? rope_positions.view({T}) : rope_positions; - ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + ops::qk_norm_rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, q, *mtp_.q_norm, qn, k, + *mtp_.k_norm, kn, kCfg.rms_eps, s); Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T}); if (active_sequence_batch_ != 0) { @@ -399,11 +452,11 @@ void TextContext::mtp_forward_tail(Tensor& x, const Tensor& ah, const Tensor& po ops::gqa_attention(q_batch, k_batch, v_batch, position_batch, *active_valid_columns_, *active_backend_kv_table_rows_, kAttnScale, batch_mtp_kv_->batch_layer_view(0), envelope, work_, a_batch, s); + ops::sigmoid_mul(gate, a, s); } else { ops::gqa_attention(qn, kn, v, positions, Tensor{}, io_.backend_kv_table_row, kAttnScale, - batch_mtp_kv_->batch_layer_view(0), envelope, work_, a, s); + batch_mtp_kv_->batch_layer_view(0), envelope, work_, a, s, &gate); } - ops::sigmoid_mul(gate, a, s); const auto post = workspace_recipe::mtp_post_attention(work_, T); Tensor o = post.output; @@ -530,8 +583,7 @@ void TextContext::mtp_prefill_chunk(const Tensor& ids, const Tensor& hidden, Tensor a = work_.alloc(DType::BF16, {kCfg.head_dim, kCfg.n_q, 1}); ops::gqa_attention_cached(qn, last_position, kAttnScale, mtp_kv_.layer_view(0), envelope, - work_, a, s); - ops::sigmoid_mul(gate, a, s); + work_, a, s, &gate); Tensor o = work_.alloc(DType::BF16, {kCfg.hidden, 1}); ops::linear(a.view({kCfg.q_size, 1}), *mtp_.o_proj, o, s); @@ -557,8 +609,22 @@ void TextContext::proposal_argmax(const Tensor& hidden, Tensor& logits, Tensor& Tensor proposal_logits = work_.alloc(DType::BF16, {proposal_head_n_, T}); ops::linear(hidden, *proposal_head_, proposal_logits, ctx_.stream); ops::argmax(proposal_logits, proposal_tokens, proposal_head_n_, ctx_.stream); + if (verbose_enabled() && !verbose_stream_capturing(ctx_.stream)) { + CUDA_CHECK(cudaStreamSynchronize(ctx_.stream)); + verbose_dump_device_i32("proposal_argmax: argmax index (pre-remap)", + proposal_tokens.data, static_cast(T), + ctx_.stream); + verbose_dump_device_i32("proposal_argmax: remap table sample", proposal_head_ids_, + static_cast(proposal_head_n_), ctx_.stream); + } ops::proposal_remap_token_ids(proposal_tokens, proposal_head_ids_, proposal_head_n_, ctx_.stream); + if (verbose_enabled() && !verbose_stream_capturing(ctx_.stream)) { + CUDA_CHECK(cudaStreamSynchronize(ctx_.stream)); + verbose_dump_device_i32("proposal_argmax: remapped token ids (post-remap)", + proposal_tokens.data, static_cast(T), + ctx_.stream); + } } else { Tensor output_logits = matrix_window(logits, T); ops::linear(hidden, *lm_head_, output_logits, ctx_.stream); @@ -814,14 +880,13 @@ void TextContext::attn_mix(const FullLayerW& w, Tensor& x, int fidx, Phase ph) { const auto results = workspace_recipe::text_attention_results(work_, T); Tensor qn = results.normalized_query.view({kCfg.head_dim, kCfg.n_q, T}); Tensor kn = results.normalized_key.view({kCfg.head_dim, kCfg.n_kv, T}); - ops::rmsnorm(q, *w.q_norm, kCfg.rms_eps, true, qn, s); - ops::rmsnorm(k, *w.k_norm, kCfg.rms_eps, true, kn, s); const Tensor& cache_positions = active_cache_positions_ != nullptr ? *active_cache_positions_ : io_.pos; const Tensor& rope_positions = active_rope_positions_ != nullptr ? *active_rope_positions_ : io_.rope_pos; Tensor rope_for_op = active_sequence_batch_ != 0 ? rope_positions.view({T}) : rope_positions; - ops::rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, qn, kn, s); + ops::qk_norm_rope(rope_for_op, kCfg.rotary_dim, kCfg.rope_theta, q, *w.q_norm, qn, k, + *w.k_norm, kn, kCfg.rms_eps, s); Tensor a = results.attention.view({kCfg.head_dim, kCfg.n_q, T}); const Tensor& kv_table_rows = @@ -840,12 +905,12 @@ void TextContext::attn_mix(const FullLayerW& w, Tensor& x, int fidx, Phase ph) { ops::gqa_attention(q_batch, k_batch, v_batch, position_batch, valid, kv_table_rows, kAttnScale, batch_text_kv_->batch_layer_view(fidx), *active_gqa_envelope_, work_, a_batch, s); + ops::sigmoid_mul(gate, a, s); } else { ops::gqa_attention(qn, kn, v, cache_positions, Tensor{}, kv_table_rows, kAttnScale, batch_text_kv_->batch_layer_view(fidx), *active_gqa_envelope_, work_, a, - s); + s, &gate); } - ops::sigmoid_mul(gate, a, s); Variant::attention_output_projection(a.view({kCfg.q_size, T}), *w.o_proj, x, ph, work_, s); } @@ -955,13 +1020,27 @@ void TextContext::gdn_mix(const GdnLayerW& w, Tensor& x, int gidx, Phase ph) { Variant::gdn_output_projection(on.view({kCfg.value_dim, T}), *w.out_proj, x, ph, work_, s); } -void TextContext::mlp_tail(const Tensor* post_norm, const MlpW& m, Tensor& x, Phase ph) { +ops::WeightPrefetchSpan TextContext::next_projection_prefetch(int layer) const { + // Names the next layer's projection payload so the current post-mixer can warm L2 for it + // while its own tail runs. The last layer names nothing. + const int next = layer + 1; + if (next >= kCfg.n_layers) { return {}; } + if (ModelConfig::is_full(next)) { + return Variant::projection_prefetch_span( + *full_.at(static_cast(ModelConfig::full_idx(next))).projection); + } + return Variant::projection_prefetch_span( + *gdn_.at(static_cast(ModelConfig::gdn_idx(next))).projection); +} + +void TextContext::mlp_tail(const Tensor* post_norm, const MlpW& m, Tensor& x, Phase ph, + ops::WeightPrefetchSpan next_prefetch) { cudaStream_t s = ctx_.stream; const int T = x.ne[1]; Tensor h = workspace_recipe::post_mixer_hidden(work_, T); ops::rmsnorm(x, *post_norm, kCfg.rms_eps, true, h, s); - Variant::post_mixer(h, *m.payload, x, ph, work_, s); + Variant::post_mixer(h, *m.payload, x, ph, work_, s, next_prefetch); } template @@ -986,7 +1065,7 @@ void TextContext::run_layers(Tensor& x, Phase ph, Tap& tap) { prefill ? nvtx::Name::PrefillPostMixer : nvtx::Name::VerifyPostMixer, nvtx::Category::PostMixer, static_cast(layer)); auto mlp_scope = work_.scope(); - mlp_tail(full.post_attn_norm, full.mlp, x, ph); + mlp_tail(full.post_attn_norm, full.mlp, x, ph, next_projection_prefetch(layer)); if constexpr (Tap::enabled) { tap.capture_layer(layer, x, ctx_.stream); } } } else { @@ -1007,7 +1086,7 @@ void TextContext::run_layers(Tensor& x, Phase ph, Tap& tap) { prefill ? nvtx::Name::PrefillPostMixer : nvtx::Name::VerifyPostMixer, nvtx::Category::PostMixer, static_cast(layer)); auto mlp_scope = work_.scope(); - mlp_tail(gdn.post_attn_norm, gdn.mlp, x, ph); + mlp_tail(gdn.post_attn_norm, gdn.mlp, x, ph, next_projection_prefetch(layer)); if constexpr (Tap::enabled) { tap.capture_layer(layer, x, ctx_.stream); } } } diff --git a/src/targets/qwen3_6_27b/impl/variant.cpp b/src/targets/qwen3_6_27b/impl/variant.cpp index c2036d9145..ec6b8cf256 100644 --- a/src/targets/qwen3_6_27b/impl/variant.cpp +++ b/src/targets/qwen3_6_27b/impl/variant.cpp @@ -292,7 +292,8 @@ void Variant::gdn_norm_control_projection(const Tensor& residual, const Tensor& } void Variant::post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream) { + qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream, + ops::WeightPrefetchSpan) { auto scope = workspace.scope(); Tensor activation = workspace.alloc(DType::BF16, {TextConfig::intermediate, hidden.ne[1]}); ops::linear_swiglu(hidden, weights.gate_up, activation, text_policy(weights.gate_up), workspace, diff --git a/src/targets/qwen3_6_27b/impl/variant.h b/src/targets/qwen3_6_27b/impl/variant.h index 75332671ad..92faa50a1a 100644 --- a/src/targets/qwen3_6_27b/impl/variant.h +++ b/src/targets/qwen3_6_27b/impl/variant.h @@ -1,6 +1,7 @@ #pragma once #include "targets/qwen3_6_27b/impl/config.h" +#include "ninfer/ops/sparse_moe.h" #include "targets/qwen3_6_27b/impl/load/bindings.h" #include @@ -28,6 +29,16 @@ struct Variant { using VisionWeights = qwen3_6::VisionWeights; using GraphExecutionProfile = detail::GraphExecutionProfile; + // The dense post-mixer path does not consume weight-prefetch registrations yet. + static ::ninfer::ops::WeightPrefetchSpan + projection_prefetch_span(const FullAttentionProjectionWeights&) { + return {}; + } + static ::ninfer::ops::WeightPrefetchSpan + projection_prefetch_span(const GdnProjectionWeights&) { + return {}; + } + static constexpr float attention_scale = kAttentionScale; static constexpr float gdn_scale = kGdnScale; static constexpr std::uint32_t prefill_chunk_alignment = kPrefillChunkAlignment; @@ -79,7 +90,7 @@ struct Variant { WorkspaceArena& workspace, cudaStream_t stream); static void post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, qwen3_6::TextPhase phase, WorkspaceArena& workspace, - cudaStream_t stream); + cudaStream_t stream, ops::WeightPrefetchSpan next_prefetch = {}); static void mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream); [[nodiscard]] static std::size_t diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp index 96ab3ac8e6..2330487dfc 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.cpp +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.cpp @@ -64,13 +64,14 @@ bool dflash_target_uses_chunked_small_t(std::uint32_t draft_window, std::uint32_ } void run_sparse_moe(const Tensor& hidden, const ops::SparseMoeWeights& weights, Tensor& residual, - WorkspaceArena& workspace, cudaStream_t stream) { + WorkspaceArena& workspace, cudaStream_t stream, + ops::WeightPrefetchSpan next_prefetch) { auto scope = workspace.scope(); const DeviceSpan storage = workspace.alloc_bytes(ops::sparse_moe_workspace_capacity_bytes( weights.routed_gate_up.qtype, weights.routed_down.qtype, hidden.ne[1], hidden.ne[1])); WorkspaceArena leaf_workspace(storage); ops::sparse_moe(hidden, weights, ops::SparseMoeEpilogue::AddResidual, residual, leaf_workspace, - stream); + stream, next_prefetch); } void validate_token_interval(std::int32_t first, std::int32_t last) { @@ -213,13 +214,14 @@ void Variant::gdn_norm_control_projection(const Tensor& residual, const Tensor& } void Variant::post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, - qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream) { - run_sparse_moe(hidden, weights.op, residual, workspace, stream); + qwen3_6::TextPhase, WorkspaceArena& workspace, cudaStream_t stream, + ops::WeightPrefetchSpan next_prefetch) { + run_sparse_moe(hidden, weights.op, residual, workspace, stream, next_prefetch); } void Variant::mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream) { - run_sparse_moe(hidden, weights.op, residual, workspace, stream); + run_sparse_moe(hidden, weights.op, residual, workspace, stream, {}); } std::size_t Variant::mtp_attention_projection_workspace_capacity_bytes(std::int32_t first, diff --git a/src/targets/qwen3_6_35b_a3b/impl/variant.h b/src/targets/qwen3_6_35b_a3b/impl/variant.h index c6802e43f4..4c74cf1735 100644 --- a/src/targets/qwen3_6_35b_a3b/impl/variant.h +++ b/src/targets/qwen3_6_35b_a3b/impl/variant.h @@ -1,6 +1,7 @@ #pragma once #include "targets/qwen3_6_35b_a3b/impl/config.h" +#include "ninfer/ops/sparse_moe.h" #include "targets/qwen3_6_35b_a3b/impl/load/bindings.h" #include @@ -26,6 +27,17 @@ struct Variant { using VisionWeights = qwen3_6::VisionWeights; using GraphExecutionProfile = detail::GraphExecutionProfile; + static ::ninfer::ops::WeightPrefetchSpan + projection_prefetch_span(const FullAttentionProjectionWeights& weights) { + return {weights.query_key_gate_value.qdata, + static_cast(weights.query_key_gate_value.payload_bytes)}; + } + static ::ninfer::ops::WeightPrefetchSpan + projection_prefetch_span(const GdnProjectionWeights& weights) { + return {weights.query_key_value_z.qdata, + static_cast(weights.query_key_value_z.payload_bytes)}; + } + static constexpr float attention_scale = kAttentionScale; static constexpr float gdn_scale = kGdnScale; static constexpr std::uint32_t prefill_chunk_alignment = kPrefillChunkAlignment; @@ -85,7 +97,7 @@ struct Variant { WorkspaceArena& workspace, cudaStream_t stream); static void post_mixer(const Tensor& hidden, const PostMixerWeights& weights, Tensor& residual, qwen3_6::TextPhase phase, WorkspaceArena& workspace, - cudaStream_t stream); + cudaStream_t stream, ops::WeightPrefetchSpan next_prefetch = {}); static void mtp_post_mixer(const Tensor& hidden, const MtpPostMixerWeights& weights, Tensor& residual, WorkspaceArena& workspace, cudaStream_t stream); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 674da12a62..5f01ce51ef 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -58,9 +58,11 @@ ninfer_add_test(ninfer_artifact_reader_test ninfer_add_test(ninfer_artifact_materialization_test SOURCES test_artifact_materialization.cpp LIBRARIES ninfer_artifact) -ninfer_add_test(ninfer_media_decode_test - SOURCES test_media_decode.cpp - LIBRARIES ninfer_media_decode) +if(NINFER_BUILD_MEDIA) + ninfer_add_test(ninfer_media_decode_test + SOURCES test_media_decode.cpp + LIBRARIES ninfer_media_decode) +endif() ninfer_add_test(ninfer_device_test SOURCES test_device.cpp) ninfer_add_test(ninfer_decode_graph_test SOURCES test_decode_graph.cpp) ninfer_add_test(ninfer_tensor_test SOURCES test_tensor.cpp) @@ -163,8 +165,21 @@ ninfer_add_test(ninfer_request_log_test ninfer_add_test(ninfer_http_error_handler_test SOURCES test_http_error_handler.cpp LIBRARIES ninfer_serve) +add_executable(ninfer_supervisor_test test_ninfer_supervisor.cpp) +target_include_directories(ninfer_supervisor_test PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR} + ${PROJECT_SOURCE_DIR}/apps/ninfer-supervisor + ${PROJECT_SOURCE_DIR}/third_party) +add_test(NAME ninfer_supervisor_test COMMAND ninfer_supervisor_test) target_include_directories(ninfer_http_error_handler_test PRIVATE ${PROJECT_SOURCE_DIR}/third_party/cpp-httplib) +ninfer_add_test(ninfer_executor_recovery_test + SOURCES test_executor_recovery.cpp + NEEDS_SOURCE_DIR + LIBRARIES ninfer_serve ninfer_engine) +target_include_directories(ninfer_executor_recovery_test PRIVATE + ${PROJECT_SOURCE_DIR}/src/targets/qwen3_6/export + ${PROJECT_SOURCE_DIR}/third_party/cpp-httplib) ninfer_add_test(ninfer_bench_support_test SOURCES test_ninfer_bench_support.cpp ${PROJECT_SOURCE_DIR}/bench/targets/qwen3_6_27b/ninfer_bench_support.cpp diff --git a/tests/targets/qwen3_6_27b/test_load_plan.cpp b/tests/targets/qwen3_6_27b/test_load_plan.cpp index 1ccd8c9848..d54bfe8194 100644 --- a/tests/targets/qwen3_6_27b/test_load_plan.cpp +++ b/tests/targets/qwen3_6_27b/test_load_plan.cpp @@ -172,6 +172,46 @@ int verify_rejection() { return 1; } +int verify_prefix_cache_reservation() { + ninfer::DeviceContext device(0); + ninfer::EngineOptions options; + options.max_context = 128; + options.max_concurrency = 2; + options.kv_capacity = ninfer::KvCapacityPolicy::explicit_capacity(128); + options.prefill_chunk = 128; + options.use_cuda_graph = false; + options.prefix_cache_bytes = 0; + auto planner_zero = + Package::make_sequence_planner(device, options, WeightsProfile::Qwen36GroupwiseInt); + const auto curve_zero = planner_zero.capacity_curve(); + const std::uint32_t min_pages = curve_zero.minimum_main_page_groups; + auto plan_zero = std::move(planner_zero).finalize(min_pages); + + options.prefix_cache_bytes = 4ULL << 30; + auto planner_seed = + Package::make_sequence_planner(device, options, WeightsProfile::Qwen36GroupwiseInt); + const auto curve_seed = planner_seed.capacity_curve(); + auto plan_seed = std::move(planner_seed).finalize(min_pages); + + if (plan_zero.device_reservation_bytes() == 0 || + plan_seed.device_reservation_bytes() - plan_zero.device_reservation_bytes() != + (4ULL << 30)) { + std::cerr << "seed-store bytes were not added to the device reservation\n"; + return 1; + } + if (curve_seed.minimum_device_reservation_bytes - curve_zero.minimum_device_reservation_bytes != + (4ULL << 30)) { + std::cerr << "seed-store bytes were not added to the minimum reservation\n"; + return 1; + } + if (curve_seed.bytes_per_additional_main_page_group != + curve_zero.bytes_per_additional_main_page_group) { + std::cerr << "seed-store term changed the KV capacity stride\n"; + return 1; + } + return 0; +} + int verify_profile_mismatch_rejection() { ninfer::DeviceContext device(0); ninfer::EngineOptions options; @@ -207,6 +247,7 @@ int main() { return 77; } if (const int result = verify_rejection(); result != 0) { return result; } + if (const int result = verify_prefix_cache_reservation(); result != 0) { return result; } if (const int result = verify_profile_mismatch_rejection(); result != 0) { return result; } if (const int result = verify_groupwise(groupwise); result != 0) { return result; } if (const int result = verify_nvfp4(nvfp4); result != 0) { return result; } diff --git a/tests/test_anthropic_schema.cpp b/tests/test_anthropic_schema.cpp index 1a409775e7..d738a53051 100644 --- a/tests/test_anthropic_schema.cpp +++ b/tests/test_anthropic_schema.cpp @@ -417,6 +417,23 @@ int test_tools_and_choice() { return failures; } +int test_duplicate_tool_name_rejected() { + int failures = 0; + const Json tool = + Json{{"name", "get_weather"}, + {"input_schema", Json{{"type", "object"}, + {"properties", Json{{"city", Json{{"type", "string"}}}}}, + {"required", Json::array({"city"})}}}}; + Json body = { + {"model", "m"}, + {"max_tokens", 8}, + {"tools", Json::array({tool, tool})}, + {"messages", Json::array({Json{{"role", "user"}, {"content", "weather in Paris?"}}})}}; + failures += check(throws_api([&] { (void)parse_messages_request(body, default_limits()); }), + "duplicate tool names rejected"); + return failures; +} + int test_tool_use_result_roundtrip() { int failures = 0; const Json tool = Json{{"name", "get_weather"}, {"input_schema", Json{{"type", "object"}}}}; @@ -644,6 +661,31 @@ int test_response_serialization() { return failures; } +// Regression: a string-typed tool argument (e.g. taskId="1") must survive the +// Anthropic render path as a JSON string, not be coerced to a number. The +// parser preserves the raw text into arguments_json; make_messages_response +// forwards it verbatim into the tool_use.input object. +int test_string_typed_argument_survives_render() { + int failures = 0; + const CompletionUsage usage{3, 1}; + // arguments_json carries taskId as a string, exactly as the schema-aware + // parser emits it for a string-typed parameter with a numeric-looking value. + const std::vector calls = { + ToolCall{"toolu_1", "TaskUpdate", R"({"taskId":"1"})"}}; + const Json resp = Json::parse( + make_messages_response("msg_s", "claude-x", "", "", calls, "tool_use", usage)); + const Json& content = resp.at("content"); + failures += check(content.size() == 1 && content.at(0).at("type") == "tool_use", + "render produced a single tool_use block"); + const Json& input = content.at(0).at("input"); + failures += check(input.at("taskId").is_string(), + "string-typed taskId rendered as a JSON string, not a number"); + failures += check(input.at("taskId") == "1", "string-typed taskId value preserved on the wire"); + failures += check(!input.at("taskId").is_number(), + "string-typed taskId is not a JSON number on the wire"); + return failures; +} + int test_streaming_events() { int failures = 0; std::string type; @@ -750,11 +792,13 @@ int main() { failures += test_missing_and_bad_fields(); failures += test_parse_image(); failures += test_tools_and_choice(); + failures += test_duplicate_tool_name_rejected(); failures += test_tool_use_result_roundtrip(); failures += test_thinking_and_sampling(); failures += test_reasoning_effort(); failures += test_stop_reason_mapping(); failures += test_response_serialization(); + failures += test_string_typed_argument_survives_render(); failures += test_streaming_events(); failures += test_count_tokens_and_error(); if (failures == 0) { std::cout << "ok\n"; } diff --git a/tests/test_executor_recovery.cpp b/tests/test_executor_recovery.cpp new file mode 100644 index 0000000000..7b0e215249 --- /dev/null +++ b/tests/test_executor_recovery.cpp @@ -0,0 +1,555 @@ +#include "runtime/engine/concurrent_executor.h" +#include "serve/console_log.h" +#include "serve/generation_service.h" +#include "serve/http_server.h" +#include "targets/qwen3_6/impl/frontend/test_access.h" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#endif + +namespace { + +using Json = nlohmann::json; +using namespace ninfer; +using namespace ninfer::runtime; +using namespace ninfer::serve; +using FrontendTestAccess = ninfer::targets::qwen3_6::FrontendTestAccess; +using FrontendResources = ninfer::targets::qwen3_6::FrontendResources; + +std::string read_file(const char* path) { + std::ifstream stream(path, std::ios::binary); + if (!stream) { throw std::runtime_error(std::string("failed to open test resource: ") + path); } + return std::string(std::istreambuf_iterator(stream), std::istreambuf_iterator()); +} + +std::string read_template_fixture(const char* path) { + std::string source = read_file(path); + std::string normalized; + normalized.reserve(source.size()); + for (char c : source) { + if (c != '\r') { normalized.push_back(c); } + } + if (!normalized.empty() && normalized.back() == '\n') { normalized.pop_back(); } + return normalized; +} + +nlohmann::json added(int id, std::string content, bool special = false) { + return nlohmann::json{{"id", id}, + {"content", std::move(content)}, + {"single_word", false}, + {"lstrip", false}, + {"rstrip", false}, + {"normalized", false}, + {"special", special}}; +} + +nlohmann::json decoder_added(std::string content, bool special = false) { + nlohmann::json value = added(0, std::move(content), special); + value.erase("id"); + return value; +} + +FrontendResources minimal_resources() { + FrontendResources result; + result.chat_template_jinja = read_template_fixture( + NINFER_SOURCE_DIR "/tests/fixtures/frontend/thinking_toggle_chat_template.jinja"); + const nlohmann::json tokens = nlohmann::json::array( + {added(1, "helloST"), added(2, "OPtail"), added(3, "thought\n\nanswer"), added(6, "", true), added(7, "<0.0 seconds>"), + added(30, "user\n"), added(31, "assistant\n"), added(32, "\n"), + added(248045, "<|im_start|>", true), added(248046, "<|im_end|>", true), + added(248053, "<|vision_start|>", true), added(248054, "<|vision_end|>", true), + added(248056, "<|image_pad|>", true), added(248057, "<|video_pad|>", true), + added(248068, ""), added(248069, "")}); + result.tokenizer_json = nlohmann::json{ + {"model", + {{"type", "BPE"}, + {"vocab", {{"x", 0}, {"ä", 10}, {"¸", 11}, {"Ń", 12}}}, + {"merges", nlohmann::json::array()}}}, + {"added_tokens", + tokens}}.dump(); + + nlohmann::json decoder = nlohmann::json::object(); + for (const nlohmann::json& token : tokens) { + nlohmann::json value = token; + const std::string id = std::to_string(value.at("id").get()); + value.erase("id"); + decoder[id] = std::move(value); + } + decoder["248070"] = decoder_added("<|audio_start|>", true); + decoder["248071"] = decoder_added("<|audio_end|>", true); + decoder["248072"] = decoder_added("", true); + decoder["248073"] = decoder_added("", true); + decoder["248074"] = decoder_added("", true); + decoder["248075"] = decoder_added("", true); + decoder["248076"] = decoder_added("<|audio_pad|>", true); + result.tokenizer_config_json = nlohmann::json{ + {"add_bos_token", false}, + {"add_prefix_space", false}, + {"pad_token", "<|endoftext|>"}, + {"chat_template", result.chat_template_jinja}, + {"added_tokens_decoder", + std::move(decoder)}}.dump(); + result.generation_config_json = R"({"eos_token_id":[6]})"; + result.preprocessor_config_json = + R"({"patch_size":16,"temporal_patch_size":2,"merge_size":2,"image_mean":[0.5,0.5,0.5],"image_std":[0.5,0.5,0.5],"size":{"shortest_edge":4096,"longest_edge":16777216}})"; + result.video_preprocessor_config_json = + R"({"patch_size":16,"temporal_patch_size":2,"merge_size":2,"image_mean":[0.5,0.5,0.5],"image_std":[0.5,0.5,0.5],"size":{"shortest_edge":4096,"longest_edge":25165824}})"; + return result; +} + +// Minimal fake types to instantiate a real ConcurrentExecutor in-process with zero CUDA dependencies. +struct FakeProgram; +struct FakeBasePlan; +struct FakePlan; + +struct FakePackage { + using Program = FakeProgram; + using RequestBasePlan = FakeBasePlan; + using RequestPlan = FakePlan; +}; + +struct FakeBasePlan { + RequestPlanSummary summary_{}; + [[nodiscard]] const RequestPlanSummary& summary() const noexcept { return summary_; } +}; + +struct FakePlan { + RequestPlanSummary summary_{}; + [[nodiscard]] const RequestPlanSummary& summary() const noexcept { return summary_; } +}; + +struct FakeLoaded { + targets::qwen3_6::Frontend frontend = FrontendTestAccess::create_component(minimal_resources(), false); +}; + +struct FakeRequestMemory { + ArenaMemorySummary summary() const { return {}; } + void activate(std::size_t, std::size_t) {} + void deactivate() {} + TransientRegion region() const { return {}; } + void reset_peak() {} +}; + +struct FakeProgram { + AdmissionResources admission_capacity_{ + .active_lanes = 1, + .main_kv_pages = 100, + .backend_kv_pages = 0, + }; + bool throw_logic_error = false; + bool throw_runtime_error = false; + bool throw_request_error = false; + TokenId token = 6; // eos token () + + [[nodiscard]] AdmissionResources admission_capacity() const noexcept { return admission_capacity_; } + + FakeBasePlan plan_request_base(const targets::qwen3_6::PreparedPrompt&, + const ResolvedExecutionOptions&) { + FakeBasePlan p; + p.summary_.prompt_tokens = 2; + p.summary_.reusable_prompt_tokens = 0; + p.summary_.effective_output_tokens = 1; + p.summary_.effective_limit_reason = FinishReason::OutputLimit; + p.summary_.service_work_quanta = 1; + p.summary_.admission.active_lanes = 1; + p.summary_.admission.main_kv_pages = 1; + return p; + } + + FakePlan plan_request_for_lane(std::uint32_t, + const targets::qwen3_6::PreparedPrompt&, + const FakeBasePlan&) { + FakePlan p; + p.summary_.prompt_tokens = 2; + p.summary_.reusable_prompt_tokens = 0; + p.summary_.effective_output_tokens = 1; + p.summary_.effective_limit_reason = FinishReason::OutputLimit; + p.summary_.service_work_quanta = 1; + p.summary_.admission.active_lanes = 1; + p.summary_.admission.main_kv_pages = 1; + return p; + } + + bool can_admit_lane(std::uint32_t, const FakePlan&) const { return true; } + bool can_admit_lane_after_retained_eviction(std::uint32_t, const FakePlan&) const { return true; } + bool has_retained_lane(std::uint32_t) const { return false; } + void evict_retained_lane(std::uint32_t) {} + void abort_lane(std::uint32_t) {} + + PrefillStepResult start_prefill_lane(std::uint32_t, const targets::qwen3_6::PreparedPrompt&, + const FakePlan&, TransientRegion) { + if (throw_logic_error) { + throw std::logic_error("scheduler invariant test failure"); + } + if (throw_runtime_error) { + throw std::runtime_error("CUDA driver context wedged"); + } + if (throw_request_error) { + throw RequestError(RequestErrorKind::Unavailable, + "planned prefix seed is no longer available"); + } + return PrefillStepResult{ + .round = GeneratedRound{.tokens = std::span(&token, 1)}, + .complete = true, + }; + } + + PrefillStepResult advance_prefill_lane(std::uint32_t) { + return PrefillStepResult{ + .round = GeneratedRound{.tokens = std::span(&token, 1)}, + .complete = true, + }; + } + + void resolve_prefill_lane(std::uint32_t, bool) {} + BatchedGeneratedRound decode_batch(std::span, + std::span) { + return {}; + } + void resolve_pending_batch(std::span, std::span, + std::span, std::span) {} + MemorySummary memory_summary() const { return {}; } + void reset_memory_peaks() {} + std::size_t prefix_seed_held_bytes() const { return 0; } + std::size_t release_prefix_seeds() { return 0; } + bool reclaim_prefix_seeds() { return false; } + GenerationTimings generation_timings_lane(std::uint32_t) const { return {}; } + SpeculativeStats speculative_stats_lane(std::uint32_t) const { return {}; } +}; + +struct FakeInstance { + using Package = FakePackage; + std::unique_ptr program = std::make_unique(); + FakeRequestMemory request_memory; + std::shared_ptr loaded = std::make_shared(); + KvCapacityResolution kv_capacity_resolution{}; +}; + +void test_exception_types() { + std::cout << "Testing exception type classification...\n"; + try { + throw RequestError(RequestErrorKind::Unavailable, "test transient failure"); + } catch (const RequestError& err) { + if (err.kind() != RequestErrorKind::Unavailable) { + std::cerr << "FAIL: RequestError kind mismatch\n"; + std::exit(1); + } + } catch (...) { + std::cerr << "FAIL: RequestError was not caught by const RequestError&\n"; + std::exit(1); + } + + bool logic_error_caught = false; + try { + try { + throw std::logic_error("scheduler invariant violation"); + } catch (const RequestError&) { + std::cerr << "FAIL: std::logic_error incorrectly caught as RequestError!\n"; + std::exit(1); + } + } catch (const std::logic_error&) { + logic_error_caught = true; + } + if (!logic_error_caught) { + std::cerr << "FAIL: std::logic_error was lost\n"; + std::exit(1); + } +} + +void test_http_health_route() { + std::cout << "Testing HTTP /health route status behavior...\n"; + const ApiError unavail = request_error_to_api_error( + RequestError(RequestErrorKind::Unavailable, "inference engine is unavailable")); + if (unavail.status != 503 || unavail.code != "service_unavailable") { + std::cerr << "FAIL: Unavailable error does not map to 503 service_unavailable\n"; + std::exit(1); + } +} + +void test_real_path_fatal_classification() { + std::cout << "Testing real ConcurrentExecutor fatal error classification with probe...\n"; + FakeInstance instance; + instance.program->throw_logic_error = true; + + bool probe_fired = false; + std::string captured_message; + std::string captured_detail; + + EngineOptions options; + options.max_concurrency = 1; + options.max_pending_requests = 4; + options.pending_timeout_ms = 5000; + options.on_fatal_error = [&](std::exception_ptr err, const std::string& msg) { + probe_fired = true; + captured_message = msg; + if (err) { + try { + std::rethrow_exception(err); + } catch (const std::exception& e) { + captured_detail = e.what(); + } + } + }; + + ConcurrentExecutor executor(instance, options); + targets::qwen3_6::PreparedPrompt prompt = instance.loaded->frontend.prepare_tokens({1, 2}); + PromptSummary summary{.prompt_tokens = 2}; + ResolvedRequestOptions req_options{}; + + auto submission = executor.submit(std::move(prompt), summary, 0.0, std::move(req_options), + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + + bool submission_threw = false; + try { + submission.wait(nullptr, CancellationView()); + } catch (const std::logic_error&) { + submission_threw = true; + } catch (...) {} + + if (!submission_threw) { + std::cerr << "FAIL: submission.wait did not receive the fatal logic_error\n"; + std::exit(1); + } + if (!probe_fired) { + std::cerr << "FAIL: on_fatal_error probe was NOT invoked on real executor fatal error\n"; + std::exit(1); + } + if (captured_detail != "scheduler invariant test failure") { + std::cerr << "FAIL: captured exception detail mismatch: " << captured_detail << '\n'; + std::exit(1); + } + if (captured_message.find("fatal executor failure") == std::string::npos || + captured_message.find("scheduler invariant test failure") == std::string::npos || + captured_message.find("terminating process") == std::string::npos) { + std::cerr << "FAIL: captured formatted message malformed: " << captured_message << '\n'; + std::exit(1); + } + if (executor.is_healthy()) { + std::cerr << "FAIL: executor must report unhealthy (!is_healthy()) after fatal error\n"; + std::exit(1); + } +} + +void test_real_path_negative_request_error_does_not_exit() { + std::cout << "Testing real ConcurrentExecutor NEGATIVE test (RequestError must NOT reach on_fatal_error)...\n"; + FakeInstance instance; + instance.program->throw_request_error = true; + + bool probe_fired = false; + EngineOptions options; + options.max_concurrency = 1; + options.max_pending_requests = 4; + options.pending_timeout_ms = 5000; + options.on_fatal_error = [&](std::exception_ptr, const std::string&) { + probe_fired = true; + }; + + ConcurrentExecutor executor(instance, options); + targets::qwen3_6::PreparedPrompt prompt1 = instance.loaded->frontend.prepare_tokens({1, 2}); + PromptSummary summary1{.prompt_tokens = 2}; + ResolvedRequestOptions req_options1{}; + + // Request 1: Fails with RequestError + auto submission1 = executor.submit(std::move(prompt1), summary1, 0.0, std::move(req_options1), + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + + bool req1_threw = false; + try { + submission1.wait(nullptr, CancellationView()); + } catch (const RequestError& err) { + req1_threw = true; + if (err.kind() != RequestErrorKind::Unavailable) { + std::cerr << "FAIL: Request 1 threw wrong RequestErrorKind\n"; + std::exit(1); + } + } + + if (!req1_threw) { + std::cerr << "FAIL: Request 1 did not throw expected RequestError\n"; + std::exit(1); + } + + // THE CRITICAL GUARD: on_fatal_error MUST NOT have been called! + if (probe_fired) { + std::cerr << "FAIL: on_fatal_error was incorrectly called for a request-scoped RequestError!\n"; + std::exit(1); + } + if (!executor.is_healthy()) { + std::cerr << "FAIL: executor must remain healthy after a request-scoped failure\n"; + std::exit(1); + } + + // Request 2: Normal subsequent request must succeed on the recovered lane + instance.program->throw_request_error = false; + targets::qwen3_6::PreparedPrompt prompt2 = instance.loaded->frontend.prepare_tokens({1, 2}); + PromptSummary summary2{.prompt_tokens = 2}; + ResolvedRequestOptions req_options2{}; + + auto submission2 = executor.submit(std::move(prompt2), summary2, 0.0, std::move(req_options2), + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + + GenerationResult result = submission2.wait(nullptr, CancellationView()); + (void)result; + + if (probe_fired) { + std::cerr << "FAIL: on_fatal_error was triggered during recovered execution\n"; + std::exit(1); + } + if (!executor.is_healthy()) { + std::cerr << "FAIL: executor is not healthy after successful recovered request\n"; + std::exit(1); + } +} + +[[noreturn]] void run_child_real_fatal_exit() { + FakeInstance instance; + instance.program->throw_logic_error = true; + + EngineOptions options; + options.max_concurrency = 1; + options.max_pending_requests = 4; + options.pending_timeout_ms = 5000; + // No custom on_fatal_error: uses default ConcurrentExecutor hard exit to stderr + + ConcurrentExecutor executor(instance, options); + targets::qwen3_6::PreparedPrompt prompt = instance.loaded->frontend.prepare_tokens({1, 2}); + PromptSummary summary{.prompt_tokens = 2}; + ResolvedRequestOptions req_options{}; + + auto submission = executor.submit(std::move(prompt), summary, 0.0, std::move(req_options), + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + try { + submission.wait(nullptr, CancellationView()); + } catch (...) {} + + // Wait for the worker thread to catch the fatal error and call std::_Exit(1) + std::this_thread::sleep_for(std::chrono::seconds(2)); + std::exit(2); +} + +[[noreturn]] void run_child_real_serve_fatal_exit() { + FakeInstance instance; + instance.program->throw_runtime_error = true; + + EngineOptions options; + options.max_concurrency = 1; + options.max_pending_requests = 4; + options.pending_timeout_ms = 5000; + options.on_fatal_error = [](std::exception_ptr, const std::string& message) { + write_console_log(ConsoleLogLevel::Error, message); + std::cerr.flush(); + std::cout.flush(); + std::_Exit(1); + }; + + ConcurrentExecutor executor(instance, options); + targets::qwen3_6::PreparedPrompt prompt = instance.loaded->frontend.prepare_tokens({1, 2}); + PromptSummary summary{.prompt_tokens = 2}; + ResolvedRequestOptions req_options{}; + + auto submission = executor.submit(std::move(prompt), summary, 0.0, std::move(req_options), + std::chrono::steady_clock::now() + std::chrono::seconds(5)); + try { + submission.wait(nullptr, CancellationView()); + } catch (...) {} + + std::this_thread::sleep_for(std::chrono::seconds(2)); + std::exit(2); +} + +void test_subprocess_fatal_exit(const char* binary_path, const char* flag, + const std::string& expected_needle) { + std::cout << "Testing real subprocess exit on fatal executor failure (" << flag << ")...\n"; + std::string command = std::string(binary_path) + " " + flag + " 2>&1"; +#if defined(_WIN32) + FILE* pipe = _popen(command.c_str(), "r"); +#else + FILE* pipe = popen(command.c_str(), "r"); +#endif + if (!pipe) { + std::cerr << "FAIL: failed to open pipe for child test process\n"; + std::exit(1); + } + char buffer[256]; + std::string output; + while (std::fgets(buffer, sizeof(buffer), pipe) != nullptr) { + output += buffer; + } +#if defined(_WIN32) + const int raw_status = _pclose(pipe); + if (raw_status == -1) { + std::cerr << "FAIL: _pclose failed on child process\n"; + std::exit(1); + } + const int exit_code = raw_status; +#else + const int raw_status = pclose(pipe); + if (raw_status == -1) { + std::cerr << "FAIL: pclose failed on child process\n"; + std::exit(1); + } + const int exit_code = WIFEXITED(raw_status) ? WEXITSTATUS(raw_status) : -1; +#endif + + if (exit_code != 1) { + std::cerr << "FAIL: child process exited with code " << exit_code << ", expected 1\n"; + std::cerr << "Output was:\n" << output << '\n'; + std::exit(1); + } + if (output.find(expected_needle) == std::string::npos) { + std::cerr << "FAIL: child process output missing expected needle: " << expected_needle + << "\nActual output was:\n" + << output << '\n'; + std::exit(1); + } + if (output.find("fatal executor failure") == std::string::npos || + output.find("terminating process") == std::string::npos) { + std::cerr << "FAIL: child process output missing expected fatal marker\n" + << "Actual output was:\n" + << output << '\n'; + std::exit(1); + } +} + +} // namespace + +int main(int argc, char** argv) { + if (argc > 1) { + const std::string_view arg = argv[1]; + if (arg == "--child-fatal-exit") { + run_child_real_fatal_exit(); + } else if (arg == "--child-serve-fatal-exit") { + run_child_real_serve_fatal_exit(); + } + } + + test_exception_types(); + test_http_health_route(); + test_real_path_fatal_classification(); + test_real_path_negative_request_error_does_not_exit(); + test_subprocess_fatal_exit(argv[0], "--child-fatal-exit", + "scheduler invariant test failure"); + test_subprocess_fatal_exit(argv[0], "--child-serve-fatal-exit", + "CUDA driver context wedged"); + + std::cout << "All executor recovery, fatal error logging, and process exit tests passed.\n"; + return 0; +} diff --git a/tests/test_http_error_handler.cpp b/tests/test_http_error_handler.cpp index 88ae04725e..53e84942c4 100644 --- a/tests/test_http_error_handler.cpp +++ b/tests/test_http_error_handler.cpp @@ -42,6 +42,24 @@ int main() { failures += check(cancelled.status == 499 && cancelled.code == "client_disconnected", "preparation cancellation did not retain its HTTP classification"); + const ninfer::serve::ApiError unavailable = + ninfer::serve::request_error_to_api_error(ninfer::RequestError( + ninfer::RequestErrorKind::Unavailable, "inference engine is unavailable")); + failures += check(unavailable.status == 503 && unavailable.code == "service_unavailable", + "engine unavailability did not map to HTTP 503"); + + const ninfer::serve::ApiError overloaded = + ninfer::serve::request_error_to_api_error(ninfer::RequestError( + ninfer::RequestErrorKind::Overloaded, "inference request queue is full")); + failures += check(overloaded.status == 429 && overloaded.code == "server_overloaded", + "queue overflow did not map to HTTP 429"); + + const ninfer::serve::ApiError timeout = + ninfer::serve::request_error_to_api_error(ninfer::RequestError( + ninfer::RequestErrorKind::QueueTimeout, "inference request expired while waiting for admission")); + failures += check(timeout.status == 503 && timeout.code == "request_queue_timeout", + "queue timeout did not map to HTTP 503"); + httplib::Request messages_request; messages_request.path = "/v1/messages"; httplib::Response messages_response; diff --git a/tests/test_kv_capacity.cpp b/tests/test_kv_capacity.cpp index 866b98fc0f..065f3718d5 100644 --- a/tests/test_kv_capacity.cpp +++ b/tests/test_kv_capacity.cpp @@ -43,6 +43,34 @@ int main() { explicit_capacity.runtime_reservation_bytes == 1128, "explicit KV capacity did not use page-aligned token semantics"); + constexpr std::size_t kSeedStoreBytes = 4ULL << 30; + const ninfer::runtime::SequenceCapacityCurve with_seed{ + .main_page_tokens = 64, + .minimum_main_page_groups = 2, + .maximum_main_page_groups = 6, + .minimum_device_reservation_bytes = 1000 + kSeedStoreBytes, + .bytes_per_additional_main_page_group = 128, + }; + const auto explicit_seed = ninfer::runtime::resolve_kv_capacity( + ninfer::KvCapacityPolicy::explicit_capacity(129), with_seed, 1200 + kSeedStoreBytes); + failures += check(explicit_seed.main_page_groups == 3 && + explicit_seed.runtime_reservation_bytes == 1128 + kSeedStoreBytes && + explicit_seed.planned_slack_bytes == 72, + "explicit reservation omitted the constant seed-store term"); + + const auto automatic_seed = ninfer::runtime::resolve_kv_capacity( + ninfer::KvCapacityPolicy::automatic(50), with_seed, 1360 + kSeedStoreBytes); + failures += + check(automatic_seed.main_page_groups == 4 && automatic_seed.resolved_tokens == 256 && + automatic_seed.runtime_reservation_bytes == 1256 + kSeedStoreBytes && + automatic_seed.planned_slack_bytes == 104, + "automatic KV capacity did not keep the seed-store term as a constant addend"); + + const auto zero_seed = ninfer::runtime::resolve_kv_capacity( + ninfer::KvCapacityPolicy::explicit_capacity(129), curve, 1200); + failures += check(zero_seed.runtime_reservation_bytes == 1128, + "zero extra reservation term changed explicit accounting"); + bool insufficient_rejected = false; try { (void)ninfer::runtime::resolve_kv_capacity(ninfer::KvCapacityPolicy::automatic(50), curve, diff --git a/tests/test_ninfer_bench_support.cpp b/tests/test_ninfer_bench_support.cpp index 5a6500a879..6e76420937 100644 --- a/tests/test_ninfer_bench_support.cpp +++ b/tests/test_ninfer_bench_support.cpp @@ -242,6 +242,7 @@ qb::BenchEnvironment sample_environment() { env.memory.workspace = {100000000ULL, 0, 0}; env.memory.request_transient = {50000000ULL, 0, 40000000ULL}; env.memory.cuda_graph_allowance_bytes = 150000000ULL; + env.memory.prefix_cache_bytes = 4096ULL << 20; env.memory.kv_payload_bytes = 123456ULL; env.max_context = 4096; env.prefill_chunk = 1024; @@ -286,6 +287,8 @@ int test_report_contract() { "request transient capacity"); failures += expect(report.at("memory").at("cuda_graph_allowance_bytes") == 150000000ULL, "CUDA Graph allowance"); + failures += expect(report.at("memory").at("prefix_cache_bytes") == (4096ULL << 20), + "prefix cache reservation"); failures += expect(report.at("memory").at("kv_payload_bytes") == 123456ULL, "KV payload"); failures += expect(report.at("config").at("proposal_head") == "optimized", "proposal head"); failures += expect(report.at("config").at("decode_graph_prime").at("output_tokens") == 13, diff --git a/tests/test_ninfer_supervisor.cpp b/tests/test_ninfer_supervisor.cpp new file mode 100644 index 0000000000..e09226f322 --- /dev/null +++ b/tests/test_ninfer_supervisor.cpp @@ -0,0 +1,423 @@ +#include "logic.hpp" +#include "config.hpp" +#include "insights.hpp" + +#include +#include + +namespace { + +int fail(const std::string& m) { + std::cerr << "FAIL: " << m << '\n'; + return 1; +} +int check(bool c, const std::string& m) { return c ? 0 : fail(m); } + +int test_loopback() { + using namespace ninfer::supervisor; + int f = 0; + f += check(is_loopback_host("127.0.0.1") && is_loopback_host("localhost") && + is_loopback_host("::1"), + "loopback hosts"); + f += check(!is_loopback_host("0.0.0.0") && !is_loopback_host("192.168.1.2"), + "non-loopback hosts"); + f += check(is_loopback_peer("127.0.0.1") && is_loopback_peer("::ffff:127.0.0.1"), + "loopback peers"); + f += check(!is_loopback_peer("10.0.0.8") && !is_loopback_peer(""), "off-box peers"); + return f; +} + +int test_crash_loop() { + using clock = std::chrono::steady_clock; + ninfer::supervisor::RestartPolicy p; + p.crash_loop_max = 3; + p.crash_loop_window_s = 60; + ninfer::supervisor::RestartGate g(p); + const auto t0 = clock::now(); + int f = 0; + f += check(g.note_exit(t0) && g.note_exit(t0 + std::chrono::seconds(1)), "first exits allowed"); + f += check(!g.note_exit(t0 + std::chrono::seconds(2)) && g.halted(), + "third exit in window must halt"); + g.reset_halt(); + f += check(!g.halted() && g.note_exit(t0 + std::chrono::seconds(120)), + "reset allows restart"); + return f; +} + +int test_backoff() { + ninfer::supervisor::RestartGate g; + int f = 0; + f += check(g.backoff_seconds() == 1, "initial backoff 1s"); + g.advance_backoff(); + f += check(g.backoff_seconds() == 2, "backoff 2s"); + g.advance_backoff(); + g.advance_backoff(); + g.advance_backoff(); + g.advance_backoff(); + g.advance_backoff(); + f += check(g.backoff_seconds() == 60, "backoff caps at 60s"); + g.note_healthy(); + f += check(g.backoff_seconds() == 1, "healthy resets backoff"); + return f; +} + +int test_config_bind() { + int f = 0; + const char* ok = + R"({"engine":{"executable":"C:/ninfer-serve.exe"},"supervisor":{"host":"127.0.0.1"}})"; + try { + const auto c = ninfer::supervisor::load_config_json(ok); + f += check(c.host == "127.0.0.1" && !c.bind_any, "loopback config"); + } catch (...) { f += fail("loopback config threw"); } + bool rejected = false; + try { + (void)ninfer::supervisor::load_config_json( + R"({"engine":{"executable":"x"},"supervisor":{"host":"0.0.0.0"}})"); + } catch (const std::invalid_argument&) { rejected = true; } + f += check(rejected, "0.0.0.0 without bind_any must be rejected"); + bool any_ok = false; + try { + const auto c = ninfer::supervisor::load_config_json( + R"({"engine":{"executable":"x"},"supervisor":{"host":"0.0.0.0","bind_any":true}})"); + any_ok = c.bind_any; + } catch (...) {} + f += check(any_ok, "bind_any allows 0.0.0.0"); + return f; +} + +int test_host_header() { + using namespace ninfer::supervisor; + int f = 0; + f += check(host_header_allowed("127.0.0.1:8099", 8099, "127.0.0.1", false), + "loopback ipv4 host"); + f += check(host_header_allowed("localhost:8099", 8099, "127.0.0.1", false), + "localhost host"); + f += check(host_header_allowed("[::1]:8099", 8099, "127.0.0.1", false), "ipv6 loopback host"); + f += check(host_header_allowed("127.0.0.1", 8099, "127.0.0.1", false), + "loopback host without port"); + f += check(!host_header_allowed("attacker.example", 8099, "127.0.0.1", false), + "rebinding host rejected"); + f += check(!host_header_allowed("attacker.example:8099", 8099, "127.0.0.1", false), + "rebinding host:port rejected"); + f += check(!host_header_allowed("127.0.0.1:8080", 8099, "127.0.0.1", false), + "wrong port rejected"); + f += check(!host_header_allowed("", 8099, "127.0.0.1", false), "empty host rejected"); + f += check(!host_header_allowed("192.168.1.5:8099", 8099, "0.0.0.0", true), + "bind-any 0.0.0.0 does not open Host allowlist"); + f += check(host_header_allowed("192.168.1.5:8099", 8099, "192.168.1.5", true), + "bind-any named host is allowed"); + f += check(!host_header_allowed("192.168.1.5:8099", 8099, "192.168.1.5", false), + "named host without bind_any rejected"); + // Suffix/prefix traps. These pass a naive substring or starts_with check and + // are the classic way a rebinding defense gets reintroduced as a bug: an + // attacker controls the whole label, so "localhost.evil.com" is evil.com. + f += check(!host_header_allowed("localhost.evil.com:8099", 8099, "127.0.0.1", false), + "localhost-prefixed attacker domain rejected"); + f += check(!host_header_allowed("127.0.0.1.evil.com:8099", 8099, "127.0.0.1", false), + "ip-prefixed attacker domain rejected"); + f += check(!host_header_allowed("evil-localhost:8099", 8099, "127.0.0.1", false), + "localhost-suffixed attacker domain rejected"); + f += check(!host_header_allowed("192.168.1.5.evil.com:8099", 8099, "192.168.1.5", true), + "bind-any named host is matched exactly, not as a prefix"); + f += check(supervisor_control_header_ok("1") && !supervisor_control_header_ok("") && + !supervisor_control_header_ok("true"), + "control header is exactly 1"); + return f; +} + +int test_nvidia_csv() { + using namespace ninfer::supervisor; + int f = 0; + const auto a = parse_nvidia_smi_memory_csv("0, 24576, 32607\n1, 10, 20\n", 0); + f += check(a.ok && a.used_mib == 24576 && a.total_mib == 32607, "device 0 csv"); + const auto b = parse_nvidia_smi_memory_csv("0, 1, 2\n1, 99, 100\n", 1); + f += check(b.ok && b.used_mib == 99 && b.total_mib == 100, "device 1 csv"); + const auto c = parse_nvidia_smi_memory_csv("0, 1, 2\n", 3); + f += check(!c.ok && !c.error.empty(), "missing device"); + const auto d = parse_nvidia_smi_memory_csv("", 0); + f += check(!d.ok, "empty csv"); + f += check(mib_to_bytes(1) == 1048576, "mib_to_bytes"); + return f; +} + +int test_kv_line() { + using namespace ninfer::supervisor; + int f = 0; + const char* log = + "[info] ninfer-serve: model loaded in 1.2 s\n" + "[info] ninfer-serve: KV capacity auto resolved=8192 tokens pages=1/2 " + "runtime=1 prefix-cache=2 free-after-weights=3 free-after-startup=4 " + "headroom=5 slack=6 graphs=7/8\n" + "later line\n"; + const auto line = extract_kv_capacity_line(log); + f += check(line.find("KV capacity auto resolved=8192") != std::string::npos, + "extracts last KV capacity line"); + f += check(extract_kv_capacity_line("no capacity here").empty(), "missing line"); + return f; +} + +int test_monitor_only_config() { + int f = 0; + try { + const auto c = ninfer::supervisor::load_config_json( + R"({"engine":{"unmanaged":true,"engine_port":8010},"supervisor":{"host":"127.0.0.1"}})"); + f += check(!ninfer::supervisor::manages_engine_process(c) && c.engine.unmanaged, + "unmanaged does not require executable"); + } catch (...) { f += fail("unmanaged config threw"); } + try { + const auto c = ninfer::supervisor::load_config_json( + R"({"engine":{"engine_port":8010},"supervisor":{"host":"127.0.0.1"}})", true); + f += check(c.monitor_only && !ninfer::supervisor::manages_engine_process(c), + "CLI monitor_only does not require executable"); + } catch (...) { f += fail("monitor_only cli config threw"); } + bool rejected = false; + try { + (void)ninfer::supervisor::load_config_json( + R"({"engine":{},"supervisor":{"host":"127.0.0.1"}})"); + } catch (const std::invalid_argument&) { rejected = true; } + f += check(rejected, "managed config still requires executable"); + return f; +} + +int test_insights_honesty() { + using namespace ninfer::supervisor; + int f = 0; + const auto missing = insights_from_request_log_path(""); + f += check(missing.at("source").at("request_log") == "unconfigured", "unconfigured source"); + f += check(missing.at("insights").size() >= 1 && + missing.at("insights").at(0).at("availability") == "unavailable", + "missing log is unavailable, not a zero"); + f += check(missing.at("insights").at(0).at("statement").get().find( + "no request_done records") != std::string::npos, + "unavailable statement"); + + const auto typed = analyze_request_log_jsonl( + R"({"type":"request_done","timestamp_unix_ms":1})" + "\n", + "mem"); + f += check(typed.at("insights").at(0).at("availability") == "unavailable", + "type-key records do not count as request_done"); + + const char* jsonl = + R"({"event":"request_start","server_instance_id":"a","timestamp_unix_ms":1000,"request":{"request_id":1,"enable_thinking":true,"tool_count":0,"requested_output_tokens":8}})" + "\n" + R"({"event":"request_done","server_instance_id":"a","timestamp_unix_ms":1600,"request":{"request_id":1,"enable_thinking":true,"tool_count":0,"requested_output_tokens":8},"result":{"finish_reason":"output_limit","completion_tokens":8},"timings_seconds":{"prepare":0.01,"prefill":0.02,"decode":0.02,"ttft":0.03,"total":0.1}})" + "\n" + R"({"event":"request_start","server_instance_id":"b","timestamp_unix_ms":2000,"request":{"request_id":1,"enable_thinking":false,"tool_count":0,"requested_output_tokens":16}})" + "\n" + R"({"event":"throughput","server_instance_id":"a","timestamp_unix_ms":1601,"scheduler":{"waiting":2,"running":1,"prefilling":0}})" + "\n"; + const auto r = analyze_request_log_jsonl(jsonl, "mem"); + bool saw_sat = false, saw_limit = false, saw_content = false, saw_ttft = false; + for (const auto& it : r.at("insights")) { + const auto id = it.at("id").get(); + if (id.find("latency.") == 0 && id.find("ttft") == std::string::npos) { + saw_sat = true; + f += check(it.at("availability") == "available" && it.at("confidence") == "measured", + "saturation is measured"); + f += check(it.at("measured_over").at("requests") == 1, "measured_over.requests is 1"); + f += check(it.at("evidence").at("queued") == 1, "0.5s queue wait classifies queued"); + } + if (id.find("ttft") != std::string::npos) { + saw_ttft = true; + f += check(it.at("availability") == "available", "ttft split is measured"); + f += check(it.at("evidence").contains("mean_prepare_s") && + it.at("evidence").contains("mean_prefill_s"), + "ttft evidence has the split"); + } + if (id == "client.output_limit_while_thinking") { + saw_limit = true; + f += check(it.at("evidence").at("output_limit_thinking") == 1, "output_limit counted"); + f += check(it.at("evidence").at("sample_request_ids").at(0) == 1, "request_id evidence"); + } + if (id == "client.content_fields") { + saw_content = true; + f += check(it.at("availability") == "unavailable", + "content fields unavailable, not fabricated"); + f += check(it.at("measured_over").at("requests") == 1, + "unavailable measured_over is examined count, not 0"); + f += check(!it.contains("recommendation"), "empty recommendation is omitted"); + } + } + f += check(saw_sat && saw_limit && saw_content && saw_ttft, "required insight ids present"); + return f; +} + +int test_admin_vram_markers() { + using namespace ninfer::supervisor; + AdminVramCursor c; + std::string kind; + int f = 0; + f += check(!c.observe("", "", kind), "first poll is baseline, does not fire"); + f += check(c.observe("release", "seed store released", kind) && kind == "vram_release", + "empty -> release fires vram_release"); + f += check(c.observe("reclaim", "seed store reclaimed", kind) && kind == "vram_reclaim", + "release -> reclaim fires vram_reclaim"); + f += check(!c.observe("reclaim", "seed store reclaimed", kind), "unchanged does not fire"); + return f; +} + +int test_insights_pinned_tier() { + using namespace ninfer::supervisor; + nlohmann::json report = {{"insights", nlohmann::json::array()}}; + nlohmann::json admin = { + {"last_transition", ""}, + {"last_reason", ""}, + {"tiers", + nlohmann::json::array( + {{{"name", "seed"}, + {"min_bytes", 4294967296ull}, + {"max_bytes", 4294967296ull}, + {"reclaimable_bytes", 0}, + {"released", false}}})}}; + append_admin_vram_insights(report, admin, ""); + int f = 0; + bool saw = false; + for (const auto& it : report.at("insights")) { + if (it.at("id") == "vram.tier_pinned_unreleasable") { + saw = true; + f += check(it.at("severity") == "warning", "pinned tier is warning"); + f += check(it.at("availability") == "available", "config trap is measured"); + } + } + f += check(saw, "pinned min==max insight present"); + + nlohmann::json elastic_report = {{"insights", nlohmann::json::array()}}; + nlohmann::json elastic = { + {"last_transition", ""}, + {"last_reason", ""}, + {"tiers", + nlohmann::json::array( + {{{"name", "seed"}, + {"min_bytes", 0}, + {"max_bytes", 4294967296ull}, + {"reclaimable_bytes", 4294967296ull}, + {"released", false}}})}}; + append_admin_vram_insights(elastic_report, elastic, ""); + bool elastic_fired = false; + for (const auto& it : elastic_report.at("insights")) { + if (it.at("id") == "vram.tier_pinned_unreleasable") { elastic_fired = true; } + } + f += check(!elastic_fired, "elastic min!=max does not fire pinned insight"); + return f; +} + +int test_insights_prefix() { + using namespace ninfer::supervisor; + const char* jsonl = + R"({"event":"request_start","server_instance_id":"a","timestamp_unix_ms":1,"request":{"request_id":1,"message_count":1}})" + "\n" + R"({"event":"request_done","server_instance_id":"a","timestamp_unix_ms":2,"request":{"request_id":1,"message_count":1},"result":{"prefix_reuse_path":"full_reset","prompt_tokens":100,"prefix_cache_hit_tokens":0},"timings_seconds":{"total":0.05,"prepare":0.01,"prefill":0.02,"decode":0.02,"ttft":0.03,"vision":0}})" + "\n" + R"({"event":"request_start","server_instance_id":"a","timestamp_unix_ms":3,"request":{"request_id":2,"message_count":5}})" + "\n" + R"({"event":"request_done","server_instance_id":"a","timestamp_unix_ms":4,"request":{"request_id":2,"message_count":5},"result":{"prefix_reuse_path":"full_reset","prompt_tokens":7749,"prefix_cache_hit_tokens":0},"timings_seconds":{"total":0.05,"prepare":0.01,"prefill":0.02,"decode":0.02,"ttft":0.03,"vision":0}})" + "\n"; + const auto r = analyze_request_log_jsonl(jsonl, "mem"); + int f = 0; + bool mix = false, miss = false; + for (const auto& it : r.at("insights")) { + const auto id = it.at("id").get(); + if (id == "prefix.reuse_mix") { + mix = true; + f += check(it.at("evidence").at("full_reset_single_turn") == 1, "single-turn reset"); + f += check(it.at("evidence").at("full_reset_multi_turn") == 1, "multi-turn reset"); + f += check(it.at("availability") == "available", "mix is measured"); + } + if (id == "prefix.multiturn_full_reset") { + miss = true; + f += check(it.at("severity") == "warning", "multi-turn reset is a warning"); + f += check(it.at("evidence").at("samples").at(0).at("prompt_tokens") == 7749, + "live-shaped sample"); + f += check(it.at("measured_over").at("requests") == 2, "examined both dones"); + } + } + f += check(mix && miss, "prefix insights present"); + return f; +} + +int test_jsonl_event_key() { + using namespace ninfer::supervisor; + int f = 0; + f += check(jsonl_event_is(R"({"event":"request_done","result":{}})", "request_done"), + "compact event key"); + f += check(jsonl_event_is(R"({"event": "request_done"})", "request_done"), "spaced event key"); + f += check(!jsonl_event_is(R"({"type":"request_done"})", "request_done"), + "type key is not the event key"); + f += check(!jsonl_event_is(R"({"event":"server_start"})", "request_done"), "other event"); + return f; +} + +int test_series_persist() { + using namespace ninfer::supervisor; + int f = 0; + VramSample s{}; + f += check(parse_series_sample_line(R"({"t_ms":10,"budget_bytes":20,"nvidia_used_bytes":30})", s) && + s.t_ms == 10 && s.budget_bytes == 20 && s.nvidia_used_bytes == 30, + "parse sample"); + VramSeriesEvent e{}; + f += check(parse_series_event_line(R"({"t_ms":11,"kind":"vram_release","label":"seed store released"})", e) && + e.kind == "vram_release", + "parse event"); + VramSeriesRing r(8); + r.load_jsonl("{\"t_ms\":1,\"budget_bytes\":2,\"nvidia_used_bytes\":3}\n" + "{\"t_ms\":4,\"kind\":\"engine_up\",\"label\":\"health 200\"}\n"); + f += check(r.size() == 1 && r.events().size() == 1 && r.events()[0].kind == "engine_up", + "load_jsonl restores samples and events"); + return f; +} + +int test_series_ring() { + ninfer::supervisor::VramSeriesRing r(3); + int f = 0; + r.push({1, 10, 4}); + r.push({2, 20, 5}); + r.push({3, 30, 6}); + r.push({4, 40, 7}); + const auto s = r.samples(); + f += check(s.size() == 3 && s[0].t_ms == 2 && s[2].t_ms == 4 && s[2].budget_bytes == 40, + "ring drops oldest, keeps raw values"); + r.push_event({4, "admin_vram", "release"}, 2); + r.push_event({5, "engine_down", "health 0"}, 2); + r.push_event({6, "engine_up", "health 200"}, 2); + const auto e = r.events(); + f += check(e.size() == 2 && e[0].kind == "engine_down" && e[1].kind == "engine_up", + "event ring cap"); + return f; +} + +int test_health_threshold() { + ninfer::supervisor::RestartPolicy p; + p.health_fail_threshold = 3; + ninfer::supervisor::RestartGate g(p); + int f = 0; + f += check(!g.note_health_fail() && !g.note_health_fail(), "below threshold"); + f += check(g.note_health_fail(), "threshold trips restart"); + g.note_healthy(); + f += check(!g.note_health_fail(), "healthy clears fail count"); + return f; +} + +} // namespace + +int main() { + int failures = 0; + failures += test_loopback(); + failures += test_crash_loop(); + failures += test_backoff(); + failures += test_config_bind(); + failures += test_host_header(); + failures += test_nvidia_csv(); + failures += test_kv_line(); + failures += test_monitor_only_config(); + failures += test_insights_honesty(); + failures += test_insights_prefix(); + failures += test_admin_vram_markers(); + failures += test_insights_pinned_tier(); + failures += test_jsonl_event_key(); + failures += test_series_persist(); + failures += test_series_ring(); + failures += test_health_threshold(); + if (failures == 0) { std::cout << "ok\n"; } + return failures == 0 ? 0 : 1; +} diff --git a/tests/test_openai_schema.cpp b/tests/test_openai_schema.cpp index 6fef3ecf0f..bca806881a 100644 --- a/tests/test_openai_schema.cpp +++ b/tests/test_openai_schema.cpp @@ -11,6 +11,7 @@ #include #include +#include #include #include #include @@ -176,6 +177,122 @@ int test_preserve_thinking_options() { return failures; } +bool thinking_disabled(const ResolvedPromptSemantics& semantics) { + return !semantics.enable_thinking && !semantics.reasoning_effort; +} + +int test_enable_thinking_dialect() { + const Json base = { + {"model", "m"}, + {"messages", Json::array({Json{{"role", "user"}, {"content", "hello"}}})}, + }; + int failures = 0; + + const ResolvedPromptSemantics omitted = + resolve_prompt_semantics(parse_chat_completion_request(base, default_limits()), + default_server(), effort_capabilities()); + failures += check(omitted.enable_thinking && !omitted.reasoning_effort, + "omitted thinking did not use the server default on"); + + Json none = base; + none["reasoning_effort"] = "none"; + const ResolvedPromptSemantics none_semantics = + resolve_prompt_semantics(parse_chat_completion_request(none, default_limits()), + default_server(), effort_capabilities()); + failures += check(thinking_disabled(none_semantics), + "reasoning_effort none did not disable thinking"); + + Json kwargs = base; + kwargs["chat_template_kwargs"] = Json{{"enable_thinking", false}}; + const GenerationRequest kwargs_request = + parse_chat_completion_request(kwargs, default_limits()); + failures += check(kwargs_request.enable_thinking == false, + "chat_template_kwargs enable_thinking was not parsed"); + const ResolvedPromptSemantics kwargs_semantics = + resolve_prompt_semantics(kwargs_request, default_server(), effort_capabilities()); + failures += check(thinking_disabled(kwargs_semantics), + "chat_template_kwargs enable_thinking false did not match reasoning_effort none"); + + Json top = base; + top["enable_thinking"] = false; + const GenerationRequest top_request = parse_chat_completion_request(top, default_limits()); + failures += + check(top_request.enable_thinking == false, "top-level enable_thinking was not parsed"); + const ResolvedPromptSemantics top_semantics = + resolve_prompt_semantics(top_request, default_server(), effort_capabilities()); + failures += check(thinking_disabled(top_semantics), + "top-level enable_thinking false did not match reasoning_effort none"); + + Json both = kwargs; + both["enable_thinking"] = false; + failures += check(parse_chat_completion_request(both, default_limits()).enable_thinking == false, + "matching enable_thinking values were rejected"); + + Json conflict = kwargs; + conflict["enable_thinking"] = true; + failures += check(api_code([&] { + (void)parse_chat_completion_request(conflict, default_limits()); + }) == "conflicting_template_option", + "conflicting enable_thinking values were accepted"); + + Json on = base; + on["chat_template_kwargs"] = Json{{"enable_thinking", true}}; + ServeOptions no_think = default_server(); + no_think.enable_thinking = false; + failures += check(resolve_prompt_semantics(parse_chat_completion_request(on, default_limits()), + no_think, effort_capabilities()) + .enable_thinking, + "chat_template_kwargs enable_thinking true did not override server default off"); + + Json misspelled = base; + misspelled["chat_template_kwargs"] = Json{{"enable_thinkng", false}}; + failures += check(api_code([&] { + (void)parse_chat_completion_request(misspelled, default_limits()); + }) == "chat_template_option_not_supported", + "misspelled enable_thinking was not rejected"); + + Json bad = base; + bad["chat_template_kwargs"] = Json{{"enable_thinking", "no"}}; + failures += check( + throws_api([&] { (void)parse_chat_completion_request(bad, default_limits()); }), + "non-boolean chat_template_kwargs.enable_thinking was accepted"); + + Json combined = base; + combined["chat_template_kwargs"] = Json{{"enable_thinking", false}, {"preserve_thinking", true}}; + const GenerationRequest combined_request = + parse_chat_completion_request(combined, default_limits()); + failures += check(combined_request.enable_thinking == false && + combined_request.preserve_thinking == true, + "enable_thinking and preserve_thinking were not accepted together"); + ServeOptions preserved = default_server(); + preserved.preserve_thinking = true; + const ResolvedPromptSemantics combined_semantics = + resolve_prompt_semantics(combined_request, default_server(), effort_capabilities()); + failures += check(thinking_disabled(combined_semantics) && combined_semantics.preserve_thinking, + "enable_thinking false cleared request preserve_thinking"); + const ResolvedPromptSemantics server_preserved = + resolve_prompt_semantics(kwargs_request, preserved, effort_capabilities()); + failures += check(thinking_disabled(server_preserved) && server_preserved.preserve_thinking, + "enable_thinking false cleared server preserve_thinking"); + + Json agree = none; + agree["enable_thinking"] = false; + failures += check(thinking_disabled(resolve_prompt_semantics( + parse_chat_completion_request(agree, default_limits()), default_server(), + effort_capabilities())), + "enable_thinking false with reasoning_effort none was rejected"); + + Json effort_conflict = none; + effort_conflict["enable_thinking"] = true; + failures += check(api_code([&] { + (void)resolve_prompt_semantics( + parse_chat_completion_request(effort_conflict, default_limits()), + default_server(), effort_capabilities()); + }) == "conflicting_template_option", + "enable_thinking true with reasoning_effort none was accepted"); + return failures; +} + int test_reasoning_effort() { const Json base = { {"model", "m"}, @@ -437,6 +554,13 @@ int test_parse_function_tools_and_choices() { failures += check(throws_api([&] { (void)parse_chat_completion_request(unknown, default_limits()); }), "unknown named tool_choice rejected"); + + // Duplicate function tool names are rejected (matches Responses behavior). + Json dup = base; + dup["tools"] = Json::array({tool, tool}); + failures += + check(throws_api([&] { (void)parse_chat_completion_request(dup, default_limits()); }), + "duplicate function tool names rejected"); return failures; } @@ -550,6 +674,12 @@ int test_response_serialization() { failures += check(j.at("usage").at("prompt_tokens") == 10, "usage prompt_tokens"); failures += check(j.at("usage").at("completion_tokens") == 3, "usage completion_tokens"); failures += check(j.at("usage").at("total_tokens") == 13, "usage total_tokens"); + failures += check(j.at("usage").at("prompt_tokens_details").at("cached_tokens") == 0, + "default cached_tokens is additive and zero"); + failures += check(j.at("usage").at("prefix_cache_hit_tokens") == 0, + "default prefix_cache_hit_tokens matches the log field"); + failures += check(j.at("usage").at("prefix_reuse_path") == "full_reset", + "default prefix_reuse_path matches the log field"); // Non-empty reasoning is attached as message.reasoning_content, content stays answer-only. const Json jr = Json::parse(make_chat_completion_response("id-2", "m", 111, "the answer", @@ -639,11 +769,59 @@ int test_chunk_serialization() { failures += check(usage_chunk.at("usage").at("prompt_tokens") == 2, "usage chunk prompt_tokens"); failures += check(usage_chunk.at("usage").at("total_tokens") == 7, "usage chunk total"); + failures += check(usage_chunk.at("usage").at("prompt_tokens_details").at("cached_tokens") == 0, + "usage chunk cached_tokens additive"); failures += check(sse_done() == "data: [DONE]\n\n", "done sentinel"); return failures; } +int test_usage_prefix_observability() { + int failures = 0; + CompletionUsage usage; + usage.prompt_tokens = 100; + usage.completion_tokens = 8; + usage.cached_prompt_tokens = 60; + usage.prefix_reuse_path = ninfer::PrefixReusePath::SeedPrefixCache; + const Json j = + Json::parse(make_chat_completion_response("id", "m", 1, "pong", "", "stop", usage)); + const Json& u = j.at("usage"); + failures += check(u.at("prompt_tokens") == 100 && u.at("completion_tokens") == 8 && + u.at("total_tokens") == 108, + "OpenAI usage totals unchanged"); + failures += check(u.at("prompt_tokens_details").at("cached_tokens") == 60, + "OpenAI cached_tokens is a subset of prompt_tokens"); + failures += check(u.at("prefix_cache_hit_tokens") == 60, + "log-named cached count matches prompt_tokens_details"); + failures += check(u.at("prefix_reuse_path") == "seed_prefix", "log-named reuse path"); + failures += check(u.at("total_tokens") == u.at("prompt_tokens").get() + + u.at("completion_tokens").get(), + "cached_tokens is not an addend of total_tokens"); + + CompletionUsage clamped; + clamped.prompt_tokens = 10; + clamped.cached_prompt_tokens = 99; + clamped.prefix_reuse_path = ninfer::PrefixReusePath::AppendAtFrontier; + const Json c = + Json::parse(make_chat_completion_response("id", "m", 1, "x", "", "stop", clamped)).at("usage"); + failures += check(c.at("prompt_tokens_details").at("cached_tokens") == 10 && + c.at("prefix_cache_hit_tokens") == 10, + "cached_tokens clamped to prompt_tokens"); + failures += check(c.at("prefix_reuse_path") == "append_frontier", "append_frontier wire name"); + + for (const auto& [path, name] : + std::array, 5>{ + {{ninfer::PrefixReusePath::FullReset, "full_reset"}, + {ninfer::PrefixReusePath::AppendAtFrontier, "append_frontier"}, + {ninfer::PrefixReusePath::RestoreTurnCheckpoint, "restore_turn_checkpoint"}, + {ninfer::PrefixReusePath::RestoreResponseCheckpoint, "restore_response_checkpoint"}, + {ninfer::PrefixReusePath::SeedPrefixCache, "seed_prefix"}}}) { + failures += check(std::string(prefix_reuse_path_name(path)) == name, + std::string("reuse path wire name ") + name); + } + return failures; +} + int test_tool_chunk_serialization() { int failures = 0; const std::vector calls = { @@ -668,16 +846,29 @@ int test_tool_chunk_serialization() { } int test_models_and_error() { - int failures = 0; - const Json list = Json::parse(make_models_list("qwen3.6-27b", 1)); + int failures = 0; + constexpr std::uint32_t configured_context = 131072; + const Json list = Json::parse(make_models_list("qwen3.6-27b", 1, configured_context)); failures += check(list.at("object") == "list", "models list object"); failures += check(list.at("data").at(0).at("id") == "qwen3.6-27b", "models list id"); failures += check(list.at("data").at(0).at("object") == "model", "models list entry object"); failures += check(list.at("data").at(0).at("owned_by") == "ninfer", "models list owner"); + failures += check(list.at("data").at(0).at("max_model_len") == configured_context, + "models list configured context"); - const Json one = Json::parse(make_model_object("qwen3.6-27b", 1)); + const Json one = Json::parse(make_model_object("qwen3.6-27b", 1, configured_context)); failures += check(one.at("id") == "qwen3.6-27b" && one.at("object") == "model", "model object"); failures += check(one.at("owned_by") == "ninfer", "model owner"); + failures += check(one.at("max_model_len") == configured_context, + "model object configured context"); + + constexpr std::uint32_t long_context = 262144; + const Json long_list = Json::parse(make_models_list("qwen3.8-27b", 1, long_context)); + failures += check(long_list.at("data").at(0).at("max_model_len") == long_context, + "models list 262144 context"); + failures += check(Json::parse(make_model_object("qwen3.8-27b", 1, long_context)) + .at("max_model_len") == long_context, + "model object 262144 context"); ApiError error; error.status = 400; @@ -698,18 +889,74 @@ int test_finish_reason_wire() { "stop token wire"); failures += check(std::string(finish_reason_wire(ninfer::FinishReason::OutputLimit)) == "length", - "output limit wire"); + "output limit wire"); failures += check(std::string(finish_reason_wire(ninfer::FinishReason::Cancelled)) == "stop", "cancelled maps to stop"); return failures; } +int test_parse_content_parts_allowed_types() { + int failures = 0; + ChatTurn turn; + const Json valid_text = Json::array({Json{{"type", "text"}, {"text", "hello"}}}); + parse_content_parts(valid_text, turn, 0, {"text"}); + failures += check(turn.content.size() == 1, "text parsed with allowed_types"); + + ChatTurn media_turn; + const Json media_parts = Json::array({ + Json{{"type", "text"}, {"text", "result"}}, + Json{{"type", "image_url"}, {"image_url", Json{{"url", "data:image/png;base64,AA=="}}}} + }); + parse_content_parts(media_parts, media_turn, 0, {"text", "image_url"}); + failures += check(media_turn.content.size() == 2, "text+image parsed with allowed_types"); + + bool rejected = false; + std::string error_message; + try { + ChatTurn rejected_turn; + parse_content_parts(media_parts, rejected_turn, 0, {"text"}); + } catch (const ApiException& e) { + rejected = true; + error_message = e.error().message; + } + failures += check(rejected, "disallowed media part was rejected"); + failures += check(error_message.find("content parts must have type 'text'") != std::string::npos, + "error message lists allowed types"); + + return failures; +} + +int test_parse_tool_message_content_parts() { + int failures = 0; + const Json body = { + {"model", "m"}, + {"messages", Json::array({ + Json{{"role", "user"}, {"content", "run screenshot"}}, + Json{{"role", "assistant"}, {"content", nullptr}, {"tool_calls", Json::array({ + Json{{"id", "call_1"}, {"type", "function"}, {"function", Json{{"name", "screenshot"}, {"arguments", "{}"}}}} + })}}, + Json{{"role", "tool"}, {"tool_call_id", "call_1"}, {"content", Json::array({ + Json{{"type", "text"}, {"text", "captured:"}}, + Json{{"type", "image_url"}, {"image_url", Json{{"url", "data:image/png;base64,AA=="}}}} + })}} + })} + }; + const GenerationRequest req = parse_chat_completion_request(body, default_limits()); + failures += check(req.messages.size() == 3, "parsed 3 messages"); + failures += check(req.messages[2].role == ninfer::ChatRole::Tool, "third message is tool role"); + failures += check(req.messages[2].content.size() == 2, "tool message has 2 content parts"); + failures += check(req.messages[2].content[0].kind == ContentKind::Text, "tool content part 0 is text"); + failures += check(req.messages[2].content[1].kind == ContentKind::Image, "tool content part 1 is image"); + return failures; +} + } // namespace int main() { int failures = 0; failures += test_parse_string_content(); failures += test_preserve_thinking_options(); + failures += test_enable_thinking_dialect(); failures += test_reasoning_effort(); failures += test_parse_parts_and_flatten(); failures += test_instruction_roles_preserved(); @@ -722,9 +969,12 @@ int main() { failures += test_response_serialization(); failures += test_tool_response_serialization(); failures += test_chunk_serialization(); + failures += test_usage_prefix_observability(); failures += test_tool_chunk_serialization(); failures += test_models_and_error(); failures += test_finish_reason_wire(); + failures += test_parse_content_parts_allowed_types(); + failures += test_parse_tool_message_content_parts(); if (failures == 0) { std::cout << "ok\n"; } return failures == 0 ? 0 : 1; } diff --git a/tests/test_request_log.cpp b/tests/test_request_log.cpp index d557ace416..15f0d2b7fa 100644 --- a/tests/test_request_log.cpp +++ b/tests/test_request_log.cpp @@ -55,6 +55,7 @@ int main() { options.enable_vision = false; options.allow_prefix_reuse = false; options.preserve_thinking = true; + options.tolerant_tool_calls = true; options.sampling_overrides.temperature = 0.6F; options.startup_argv = {"ninfer-serve", options.artifact_path, "--api-key", ""}; @@ -94,6 +95,7 @@ int main() { memory.planned_slack_bytes = 100; memory.cuda_graph_allowance_bytes = 600; memory.cuda_graph_observed_bytes = 550; + memory.prefix_cache_bytes = 4096ULL << 20; memory.kv_payload_bytes = 400; ServerLogEnvironment environment; @@ -141,6 +143,8 @@ int main() { check(server.at("engine").at("prefix_reuse") == false, "prefix-reuse state missing"); failures += check(server.at("server").at("default_preserve_thinking") == true, "server preserve-thinking default missing"); + failures += check(server.at("server").at("tolerant_tool_calls") == true, + "tolerant tool-call setting missing"); failures += check(server.at("sampling_defaults").at("thinking").at("temperature") == 1.0 && server.at("sampling_defaults").at("non_thinking").at("presence_penalty") == 1.5, @@ -162,7 +166,8 @@ int main() { server.at("memory").at("available_after_startup_bytes") == 180 && server.at("memory").at("kv_capacity_headroom_bytes") == 0 && server.at("memory").at("planned_slack_bytes") == 100 && - server.at("memory").at("cuda_graph_observed_bytes") == 550, + server.at("memory").at("cuda_graph_observed_bytes") == 550 && + server.at("memory").at("prefix_cache_bytes") == (4096ULL << 20), "adaptive KV memory ledger missing"); failures += check(server.dump().find("must-not-appear") == std::string::npos, "server JSON leaked the API key"); diff --git a/tests/test_responses_schema.cpp b/tests/test_responses_schema.cpp index bd685df446..81593cd5b8 100644 --- a/tests/test_responses_schema.cpp +++ b/tests/test_responses_schema.cpp @@ -246,6 +246,75 @@ int test_preserve_thinking_options_and_inheritance() { return failures; } +bool thinking_disabled(const ResolvedPromptSemantics& semantics) { + return !semantics.enable_thinking && !semantics.reasoning_effort; +} + +int test_enable_thinking_dialect() { + const Json base = {{"model", "m"}, {"input", "hello"}, {"max_output_tokens", 32}}; + int failures = 0; + + Json none = base; + none["reasoning"] = Json{{"effort", "none"}}; + const ResolvedPromptSemantics none_semantics = resolve_prompt_semantics( + parse_responses_request(none, limits()).generation, ServeOptions{}, effort_capabilities()); + failures += + check(thinking_disabled(none_semantics), "Responses reasoning.effort none did not disable thinking"); + + Json kwargs = base; + kwargs["chat_template_kwargs"] = Json{{"enable_thinking", false}}; + const GenerationRequest kwargs_request = parse_responses_request(kwargs, limits()).generation; + failures += check(kwargs_request.enable_thinking == false, + "Responses chat_template_kwargs enable_thinking was not parsed"); + failures += check(thinking_disabled(resolve_prompt_semantics(kwargs_request, ServeOptions{}, + effort_capabilities())), + "Responses kwargs enable_thinking false did not match reasoning.effort none"); + + Json top = base; + top["enable_thinking"] = false; + failures += check(parse_responses_request(top, limits()).generation.enable_thinking == false, + "Responses top-level enable_thinking was not parsed"); + + Json both = kwargs; + both["enable_thinking"] = false; + failures += + check(parse_responses_request(both, limits()).generation.enable_thinking == false, + "Responses matching enable_thinking values were rejected"); + + Json conflict = kwargs; + conflict["enable_thinking"] = true; + failures += check(api_code([&] { (void)parse_responses_request(conflict, limits()); }) == + "conflicting_template_option", + "Responses conflicting enable_thinking values were accepted"); + + Json misspelled = base; + misspelled["chat_template_kwargs"] = Json{{"enable_thinkng", false}}; + failures += check(api_code([&] { (void)parse_responses_request(misspelled, limits()); }) == + "chat_template_option_not_supported", + "Responses misspelled enable_thinking was not rejected"); + + Json combined = base; + combined["chat_template_kwargs"] = Json{{"enable_thinking", false}, {"preserve_thinking", true}}; + const GenerationRequest combined_request = + parse_responses_request(combined, limits()).generation; + failures += check(combined_request.enable_thinking == false && + combined_request.preserve_thinking == true, + "Responses enable_thinking and preserve_thinking were not accepted together"); + + Json tokens = {{"model", "m"}, {"input", "hello"}}; + tokens["chat_template_kwargs"] = Json{{"enable_thinking", false}}; + failures += check(parse_response_input_tokens_request(tokens, limits()) + .generation.enable_thinking == false, + "Responses input_tokens rejected chat_template_kwargs.enable_thinking"); + tokens.erase("chat_template_kwargs"); + tokens["enable_thinking"] = false; + failures += + check(parse_response_input_tokens_request(tokens, limits()).generation.enable_thinking == + false, + "Responses input_tokens rejected top-level enable_thinking"); + return failures; +} + int test_typed_items_and_tools() { const Json function = {{"type", "function"}, {"name", "weather"}, @@ -349,6 +418,12 @@ int test_explicit_rejections() { failures += check(api_code([&] { (void)parse_responses_request(too_small, limits()); }) == "invalid_value", "OpenAI minimum max_output_tokens enforced"); + + Json dup = base; + dup["tools"] = Json::array({Json{{"type", "function"}, {"name", "f"}, {"parameters", Json::object()}, {"strict", false}}, + Json{{"type", "function"}, {"name", "f"}, {"parameters", Json::object()}, {"strict", false}}}); + failures += check(throws_api([&] { (void)parse_responses_request(dup, limits()); }), + "duplicate function tool names rejected"); return failures; } @@ -391,7 +466,9 @@ int test_response_object() { failures += check(response.at("usage").at("input_tokens_details").at("cached_tokens") == 4 && response.at("usage").at("output_tokens_details").at("reasoning_tokens") == 3 && - response.at("usage").at("total_tokens") == 18, + response.at("usage").at("total_tokens") == 18 && + response.at("usage").at("prefix_cache_hit_tokens") == 4 && + response.at("usage").at("prefix_reuse_path") == "full_reset", "Responses usage details serialized"); failures += check(built.output_history.size() == 1 && built.output_history[0].reasoning_content == "thought" && @@ -520,6 +597,7 @@ int main() { failures += test_basic_request(); failures += test_instruction_message_order(); failures += test_preserve_thinking_options_and_inheritance(); + failures += test_enable_thinking_dialect(); failures += test_reasoning_effort(); failures += test_typed_items_and_tools(); failures += test_explicit_rejections(); diff --git a/tests/test_serve_corpus.py b/tests/test_serve_corpus.py index fa55b8897c..53bf2f947d 100644 --- a/tests/test_serve_corpus.py +++ b/tests/test_serve_corpus.py @@ -9,15 +9,15 @@ ) -def test_request_log_v9_identity_is_accepted() -> None: +def test_request_log_v10_identity_is_accepted() -> None: current = { "artifact_type": "ninfer_serve_request_log", - "schema_version": 9, + "schema_version": 10, "event": "server_start", } require_server_log_identity(current, "server_start") - stale = dict(current, schema_version=8) + stale = dict(current, schema_version=9) with pytest.raises(CampaignError): require_server_log_identity(stale, "server_start") diff --git a/tests/test_serve_options.cpp b/tests/test_serve_options.cpp index 65231dff94..06ba38f24f 100644 --- a/tests/test_serve_options.cpp +++ b/tests/test_serve_options.cpp @@ -33,6 +33,8 @@ int main() { failures += check(!defaults.preserve_thinking, "thinking history is unexpectedly preserved by default"); failures += check(!defaults.enable_vision, "Vision is not disabled by default"); + failures += check(!defaults.tolerant_tool_calls, + "tolerant tool-call recovery is not disabled by default"); failures += check(defaults.request_log_jsonl.empty(), "request JSONL logging is not disabled by default"); failures += check(defaults.log_stats_interval_ms == 5000, @@ -41,6 +43,10 @@ int main() { defaults.media_live_bytes == ninfer::kDefaultMediaLiveBytes && defaults.media_preprocess_threads == 0, "media preparation resource defaults mismatch"); + failures += check(defaults.prefix_cache_bytes == 0 && defaults.prefix_cache_min_bytes == 0 && + defaults.prefix_cache_max_bytes == 0 && + defaults.vram_idle_release_after_s == 0 && !defaults.vram_observe_only, + "prefix cache is not disabled by default"); failures += check(defaults.kv_capacity.mode == ninfer::KvCapacityMode::Explicit && defaults.kv_capacity.explicit_tokens == defaults.max_context, "default KV capacity does not follow max context"); @@ -111,6 +117,7 @@ int main() { "--log-stats-interval-ms", "0", "--preserve-thinking", + "--tolerant-tool-calls", "--media-cache-mib", "256", "--media-live-mib", @@ -122,6 +129,8 @@ int main() { failures += check(configured.enable_vision, "--vision did not enable Vision"); failures += check(configured.preserve_thinking, "--preserve-thinking did not reach serving options"); + failures += check(configured.tolerant_tool_calls, + "--tolerant-tool-calls did not reach serving options"); failures += check(configured.max_concurrency == 4, "--max-concurrency did not reach serving options"); failures += check(configured.max_context == 4096 && @@ -139,6 +148,58 @@ int main() { configured.media_preprocess_threads == 6, "media preparation limits did not reach serving options"); + const ServeOptions prefix_disabled = + parse({"ninfer-serve", "model.ninfer", "--prefix-cache-mib", "0"}); + failures += check(prefix_disabled.prefix_cache_bytes == 0, + "--prefix-cache-mib 0 did not keep the seed store disabled"); + const ServeOptions prefix_enabled = + parse({"ninfer-serve", "model.ninfer", "--prefix-cache-mib", "4096"}); + failures += check(prefix_enabled.prefix_cache_bytes == (4096ULL << 20) && + prefix_enabled.prefix_cache_min_bytes == (4096ULL << 20) && + prefix_enabled.prefix_cache_max_bytes == (4096ULL << 20), + "--prefix-cache-mib 4096 did not reserve a fixed 4 GiB range"); + + const ServeOptions prefix_range = + parse({"ninfer-serve", "model.ninfer", "--prefix-cache-mib-min", "0", + "--prefix-cache-mib-max", "4096"}); + failures += check(prefix_range.prefix_cache_bytes == (4096ULL << 20) && + prefix_range.prefix_cache_min_bytes == 0 && + prefix_range.prefix_cache_max_bytes == (4096ULL << 20), + "prefix-cache range did not boot at max with min 0"); + + const ServeOptions kv_range = + parse({"ninfer-serve", "model.ninfer", "--max-context", "8192", "--kv-capacity-min", "8192", + "--kv-capacity-max", "65536"}); + failures += check(kv_range.kv_capacity.mode == ninfer::KvCapacityMode::Explicit && + kv_range.kv_capacity.explicit_tokens == 65536 && + kv_range.kv_capacity_min_tokens == 8192 && + kv_range.kv_capacity_max_tokens == 65536, + "kv-capacity range did not boot at max"); + + const ServeOptions idle = parse({"ninfer-serve", "model.ninfer", "--vram-idle-release-after-s", + "30", "--vram-observe-only"}); + failures += check(idle.vram_idle_release_after_s == 30 && idle.vram_observe_only, + "idle-release and observe-only flags did not parse"); + failures += check(!defaults.enable_admin_vram, "admin VRAM routes are not disabled by default"); + bool admin_without_key_rejected = false; + try { + (void)parse({"ninfer-serve", "model.ninfer", "--admin-vram"}); + } catch (const std::invalid_argument&) { admin_without_key_rejected = true; } + failures += check(admin_without_key_rejected, "--admin-vram without --api-key was accepted"); + const ServeOptions admin = + parse({"ninfer-serve", "model.ninfer", "--admin-vram", "--api-key", "secret"}); + failures += check(admin.enable_admin_vram && admin.api_key == "secret", + "--admin-vram with --api-key did not enable admin routes"); + failures += check(serve_usage_text("ninfer-serve").find("--admin-vram") != std::string::npos, + "serve help omits --admin-vram"); + + bool inverted_prefix_rejected = false; + try { + (void)parse({"ninfer-serve", "model.ninfer", "--prefix-cache-mib-min", "8", + "--prefix-cache-mib-max", "4"}); + } catch (const std::invalid_argument&) { inverted_prefix_rejected = true; } + failures += check(inverted_prefix_rejected, "inverted prefix-cache range was accepted"); + const ServeOptions response_store = parse({"ninfer-serve", "model.ninfer", "--response-store-max-records", "42", "--response-store-max-mib", "8"}); @@ -189,6 +250,9 @@ int main() { failures += check(serve_usage_text("ninfer-serve").find("--preserve-thinking") != std::string::npos, "serve help omits --preserve-thinking"); + failures += check(serve_usage_text("ninfer-serve").find("--tolerant-tool-calls") != + std::string::npos, + "serve help omits --tolerant-tool-calls"); failures += check(serve_usage_text("ninfer-serve").find("--vision") != std::string::npos, "serve help omits --vision"); failures += @@ -199,6 +263,10 @@ int main() { "serve help omits media preparation controls"); failures += check(serve_usage_text("ninfer-serve").find("--kv-capacity") != std::string::npos, "serve help omits --kv-capacity"); + failures += check(serve_usage_text("ninfer-serve") + .find("(bounded by max-context * max-concurrency)") != + std::string::npos, + "serve help omits auto kv-capacity bounding explanation"); failures += check(serve_usage_text("ninfer-serve").find("--response-store-max-mib") != std::string::npos, "serve help omits Responses store limits"); @@ -232,8 +300,16 @@ int main() { secret_present = secret_present || argument == "do-not-log"; redaction_present = redaction_present || argument == ""; } - failures += check(!secret_present, "startup argv retained the API key"); - failures += check(redaction_present, "startup argv omitted the API-key redaction marker"); + failures += check(defaults.boot_watchdog_timeout_s == 120, + "boot watchdog timeout default mismatch"); + + const ServeOptions watchdog_opt = + parse({"ninfer-serve", "model.ninfer", "--boot-watchdog-timeout-s", "300"}); + failures += check(watchdog_opt.boot_watchdog_timeout_s == 300, + "--boot-watchdog-timeout-s did not parse value"); + failures += check(serve_usage_text("ninfer-serve").find("--boot-watchdog-timeout-s") != + std::string::npos, + "serve help omits --boot-watchdog-timeout-s"); if (failures == 0) { std::cout << "ok\n"; } return failures == 0 ? 0 : 1; diff --git a/tests/test_tool_call_parser.cpp b/tests/test_tool_call_parser.cpp index 8c96b73e1e..11dd77cfeb 100644 --- a/tests/test_tool_call_parser.cpp +++ b/tests/test_tool_call_parser.cpp @@ -1,9 +1,11 @@ #include "serve/tool_call_parser.h" +#include "serve/request.h" #include #include #include +#include namespace { @@ -16,30 +18,52 @@ int fail(const std::string& message) { int check(bool condition, const std::string& message) { return condition ? 0 : fail(message); } +// Build a ToolDefinition with the given name and a JSON Schema parameters +// object string, so tests exercise the real build_tool_param_type_map path +// rather than hand-fabricating the ToolParamTypeMap. +ninfer::serve::ToolDefinition make_tool(const std::string& name, const std::string& schema_json) { + ninfer::serve::ToolDefinition tool; + tool.name = name; + tool.parameters_json = schema_json; + return tool; +} + int test_single_call() { + // build_tool_param_type_map records only non-string types; city (string) + // and any unknown param are absent, so the parser preserves raw text. + ninfer::serve::ToolParamTypeMap map; + map["get_weather"]["days"] = {"integer"}; + const ninfer::serve::ParsedToolCallOutput parsed = ninfer::serve::parse_qwen_tool_call_output("Calling weather.\n" - "\n" + " \n" "\n" "\nParis\n\n" "\n2\n\n" "\n" "", - 64); + 64, map); int failures = 0; failures += check(parsed.is_tool_call_response, "single call parsed as tool response"); - failures += check(parsed.content == "Calling weather.", "content prefix trimmed"); + // The parser RETAINS the pre-call prefix; suppression happens in the caller, + // which alone knows whether those bytes were already streamed. + failures += check(parsed.content == "Calling weather.", + "parser retains the pre-call preamble for the caller to suppress"); failures += check(parsed.tool_calls.size() == 1, "one parsed call"); failures += check(parsed.tool_calls[0].id.rfind("call_", 0) == 0, "generated call id prefix"); failures += check(parsed.tool_calls[0].name == "get_weather", "function name parsed"); const Json args = Json::parse(parsed.tool_calls[0].arguments_json); failures += check(args.at("city") == "Paris", "string parameter parsed"); - failures += check(args.at("days") == 2, "number parameter parsed"); + failures += check(args.at("days") == 2, "integer-typed days deserialized to number"); return failures; } int test_multiple_calls_and_json_values() { + // payload is object => recorded; value is string => absent. + ninfer::serve::ToolParamTypeMap map; + map["first"]["payload"] = {"object"}; + const ninfer::serve::ParsedToolCallOutput parsed = ninfer::serve::parse_qwen_tool_call_output( "\n" "\n" @@ -51,7 +75,7 @@ int test_multiple_calls_and_json_values() { "\nplain text\n\n" "\n" "", - 64); + 64, map); int failures = 0; failures += check(parsed.is_tool_call_response, "multiple calls parsed as tool response"); @@ -66,10 +90,57 @@ int test_multiple_calls_and_json_values() { return failures; } +int test_string_param_keeps_numeric_looking_value() { + // priority is integer => recorded; taskId and status are string => + // absent (exactly what build_tool_param_type_map produces). The tool + // is known, yet its string-typed params still preserve raw text. + ninfer::serve::ToolParamTypeMap map; + map["TaskUpdate"]["priority"] = {"integer"}; + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "deleted\n" + "1\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "string-typed call parsed as tool response"); + failures += check(parsed.tool_calls.size() == 1, "one parsed string-typed call"); + failures += check(parsed.tool_calls[0].name == "TaskUpdate", "string-typed call name"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("status") == "deleted", "string status preserved"); + failures += check(args.at("taskId").is_string(), "taskId is a string, not a number"); + failures += check(args.at("taskId") == "1", "string-typed taskId keeps numeric-looking value"); + return failures; +} + +int test_unknown_param_defaults_to_string() { + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "7\n" + "\n" + "", + 64, {}); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "unknown-schema call parsed as tool response"); + failures += check(parsed.tool_calls.size() == 1, "one parsed unknown-schema call"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("count").is_string(), "unknown-schema count defaults to string"); + failures += check(args.at("count") == "7", "unknown-schema count value preserved"); + return failures; +} + int test_malformed_falls_back_to_text() { - const std::string text = "\n\n"; + const std::string text = " \n\n"; const ninfer::serve::ParsedToolCallOutput parsed = - ninfer::serve::parse_qwen_tool_call_output(text, 64); + ninfer::serve::parse_qwen_tool_call_output(text, 64, {}); int failures = 0; failures += check(!parsed.is_tool_call_response, "malformed xml is not tool response"); failures += check(parsed.content == text, "malformed xml preserved as text"); @@ -77,33 +148,40 @@ int test_malformed_falls_back_to_text() { return failures; } -int test_suffix_after_tool_falls_back_to_text() { - const std::string text = "\n" +int test_partial_recovery_trailing_suffix() { + const std::string text = " \n" "\n" "\nParis\n\n" "\n" "\n" "extra answer"; - const ninfer::serve::ParsedToolCallOutput parsed = - ninfer::serve::parse_qwen_tool_call_output(text, 64); int failures = 0; - failures += check(!parsed.is_tool_call_response, "non-whitespace suffix falls back to text"); - failures += check(parsed.content == text, "suffix fallback preserves text"); + for (const bool tolerant : {false, true}) { + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response, "trailing suffix retains completed tool call"); + failures += check(parsed.tool_calls.size() == 1, "one call recovered despite trailing suffix"); + if (parsed.tool_calls.size() == 1) { + failures += check(parsed.tool_calls[0].name == "get_weather", "correct tool name"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("city") == "Paris", "correct parameter value"); + } + } return failures; } int test_configured_name_limit() { const std::string name(128, 'a'); - const std::string text = "\n\n\n"; + const std::string text = " \n\n\n"; const ninfer::serve::ParsedToolCallOutput anthropic = - ninfer::serve::parse_qwen_tool_call_output(text, 128); + ninfer::serve::parse_qwen_tool_call_output(text, 128, {}); const ninfer::serve::ParsedToolCallOutput openai = - ninfer::serve::parse_qwen_tool_call_output(text, 64); + ninfer::serve::parse_qwen_tool_call_output(text, 64, {}); const std::string too_long_text = - "\n\n\n"; + " \n\n\n"; const ninfer::serve::ParsedToolCallOutput too_long = - ninfer::serve::parse_qwen_tool_call_output(too_long_text, 128); + ninfer::serve::parse_qwen_tool_call_output(too_long_text, 128, {}); int failures = 0; failures += check(anthropic.is_tool_call_response && anthropic.tool_calls.size() == 1 && @@ -125,9 +203,17 @@ int test_incremental_filter_valid_tool() { visible += filter.finish(true); int failures = 0; failures += check(visible == "Calling weather.", - "valid tool filter did not stream the trimmed content prefix"); + "split-marker stream may emit prefix before is recognized"); failures += check(filter.emitted_bytes() == visible.size(), "valid tool filter byte count mismatch"); + + ninfer::serve::ToolCallStreamFilter oneshot; + const std::string full = "Calling weather. \n\n\n" + "\n"; + std::string held; + held += oneshot.feed(full); + held += oneshot.finish(true); + failures += check(held.empty(), "complete tool payload in one feed emits no preamble"); return failures; } @@ -151,17 +237,915 @@ int test_incremental_filter_fallback() { return failures; } +int test_tolerant_recovery() { + int failures = 0; + + // Upstream #10 test: duplicate closing tags and extra suffix after complete function + const std::string drifted = "Thought before the call.\n" + "\n" + "\n" + "\n" + "/home/matt/Projects/gamemanager/src-tauri/src/main.rs\n" + "\n" + "\n15\n\n" + "\n15\n\n" + "\n" + "\n" + "\n" + "\n" + "extra suffix"; + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(drifted, 64, true); + failures += check(parsed.is_tool_call_response, "tolerant parser recovered drifted call"); + failures += check(!parsed.content.empty(), + "tolerant parser retains the preamble for caller-side suppression"); + failures += check(parsed.tool_calls.size() == 1, "tolerant parser recovered one call"); + failures += check(parsed.tool_calls[0].name == "read", "tolerant parser recovered function"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("filePath") == "/home/matt/Projects/gamemanager/src-tauri/src/main.rs", + "tolerant parser recovered filePath"); + failures += check(args.at("limit") == "15", "tolerant parser recovered limit as raw text"); + failures += check(args.at("offset") == "15", "tolerant parser recovered offset as raw text"); + + // Missing outer + const std::string missing_outer = "\n" + "\n" + "\ntrue\n\n" + ""; + const auto recovered_missing_outer = + ninfer::serve::parse_qwen_tool_call_output(missing_outer, 64, true); + failures += check(recovered_missing_outer.is_tool_call_response && + recovered_missing_outer.tool_calls.size() == 1, + "tolerant parser recovered missing outer close"); + + const auto strict_missing_outer = + ninfer::serve::parse_qwen_tool_call_output(missing_outer, 64, false); + failures += check(!strict_missing_outer.is_tool_call_response, + "strict parser rejected missing outer close"); + + // Negative tests: incomplete/truncated parameters or functions must NOT be recovered + const std::string truncated_param = "\n" + "\n" + "\nhttps://example.com/api"; + const auto truncated_parsed = + ninfer::serve::parse_qwen_tool_call_output(truncated_param, 64, true); + failures += check(!truncated_parsed.is_tool_call_response, + "tolerant parser rejected truncated parameter (not executed)"); + + // Negative tests: near-miss tags must NOT be recovered into fabricated calls + const std::string near_miss_fn = "\n" + "\n" + "\nls -la\n\n" + "\n" + ""; + const auto near_miss_parsed = + ninfer::serve::parse_qwen_tool_call_output(near_miss_fn, 64, true); + failures += check(!near_miss_parsed.is_tool_call_response, + "tolerant parser rejected near-miss function name= tag"); + + // Negative tests: bare function without must NOT be recovered + const std::string bare_fn = "\n\n1\n\n"; + const auto bare_parsed = ninfer::serve::parse_qwen_tool_call_output(bare_fn, 64, true); + failures += check(!bare_parsed.is_tool_call_response, + "tolerant parser rejected bare function tag"); + + // Negative tests: schema/echoed tags (, ) must NOT fabricate calls + const std::string schema_echo = + "\n\nfoo\n\n"; + const auto schema_parsed = ninfer::serve::parse_qwen_tool_call_output(schema_echo, 64, true); + failures += check(!schema_parsed.is_tool_call_response, + "tolerant parser rejected schema echo tags"); + + return failures; +} + +int test_pass_through_adversarial_values() { + int failures = 0; + + // A valid tool call whose parameter value contains XML fragments and tag-like strings + const std::string adversarial = + "\n" + "\n" + "\n" + "value\n" + "\n" + "\n" + ""; + + const auto strict = ninfer::serve::parse_qwen_tool_call_output(adversarial, 64, false); + const auto tolerant = ninfer::serve::parse_qwen_tool_call_output(adversarial, 64, true); + + failures += check(strict.is_tool_call_response, "strict mode recognized adversarial value"); + failures += check(tolerant.is_tool_call_response, "tolerant mode recognized adversarial value"); + if (!strict.is_tool_call_response || !tolerant.is_tool_call_response || + strict.tool_calls.empty() || tolerant.tool_calls.empty()) { + return failures; + } + failures += check(strict.tool_calls.size() == 1 && tolerant.tool_calls.size() == 1, + "both parsed 1 call"); + failures += check(strict.tool_calls[0].name == "process_xml" && + tolerant.tool_calls[0].name == "process_xml", + "both parsed exact name"); + failures += check(strict.tool_calls[0].arguments_json == tolerant.tool_calls[0].arguments_json, + "strict and tolerant produced byte-identical argument JSON"); + const Json args = Json::parse(strict.tool_calls[0].arguments_json); + failures += check( + args.at("payload") == + "value", + "parameter value preserved exactly without premature truncation"); + + return failures; +} + +int test_streaming_consistency() { + int failures = 0; + + // Verify stream filter emission matches parsed tool call content prefix + const std::string response = "I will check that for you.\n" + "\n" + "\n" + "\ntest\n\n" + "\n" + "\n" + "\n"; + + ninfer::serve::ToolCallStreamFilter filter; + std::string streamed; + streamed += filter.feed(response.substr(0, 15)); + streamed += filter.feed(response.substr(15)); + streamed += filter.finish(true); + + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(response, 64, true); + failures += check(parsed.is_tool_call_response, "parsed as tool response"); + failures += check(parsed.content == "I will check that for you.", + "parser retains the pre-call preamble"); + failures += check(filter.emitted_bytes() <= parsed.content.size(), + "terminal content is never shorter than what streaming emitted"); + failures += check(streamed == "I will check th", + "prefix streamed before the tool marker is recognized cannot be recalled"); + + return failures; +} + +int test_multi_tool_discrimination_and_parallel() { + int failures = 0; + + // Parallel calls with trailing suffix after the last call + const std::string text = "\n" + "\n" + "\nTokyo\n\n" + "\n" + "\n" + "\n" + "\n" + "\nTokyo\n\n" + "\n" + "\n" + "\n" + "Done!"; + + const auto tolerant = ninfer::serve::parse_qwen_tool_call_output(text, 64, true); + failures += check(tolerant.is_tool_call_response, "tolerant parsed parallel calls"); + failures += check(tolerant.tool_calls.size() == 2, "2 calls recovered"); + failures += check(tolerant.tool_calls[0].name == "get_temperature", "first name"); + failures += check(tolerant.tool_calls[1].name == "get_wind", "second name"); + + const Json arg0 = Json::parse(tolerant.tool_calls[0].arguments_json); + const Json arg1 = Json::parse(tolerant.tool_calls[1].arguments_json); + failures += check(arg0.at("location") == "Tokyo", "first arg"); + failures += check(arg1.at("location") == "Tokyo", "second arg"); + return failures; +} + +// Schema-driven coverage: construct real ToolDefinition schemas and exercise +// build_tool_param_type_map end-to-end instead of hand-fabricating the map. +int test_schema_driven_type_map() { + // taskId is string; days is integer; count is nullable integer; note has a + // misspelled "strnig" type; flag is boolean; payload is object. + const ninfer::serve::ToolDefinition tool = make_tool( + "TaskUpdate", + R"({"type":"object","properties":{)" + R"("taskId":{"type":"string"},)" + R"("days":{"type":"integer"},)" + R"("count":{"type":["integer","null"]},)" + R"("note":{"type":"strnig"},)" + R"("flag":{"type":"boolean"},)" + R"("payload":{"type":"object"})" + R"(}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + int failures = 0; + failures += check(map.count("TaskUpdate") == 1, "tool recorded"); + const auto& inner = map.at("TaskUpdate"); + failures += check(inner.count("taskId") == 0, "string-typed taskId not recorded"); + failures += check(inner.count("days") == 1, "integer-typed days recorded"); + failures += check(inner.count("count") == 1, "nullable integer count recorded"); + failures += check(inner.count("note") == 0, "misspelled strnig type not recorded"); + failures += check(inner.count("flag") == 1, "boolean-typed flag recorded"); + failures += check(inner.count("payload") == 1, "object-typed payload recorded"); + return failures; +} + +// (a) numeric-looking string param (taskId=1 -> "1" string). +int test_schema_string_param_keeps_numeric_looking_value() { + const ninfer::serve::ToolDefinition tool = make_tool( + "TaskUpdate", R"({"type":"object","properties":{"taskId":{"type":"string"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "1\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "schema string call parsed as tool response"); + failures += check(parsed.tool_calls.size() == 1, "one schema string call"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("taskId").is_string(), "schema taskId is a string, not a number"); + failures += check(args.at("taskId") == "1", "schema string taskId keeps numeric-looking value"); + return failures; +} + +// (b) genuine integer (days=2 -> 2 number). +int test_schema_integer_param_deserializes() { + const ninfer::serve::ToolDefinition tool = make_tool( + "get_weather", R"({"type":"object","properties":{"days":{"type":"integer"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "2\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "schema integer call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("days").is_number(), "schema integer days is a number"); + failures += check(args.at("days") == 2, "schema integer days deserialized to number 2"); + return failures; +} + +// (c) valid nullable integer (["integer","null"] count=7 -> 7 number). +int test_schema_nullable_integer_deserializes() { + const ninfer::serve::ToolDefinition tool = make_tool( + "get_items", R"({"type":"object","properties":{"count":{"type":["integer","null"]}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "7\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "nullable integer call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("count").is_number(), "nullable count deserialized to number"); + failures += check(args.at("count") == 7, "nullable count value 7 preserved as number"); + return failures; +} + +// ["string","null"] => string allowed => not recorded; 5 -> "5". +int test_schema_nullable_string_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "get_opt", R"({"type":"object","properties":{"opt":{"type":["string","null"]}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "5\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "nullable string call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("opt").is_string(), "nullable string opt stays a string"); + failures += check(args.at("opt") == "5", "nullable string opt value preserved as text"); + return failures; +} + +// ["integer","string"] => string allowed => not recorded; 9 -> "9". +int test_schema_mixed_integer_string_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "mix", R"({"type":"object","properties":{"v":{"type":["integer","string"]}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "9\n" + "\n" + "", + 64, map); + + int failures = 0; + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("v").is_string(), "mixed integer/string v stays a string"); + failures += check(args.at("v") == "9", "mixed integer/string v value preserved as text"); + return failures; +} + +// (d) invalid type spelling ("strnig" -> raw text). +int test_schema_invalid_type_spelling_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "bad", R"({"type":"object","properties":{"note":{"type":"strnig"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "hi\n" + "\n" + "", + 64, map); + + int failures = 0; + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("note").is_string(), "invalid-type note stays a string"); + failures += check(args.at("note") == "hi", "invalid-type note value preserved as text"); + return failures; +} + +// (d) boolean param: Python-style scalars coerce to JSON booleans +// (vLLM qwen3coder coercion); non-boolean text stays raw. +int test_schema_boolean_param_coerces_python_scalars() { + const ninfer::serve::ToolDefinition tool = make_tool( + "set_flags", + R"({"type":"object","properties":{)" + R"("a":{"type":"boolean"},"b":{"type":"boolean"},"c":{"type":"boolean"},)" + R"("d":{"type":"boolean"},"e":{"type":"boolean"},"f":{"type":"boolean"},)" + R"("g":{"type":"boolean"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "True\n" + "1\n" + "False\n" + "0\n" + "maybe\n" + "true\n" + " TRUE \n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "schema boolean call parsed as tool response"); + failures += check(parsed.tool_calls.size() == 1, "one schema boolean call"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("a").is_boolean() && args.at("a") == true, + "boolean param True coerces to true"); + failures += check(args.at("b").is_boolean() && args.at("b") == true, + "boolean param 1 coerces to true"); + failures += check(args.at("c").is_boolean() && args.at("c") == false, + "boolean param False coerces to false"); + failures += check(args.at("d").is_boolean() && args.at("d") == false, + "boolean param 0 coerces to false"); + failures += check(args.at("e").is_string() && args.at("e") == "maybe", + "non-boolean text for a boolean param stays raw"); + failures += check(args.at("f").is_boolean() && args.at("f") == true, + "JSON true for a boolean param still coerces"); + failures += check(args.at("g").is_boolean() && args.at("g") == true, + "padded all-caps TRUE coerces to true"); + return failures; +} + +// (g) nullable boolean: Python scalars coerce, the literal null is JSON +// null, and the result does not depend on the type-array order. +int test_schema_nullable_boolean_param() { + const ninfer::serve::ToolDefinition tool = make_tool( + "flags", + R"({"type":"object","properties":{)" + R"("a":{"type":["boolean","null"]},"b":{"type":["null","boolean"]},)" + R"("c":{"type":"boolean"},"d":{"type":["boolean","null"]},"e":{"type":["null","boolean"]},"f":{"type":["boolean","null"]}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "True\n" + "null\n" + "null\n" + "maybe\n" + "False\n" + "Null\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.is_tool_call_response, "nullable boolean call parsed as tool response"); + failures += check(parsed.tool_calls.size() == 1, "one nullable boolean call"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("a").is_boolean() && args.at("a") == true, + "nullable boolean True coerces to true"); + failures += check(args.at("b").is_null(), + "nullable boolean null is JSON null (null listed first)"); + failures += check(args.at("c").is_null(), + "plain boolean null is JSON null"); + failures += check(args.at("d").is_string() && args.at("d") == "maybe", + "non-boolean text for a nullable boolean stays raw"); + failures += check(args.at("e").is_boolean() && args.at("e") == false, + "nullable boolean False coerces to false (null listed first)"); + failures += check(args.at("f").is_null(), + "capitalized Null is JSON null (case-insensitive)"); + return failures; +} + +// (e) boolean true for a string param -> raw text "true". +// (f) null for a string param -> raw text "null". +int test_schema_string_param_bool_and_null_preserve_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "flaggy", R"({"type":"object","properties":{"s":{"type":"string"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "true\n" + "\n" + "\n" + "\n" + "\n" + "null\n" + "\n" + "", + 64, map); + + int failures = 0; + failures += check(parsed.tool_calls.size() == 2, "two string-param calls parsed"); + const Json a1 = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(a1.at("s").is_string(), "string param bool stays a string"); + failures += check(a1.at("s") == "true", "string param bool value preserved as text"); + const Json a2 = Json::parse(parsed.tool_calls[1].arguments_json); + failures += check(a2.at("s").is_string(), "string param null stays a string"); + failures += check(a2.at("s") == "null", "string param null value preserved as text"); + return failures; +} + +// (g) object-looking text for a string param -> raw text. +int test_schema_string_param_object_text_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "obj", R"({"type":"object","properties":{"s":{"type":"string"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "{\"k\":1}\n" + "\n" + "", + 64, map); + + int failures = 0; + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("s").is_string(), "string param object text stays a string"); + failures += check(args.at("s") == "{\"k\":1}", "string param object text preserved verbatim"); + return failures; +} + +// (h) empty type array ("type":[] -> raw text, no crash). +int test_schema_empty_type_array_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "emptytype", R"({"type":"object","properties":{"n":{"type":[]}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + int failures = 0; + failures += check(map.at("emptytype").count("n") == 0, + "empty type array leaves n unrecorded"); + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "3\n" + "\n" + "", + 64, map); + failures += check(parsed.is_tool_call_response, "empty-array call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("n").is_string(), "empty type array n stays a string"); + failures += check(args.at("n") == "3", "empty type array n value preserved as text"); + return failures; +} + +// (i) non-string non-array "type" (e.g. "type":5) preserves raw text. +int test_schema_non_string_non_array_type_preserves_raw() { + const ninfer::serve::ToolDefinition tool = make_tool( + "numtype", R"({"type":"object","properties":{"n":{"type":5}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({tool}); + + int failures = 0; + failures += check(map.at("numtype").count("n") == 0, + "non-string non-array type leaves n unrecorded"); + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "3\n" + "\n" + "", + 64, map); + failures += check(parsed.is_tool_call_response, "non-string-type call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("n").is_string(), "non-string-type n stays a string"); + failures += check(args.at("n") == "3", "non-string-type n value preserved as text"); + return failures; +} + +// A second same-name definition whose object schema has no "properties" +// must replace the first definition's recorded (integer) permissions, +// leaving the param unrecorded (raw text) instead of leaking the first. +int test_duplicate_tool_definition_no_properties_replaces() { + const ninfer::serve::ToolDefinition first = make_tool( + "dup2", R"({"type":"object","properties":{"count":{"type":"integer"}}})"); + const ninfer::serve::ToolDefinition second = make_tool( + "dup2", R"({"type":"object"})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({first, second}); + + int failures = 0; + failures += check(map.count("dup2") == 1, "no-properties duplicate has one entry"); + failures += check(map.at("dup2").count("count") == 0, + "second no-properties definition replaced the first (count not recorded)"); + const ninfer::serve::ParsedToolCallOutput parsed = + ninfer::serve::parse_qwen_tool_call_output( + " \n" + "\n" + "3\n" + "\n" + "", + 64, map); + failures += check(parsed.is_tool_call_response, "no-properties dup call parsed"); + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("count").is_string(), "no-properties dup count stays a string"); + failures += check(args.at("count") == "3", "no-properties dup count value preserved as text"); + return failures; +} + +// A redefinition of the same tool name must replace the prior entry so a +// second (string) definition cannot leak the first's integer permission. +int test_duplicate_tool_definition_replaced() { + const ninfer::serve::ToolDefinition first = make_tool( + "dup", R"({"type":"object","properties":{"count":{"type":"integer"}}})"); + const ninfer::serve::ToolDefinition second = make_tool( + "dup", R"({"type":"object","properties":{"count":{"type":"string"}}})"); + const ninfer::serve::ToolParamTypeMap map = + ninfer::serve::build_tool_param_type_map({first, second}); + + int failures = 0; + failures += check(map.count("dup") == 1, "duplicate tool name has one entry"); + failures += check(map.at("dup").count("count") == 0, + "second (string) definition replaced the first (integer) entry"); + return failures; +} + +int test_json_argument_object_both_modes() { + const std::string text = "\n" + "\n" + "{\"query\":\"scheduling\"}\n" + "\n" + ""; + int failures = 0; + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response, "JSON-args form recovered"); + failures += check(parsed.content.empty(), "JSON-args form has no visible content"); + failures += check(parsed.tool_calls.size() == 1 && parsed.tool_calls[0].name == "search", + "JSON-args form recovered search"); + if (parsed.tool_calls.empty()) { continue; } + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("query") == "scheduling", "JSON-args query preserved"); + } + return failures; +} + +int test_json_argument_object_schema_typing() { + const ninfer::serve::ToolDefinition tool = make_tool( + "write_file", + R"({"type":"object","properties":{)" + R"("path":{"type":"string"},)" + R"("content":{"type":"string"},)" + R"("overwrite":{"type":"boolean"}}})"); + const auto map = ninfer::serve::build_tool_param_type_map({tool}); + const std::string text = + "\n\n" + "{\"path\":\"config.json\",\"content\":{\"a\":1,\"b\":true,\"name\":\"svc\"}," + "\"overwrite\":true}\n" + "\n"; + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, map, false); + int failures = 0; + failures += check(parsed.is_tool_call_response && parsed.tool_calls.size() == 1, + "JSON-args write_file recovered"); + if (parsed.tool_calls.empty()) { return failures; } + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("path").is_string() && args.at("path") == "config.json", + "JSON-args path stays string"); + failures += check(args.at("content").is_string(), + "JSON-args content declared string stays string, not object"); + failures += check(args.at("overwrite").is_boolean() && args.at("overwrite") == true, + "JSON-args overwrite stays boolean"); + return failures; +} + +int test_json_argument_boolean_python_scalar() { + const ninfer::serve::ToolDefinition tool = make_tool( + "write_file", + R"({"type":"object","properties":{"overwrite":{"type":"boolean"}}})"); + const auto map = ninfer::serve::build_tool_param_type_map({tool}); + const std::string text = + "\n\n{\"overwrite\":\"True\"}\n\n"; + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, map); + int failures = 0; + failures += check(parsed.is_tool_call_response, "JSON-args True recovered"); + if (!parsed.tool_calls.empty()) { + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("overwrite").is_boolean() && args.at("overwrite") == true, + "JSON-args Python True coerced to boolean"); + } + return failures; +} + +int test_near_miss_forms_do_not_fabricate() { + int failures = 0; + const std::vector texts = { + "[tool_use: search]{\"query\":\"scheduling\"}", + "tool_call: search\narguments: {\"query\":\"scheduling\"}", + "You can emit to look things up.", + "use {\"query\":\"x\"} in your reply", + "\n{\"query\":\"x\"}\n", + }; + for (const auto& text : texts) { + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(!parsed.is_tool_call_response && parsed.tool_calls.empty() && + parsed.content == text, + "near-miss form fabricated a call"); + } + } + return failures; +} + +int test_json_args_adversarial_string_roundtrip() { + const std::string payload = "x raw"; + const std::string text = + "\n\n{\"payload\":\"x raw\"}\n" + "\n"; + int failures = 0; + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response && parsed.tool_calls.size() == 1, + "JSON-args adversarial recovered"); + if (parsed.tool_calls.empty()) { continue; } + const Json args = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(args.at("payload") == payload, "JSON-args adversarial value round-trip"); + } + return failures; +} + } // namespace + +// Regression: a preamble followed by a tool call must never leave the terminal +// body shorter than what streaming already emitted. When the marker +// straddles two chunks the prefix is already on the wire, and a parser that +// dropped it would trip unstreamed_content and abort the request mid-stream. +int test_stream_terminal_consistency() { + int failures = 0; + const std::string preamble = "I need to search the knowledge base. "; + const std::string call = + "\n\n\nx\n\n\n" + ""; + + // split marker: prefix is emitted before the call is recognised + { + ninfer::serve::ToolCallStreamFilter filter; + std::string streamed; + streamed += filter.feed(preamble + "\n" + "styles.css\n" + "--temper-sans: Inter\n" + "--font: Arial\n" + "\n" + "\n" + "\n" + "\n" + "index.html\n" + "\n" + "\n" + "\n" + "\n" + "I have finished updating both files."; + + int failures = 0; + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response, "two calls with trailing prose parsed as tool response"); + failures += check(parsed.tool_calls.size() == 2, "two calls recovered despite trailing prose"); + if (parsed.tool_calls.size() == 2) { + failures += check(parsed.tool_calls[0].name == "replace_string_in_file", "first call name"); + failures += check(parsed.tool_calls[1].name == "replace_string_in_file", "second call name"); + const Json arg0 = Json::parse(parsed.tool_calls[0].arguments_json); + const Json arg1 = Json::parse(parsed.tool_calls[1].arguments_json); + failures += check(arg0.at("filePath") == "styles.css", "first call filePath"); + failures += check(arg1.at("filePath") == "index.html", "second call filePath"); + } + } + return failures; +} + +int test_partial_recovery_text_between_calls() { + const std::string text = "\n" + "\n" + "Tokyo\n" + "\n" + "\n" + "Next, checking the second city:\n" + "\n" + "\n" + "Kyoto\n" + "\n" + ""; + + int failures = 0; + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response, "calls separated by text parsed as tool response"); + failures += check(parsed.tool_calls.size() == 2, "both calls recovered across intermediate text"); + if (parsed.tool_calls.size() == 2) { + const Json arg0 = Json::parse(parsed.tool_calls[0].arguments_json); + const Json arg1 = Json::parse(parsed.tool_calls[1].arguments_json); + failures += check(arg0.at("city") == "Tokyo", "first city Tokyo"); + failures += check(arg1.at("city") == "Kyoto", "second city Kyoto"); + } + } + return failures; +} + +int test_partial_recovery_truncated_tail() { + const std::string text = "\n" + "\n" + "styles.css\n" + "--temper-sans: Inter\n" + "--font: Arial\n" + "\n" + "\n" + "\n" + "\n" + "index.html\n" + ""; + + int failures = 0; + for (const bool tolerant : {false, true}) { + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, tolerant); + failures += check(parsed.is_tool_call_response, "truncated tail retains complete prior call"); + failures += check(parsed.tool_calls.size() == 1, "exactly one valid call recovered"); + if (parsed.tool_calls.size() == 1) { + failures += check(parsed.tool_calls[0].name == "replace_string_in_file", "valid call name preserved"); + const Json arg0 = Json::parse(parsed.tool_calls[0].arguments_json); + failures += check(arg0.at("filePath") == "styles.css", "valid call filePath preserved"); + } + } + return failures; +} + +int test_partial_recovery_fuzz_split_boundaries() { + const std::string text = "\n" + "\n" + "123\n" + "\n" + "\n" + " intermediate prose \n" + "\n" + "\n" + "456\n" + "\n" + "\n" + " trailing prose"; + + const auto parsed = ninfer::serve::parse_qwen_tool_call_output(text, 64, {}, false); + int failures = 0; + failures += check(parsed.is_tool_call_response, "fuzz test input parsed as tool response"); + failures += check(parsed.tool_calls.size() == 2, "fuzz test input recovered 2 calls"); + + for (std::size_t split = 0; split <= text.size(); ++split) { + ninfer::serve::ToolCallStreamFilter filter; + std::string s1 = filter.feed(text.substr(0, split)); + std::string s2 = filter.feed(text.substr(split)); + std::string s3 = filter.finish(true); + std::string total = s1 + s2 + s3; + failures += check(total.empty(), "stream filter leaked during partial-recovery split"); + } + + for (std::size_t chunk_sz : {1, 2, 3, 5, 8, 13, 21}) { + ninfer::serve::ToolCallStreamFilter filter; + std::string total; + for (std::size_t i = 0; i < text.size(); i += chunk_sz) { + total += filter.feed(text.substr(i, chunk_sz)); + } + total += filter.finish(true); + failures += check(total.empty(), "stream filter leaked during chunked streaming"); + } + + return failures; +} + int main() { int failures = 0; + failures += test_stream_terminal_consistency(); failures += test_single_call(); failures += test_multiple_calls_and_json_values(); + failures += test_string_param_keeps_numeric_looking_value(); + failures += test_unknown_param_defaults_to_string(); failures += test_malformed_falls_back_to_text(); - failures += test_suffix_after_tool_falls_back_to_text(); + failures += test_partial_recovery_trailing_suffix(); + failures += test_partial_recovery_multiple_calls_with_trailing_prose(); + failures += test_partial_recovery_text_between_calls(); + failures += test_partial_recovery_truncated_tail(); + failures += test_partial_recovery_fuzz_split_boundaries(); failures += test_configured_name_limit(); failures += test_incremental_filter_valid_tool(); failures += test_incremental_filter_fallback(); + failures += test_tolerant_recovery(); + failures += test_pass_through_adversarial_values(); + failures += test_streaming_consistency(); + failures += test_multi_tool_discrimination_and_parallel(); + failures += test_schema_driven_type_map(); + failures += test_schema_string_param_keeps_numeric_looking_value(); + failures += test_schema_integer_param_deserializes(); + failures += test_schema_nullable_integer_deserializes(); + failures += test_schema_nullable_string_preserves_raw(); + failures += test_schema_boolean_param_coerces_python_scalars(); + failures += test_schema_nullable_boolean_param(); + failures += test_schema_mixed_integer_string_preserves_raw(); + failures += test_schema_invalid_type_spelling_preserves_raw(); + failures += test_schema_string_param_bool_and_null_preserve_raw(); + failures += test_schema_string_param_object_text_preserves_raw(); + failures += test_schema_empty_type_array_preserves_raw(); + failures += test_schema_non_string_non_array_type_preserves_raw(); + failures += test_duplicate_tool_definition_no_properties_replaces(); + failures += test_duplicate_tool_definition_replaced(); + failures += test_json_argument_object_both_modes(); + failures += test_json_argument_object_schema_typing(); + failures += test_json_argument_boolean_python_scalar(); + failures += test_near_miss_forms_do_not_fabricate(); + failures += test_json_args_adversarial_string_roundtrip(); if (failures == 0) { std::cout << "ok\n"; } return failures == 0 ? 0 : 1; } diff --git a/tools/bench/run_serve_corpus.py b/tools/bench/run_serve_corpus.py index cec025e22f..209de194a1 100644 --- a/tools/bench/run_serve_corpus.py +++ b/tools/bench/run_serve_corpus.py @@ -83,7 +83,7 @@ RUN_ARTIFACT_TYPE = "ninfer_serve_corpus_result" RUN_SCHEMA_VERSION = 5 SERVER_LOG_ARTIFACT_TYPE = "ninfer_serve_request_log" -SERVER_LOG_SCHEMA_VERSION = 9 +SERVER_LOG_SCHEMA_VERSION = 10 STARTUP_TIMEOUT_SECONDS = 1800.0 REQUEST_TIMEOUT_SECONDS = 24.0 * 60.0 * 60.0 LOG_EVENT_TIMEOUT_SECONDS = 10.0