fix: preserve structured draft row pitch across speculative widths

[wongkiller]

Problem and scope

Related Issue: none; Issues are disabled in this repository. The full bug report and reproducer are included below.

Draft for maintainer confirmation of scope/direction (not yet approved). This is a focused local crash fix, submitted with evidence under the draft-PR policy.

When DFlash2 uses seven drafts but the structured round reserves a wider ngram/copy window, compact device rows and fixed-pitch host rows have different layouts. A contiguous D2H copy leaves later callback rows reading the wrong tokens. With concurrent structured requests, incorrect masks can lead to generated token ... violates the structured output grammar, HTTP 500, and a failed worker returning HTTP 503.

Reproduced on base f118551fb401de073555807a48c50238e180e3b8; the report below includes the artifact, server options and HTTP reproducer. Feature provenance: Neroued#294. No claim is made that upstream master reproduces this fork's integration failure. This is separate from #4 (aborted-prefill MTP continuation checkpoints); that change is not included here.

Implementation

• Replace the contiguous copy with cudaMemcpy2DAsync, mapping the tensor's actual device row pitch into the callback's reserved host pitch.
• Validate the draft frame's dtype, dimensions, scalar stride, and capacity before enqueueing the transfer.
• Extend the existing CUDA graph test with deterministic staging initialization, compact widths 1 and 7 against a reserved width of 24, full-width coverage, lane remapping and repeated replay. Keep callback-error/reset coverage.
• Document the staging-pitch contract in the active architecture reference.

The change stays within Program-owned staging. It does not disable grammar checks, catch-and-ignore worker failures, change quantization, reduce concurrency, or change the public API.

Verification

Environment: RTX 5090 32 GB; Windows / WSL2 / Docker Desktop; Linux CUDA 13.1.2, Release sm_120a compatibility build. The tested model was Ternary Bonsai 2 27B Heretic with DFlash2-7, default ngram proposals, two active lanes, rk4v4-e8 KV and a 917,504-token configured pool. The live regression requests were short; this is not a full-context correctness claim.

Previously executed local verification:

• The same extended CUDA regression compiled against the original implementation exited 1 with compact draft row used reserved-width stride; against the fixed implementation it exited 0.
• Rebuilt ninfer, ninfer-serve, and the focused regression target successfully.
• 60 paired strict-JSON requests, 20 mixed plain/JSON streaming requests, and 10 medium-reasoning JSON streaming requests: 90/90 passed, with no worker crashes in the fixed server logs.
• Contradictory-prompt constant-schema enforcement, a tool-call/result round trip, a 32,726-token prompt, and a Hermes agent smoke test also passed.
• git diff --check passed.

The regression was compiled as a standalone target linked to the existing built runtime so the old and fixed implementations could be tested with the identical test source. Local commands used:

cmake -S /src -B /build -DCMAKE_PROJECT_INCLUDE=/patch/regression.cmake
cmake --build /build --parallel 8 --target structured_round_regression
# Save/run this binary before applying structured_round.cpp, then rebuild:
cmake --build /build --parallel 8 --target ninfer ninfer-serve structured_round_regression
/opt/ninfer-tests/structured-round-before  # expected exit 1
/opt/ninfer-tests/structured-round-after   # exit 0

The standalone target definition was:

add_executable(structured_round_regression
  ${CMAKE_SOURCE_DIR}/tests/models/qwen3_5/test_structured_round.cpp
  ${CMAKE_SOURCE_DIR}/src/models/qwen3_5/program/structured_round.cpp)
target_compile_features(structured_round_regression PRIVATE cxx_std_20)
target_include_directories(structured_round_regression PRIVATE
  ${CMAKE_SOURCE_DIR}/src ${CMAKE_SOURCE_DIR}/include ${CMAKE_SOURCE_DIR}/tests)
target_link_libraries(structured_round_regression PRIVATE ninfer_core ninfer_text)

Maintainers can instead build and run the repository's existing ninfer_qwen3_5_structured_round_test target with the modified test source in this PR.

Limitations: the full repository test suite, CUDA sanitizers, other GPU architectures, and an isolated transfer-performance benchmark were not run. This is a correctness fix; no speedup or performance-parity claim is made. Code and test preparation used AI assistance; the listed results are local test observations, not assumed CI results.

---

Full bug report and HTTP reproducer

Problem

At f118551fb401de073555807a48c50238e180e3b8, concurrent structured-output requests can crash the worker with generated token ... violates the structured output grammar. Active requests return HTTP 500; subsequent requests return HTTP 503 until restart.

This was reproduced on an RTX 5090 (32 GB), Windows / WSL2 / Docker Desktop, using a Linux CUDA 13.1.2 sm_120a compatibility build (NINFER_SM120_NATIVE=OFF) and emiltsoi/Ternary-Bonsai-2-27B-Uncensored-Heretic-NInfer. Artifact SHA-256: 99f94f44bed48892b7c5e1c4dc8349e0db8b4b44d2bd8ad7cd438cfc181de2fd.

Relevant server configuration:

--model /models/Ternary-Bonsai-2-27B-Heretic-ninfer.ninfer
--model-id bonsai2-27b-heretic
--spec dflash2 --draft-tokens 7 --ngram-draft-tokens 15
--structured-output --max-concurrency 2
--kv-dtype rk4v4-e8 --gdn-state-fp16
--max-context 917504 --kv-capacity 917504
--prefill-chunk 2048 --host-kv-mib 8192 --host-state-slots 8
--preserve-thinking --default-thinking-budget 4096 --default-max-tokens 16384
--vision --vision-residency overlay --vision-max-merged 12288

The live failure occurred on the sixth pair of concurrent strict-JSON requests. The following is a minimal HTTP workload matching that test (adjust the URL to the server):

import concurrent.futures, json, urllib.request

def request(i):
    body = {
        "model": "bonsai2-27b-heretic",
        "messages": [{"role": "user", "content": f"Test {i}. Return an object with an items list containing integers 1 to 40 and a label describing them. No commentary."}],
        "max_tokens": 600, "temperature": 0.7, "reasoning_effort": "none",
        "response_format": {"type": "json_schema", "json_schema": {
            "name": "regression", "strict": True, "schema": {
                "type": "object", "properties": {
                    "items": {"type": "array", "items": {"type": "integer"}, "minItems": 40, "maxItems": 40},
                    "label": {"type": "string"}},
                "required": ["items", "label"], "additionalProperties": False}}}}
    req = urllib.request.Request("http://127.0.0.1:18082/v1/chat/completions",
        data=json.dumps(body).encode(), headers={"Content-Type": "application/json"})
    with urllib.request.urlopen(req, timeout=180) as response:
        result = json.load(response)
    choice = result["choices"][0]
    assert choice["finish_reason"] == "stop", choice
    obj = json.loads(choice["message"]["content"])
    assert set(obj) == {"items", "label"} and isinstance(obj["label"], str)
    assert len(obj["items"]) == 40 and all(type(x) is int for x in obj["items"])

for round_id in range(30):
    with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
        list(pool.map(request, [2 * round_id, 2 * round_id + 1]))

The exact failing round/token depends on sampling and batching. A deterministic GPU regression is included in the accompanying proposed fix.

Root cause and proposed scope

StructuredRound::enqueue_dflash() copies the compact device draft frame contiguously. StructuredRound::callback() indexes the host staging rows using the reserved maximum pitch, width_ - 1. With seven neural drafts and a larger ngram reservation, later compact rows are therefore read at the wrong offset. The grammar masks are conditioned on stale/wrong draft tokens.

The proposed change is confined to Program-owned structured-round staging: use cudaMemcpy2DAsync with the actual device row pitch and reserved host pitch, validate frame bounds, and add regression coverage. Grammar enforcement and worker invariant checks remain enabled. No sampling, model, KV-capacity, or public API changes are proposed.

This is a locally reproduced bug in this consolidated fork. The related feature originates in Neroued#294; this report does not claim that upstream master or that PR independently reproduces the same mixed-width integration failure.

A local repair already passed focused tests. Issues are disabled in this repository, so this report accompanies the draft PR under PR_POLICY.md; maintainer confirmation of scope/direction is still pending.