diff --git a/docs/maintainer/engine-architecture.md b/docs/maintainer/engine-architecture.md index 0631cda163..e2c146d68b 100644 --- a/docs/maintainer/engine-architecture.md +++ b/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 diff --git a/src/models/qwen3_5/program/structured_round.cpp b/src/models/qwen3_5/program/structured_round.cpp index 06b7d41003..af83900651 100644 --- a/src/models/qwen3_5/program/structured_round.cpp +++ b/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(drafts.ne[0]) >= width_ || drafts.ne[1] <= 0 || + static_cast(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)); diff --git a/tests/models/qwen3_5/test_structured_round.cpp b/tests/models/qwen3_5/test_structured_round.cpp index 2b0c613111..c20ad77365 100644 --- a/tests/models/qwen3_5/test_structured_round.cpp +++ b/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 vocab(vocab_size); for (int i = 0; i < 128; ++i) { vocab[i] = std::string(1, static_cast(i)); } text::StructuredCompiler compiler(vocab, {128}); @@ -39,6 +49,12 @@ int main() { round.check(); const auto bits = from_device(round.device_mask(7 - pass), words * width); + std::vector 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="; @@ -46,6 +62,38 @@ int main() { 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 row_tokens(draft_width, ' '); + const std::string prefix = "{\"a\":1}"; + std::copy(prefix.begin(), prefix.end(), row_tokens.begin()); + std::vector 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(round.device_mask(5 + pass), words * width); + std::vector 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(vocab_size, "q"); dead_vocab.back().clear();