fix: preserve structured draft row pitch across speculative widths

docs/maintainer/engine-architecture.md

@@ -645,6 +645,11 @@ transfer, then target verification. The host function calls no CUDA API; it capt
 for rethrow after stream synchronization and before any token can be published. All node addresses
 belong to Program and outlive graph replay.
 
+Draft frames are compact at each round's actual proposal width, which may be narrower than the
+maximum ngram/copy window. D2H staging must preserve the fixed host-row pitch used by the grammar
+callback; a contiguous whole-frame copy cannot substitute for row-strided transfer. Graph replay
+must retain this mapping for every compact row, independently of its physical lane.
+
 SamplingConfig carries an optional bitset pointer and column stride. Column i describes the
 grammar after drafts[0..i). The mask is applied before penalties and sampling filters. After an
 illegal draft or EOS, suffix columns are unreachable and need no grammar traversal. Every reachable

src/models/qwen3_5/program/structured_round.cpp

@@ -52,10 +52,19 @@ void CUDART_CB StructuredRound::callback(void* opaque) noexcept {
 }
 
 void StructuredRound::enqueue_dflash(const Tensor& drafts, cudaStream_t stream) {
+    if (drafts.dtype != DType::I32 || drafts.ne[0] <= 0 ||
+        static_cast<std::uint32_t>(drafts.ne[0]) >= width_ || drafts.ne[1] <= 0 ||
+        static_cast<std::uint32_t>(drafts.ne[1]) > lanes_ || drafts.ne[2] != 1 ||
+        drafts.ne[3] != 1 || drafts.nb[0] != sizeof(TokenId)) {
+        throw std::logic_error("structured draft tensor exceeds its reserved frame");
+    }
     // No CUDA calls or exceptions escape the host function. The transfers and host node are
     // captured together: verification cannot observe masks until the CPU matcher has filled them.
-    CUDA_CHECK(cudaMemcpyAsync(host_drafts_.data(), drafts.data, drafts.bytes(),
-                               cudaMemcpyDeviceToHost, stream));
+    // The device frame is compact at this round's width (neural/copy widths can differ).
+    // The callback addresses fixed maximum-width host rows; preserve that pitch on every replay.
+    CUDA_CHECK(cudaMemcpy2DAsync(host_drafts_.data(), (width_ - 1U) * sizeof(TokenId),
+                                 drafts.data, drafts.nb[1], drafts.ne[0] * sizeof(TokenId),
+                                 drafts.ne[1], cudaMemcpyDeviceToHost, stream));
     CUDA_CHECK(cudaLaunchHostFunc(stream, callback, this));
     CUDA_CHECK(cudaMemcpyAsync(masks_.data, host_masks_.data(), masks_.bytes(),
                                cudaMemcpyHostToDevice, stream));

tests/models/qwen3_5/test_structured_round.cpp

@@ -9,12 +9,22 @@ using namespace ninfer::test;
 int main() {
     if (cuda_unavailable()) { return 77; }
     try {
-        constexpr int vocab_size = 129, words = 5, width = 2, lanes = 8;
+        // Reserve a wider copy window than the actual one-token neural round.
+        // Compact draft rows must not be read using the maximum reserved pitch.
+        constexpr int vocab_size = 129, words = 5, width = 25, lanes = 8;
         DeviceContext device(0);
-        DeviceArena arena(words * width * lanes * 4 + 256);
+        DeviceArena arena(words * width * lanes * 4 + 2048);
         auto masks  = arena.alloc(DType::I32, {words, width, lanes});
-        auto drafts = arena.alloc(DType::I32, {width - 1, lanes});
+        auto drafts = arena.alloc(DType::I32, {1, lanes});
+        auto wide_drafts = arena.alloc(DType::I32, {width - 1, lanes});
         models::qwen3_5::StructuredRound round(masks, vocab_size, width, lanes);
+        // Seed the maximum-width host frame deterministically. A compact copy must replace
+        // each row, rather than leaving zero draft tokens in later reserved rows.
+        CUDA_CHECK(cudaMemsetAsync(wide_drafts.data, 0, wide_drafts.bytes(), device.stream));
+        round.begin_dflash();
+        round.enqueue_dflash(wide_drafts, device.stream);
+        device.synchronize();
+        round.check();
         std::vector<std::string> vocab(vocab_size);
         for (int i = 0; i < 128; ++i) { vocab[i] = std::string(1, static_cast<char>(i)); }
         text::StructuredCompiler compiler(vocab, {128});
@@ -39,13 +49,51 @@ int main() {
             round.check();
             const auto bits =
                 from_device<std::uint32_t>(round.device_mask(7 - pass), words * width);
+            std::vector<std::uint32_t> expected(words * 2);
+            const TokenId expected_draft = '{';
+            grammar->fill_masks(expected, {&expected_draft, 1});
+            if (!std::equal(expected.begin(), expected.end(), bits.begin())) {
+                throw std::runtime_error("compact draft row used reserved-width stride");
+            }
             if (!(bits['{' / 32] & (1U << ('{' % 32))) || (bits[4] & 1U) ||
                 (bits[words + 4] & 1U)) {
                 std::cerr << "pass=" << pass << " bits=";
                 for (auto b : bits) { std::cerr << std::hex << b << " "; }
                 throw std::runtime_error("graph mask/EOS mismatch");
             }
         }
+        // Alternate neural and maximum copy widths. Both graphs share the staging buffers,
+        // but must retain their own transfer geometry across replay and lane remapping.
+        for (const int draft_width : {7, width - 1}) {
+            Tensor frame(wide_drafts.data, DType::I32, {draft_width, lanes});
+            std::vector<TokenId> row_tokens(draft_width, ' ');
+            const std::string prefix = "{\"a\":1}";
+            std::copy(prefix.begin(), prefix.end(), row_tokens.begin());
+            std::vector<TokenId> frame_tokens;
+            for (int row = 0; row < lanes; ++row) {
+                frame_tokens.insert(frame_tokens.end(), row_tokens.begin(), row_tokens.end());
+            }
+            CUDA_CHECK(cudaMemcpyAsync(frame.data, frame_tokens.data(), frame.bytes(),
+                                       cudaMemcpyHostToDevice, device.stream));
+            device.synchronize();
+            DecodeGraphDefinition varying_definition;
+            varying_definition.capture(device.stream, [&] { round.enqueue_dflash(frame, device.stream); });
+            DecodeGraphExecutable varying_executable;
+            varying_executable.instantiate(varying_definition);
+            for (int pass = 0; pass < 3; ++pass) {
+                round.begin_dflash();
+                round.set_dflash_row(2 - pass, 5 + pass, draft_width, grammar);
+                varying_executable.launch(device.stream);
+                device.synchronize();
+                round.check();
+                const auto bits = from_device<std::uint32_t>(round.device_mask(5 + pass), words * width);
+                std::vector<std::uint32_t> expected(words * (draft_width + 1));
+                grammar->fill_masks(expected, row_tokens);
+                if (!std::equal(expected.begin(), expected.end(), bits.begin())) {
+                    throw std::runtime_error("variable-width draft graph mask mismatch");
+                }
+            }
+        }
         round.begin_dflash();
         auto dead_vocab = std::vector<std::string>(vocab_size, "q");
         dead_vocab.back().clear();