diff --git a/backend/cpp/ds4/generation_limits.h b/backend/cpp/ds4/generation_limits.h new file mode 100644 index 000000000000..0985b059072f --- /dev/null +++ b/backend/cpp/ds4/generation_limits.h @@ -0,0 +1,27 @@ +// SPDX-License-Identifier: MIT +#pragma once + +#include + +namespace ds4cpp { + +inline int EffectiveGenerationLimit(int requested, int context_size, + int session_position) { + const int limit = requested > 0 ? requested : 256; + const int room = context_size - session_position; + if (room <= 1) return 0; + return std::min(limit, room - 1); +} + +inline int RemainingGenerationBudget(int effective_limit, int produced) { + if (effective_limit <= produced) return 0; + return effective_limit - produced; +} + +inline int SpeculativeAcceptedCapacity(int remaining, int draft_allowance, + int buffer_capacity) { + if (remaining <= 0 || draft_allowance < 0 || buffer_capacity <= 0) return 0; + return std::min({remaining, draft_allowance + 1, buffer_capacity}); +} + +} // namespace ds4cpp diff --git a/backend/cpp/ds4/generation_limits_test.cpp b/backend/cpp/ds4/generation_limits_test.cpp new file mode 100644 index 000000000000..ee38251f655e --- /dev/null +++ b/backend/cpp/ds4/generation_limits_test.cpp @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: MIT + +#include "generation_limits.h" + +#include + +namespace { + +int failures = 0; + +void check_equal(int got, int want, const char *name) { + if (got == want) return; + std::fprintf(stderr, "FAIL %s: got %d, want %d\n", name, got, want); + failures++; +} + +// Mutation caught: treating omitted or negative max_tokens as unlimited instead +// of preserving DS4's legacy 256-token default. +void test_nonpositive_uses_legacy_default_when_space_permits() { + check_equal(ds4cpp::EffectiveGenerationLimit(0, 4096, 100), 256, + "zero max_tokens uses legacy default"); + check_equal(ds4cpp::EffectiveGenerationLimit(-1, 4096, 100), 256, + "negative max_tokens uses legacy default"); +} + +// Mutation caught: applying the legacy default without clamping it to the +// post-prefill context room and reserved slot. +void test_legacy_default_is_clamped_by_context() { + check_equal(ds4cpp::EffectiveGenerationLimit(0, 300, 100), 199, + "legacy default is context-clamped"); +} + +// Mutation caught: allowing an explicitly large request to overrun the +// post-prefill context boundary. +void test_large_positive_limit_is_clamped_to_context() { + check_equal(ds4cpp::EffectiveGenerationLimit(32768, 32768, 100), 32667, + "large positive is context-clamped"); +} + +// Mutation caught: replacing every positive request with the legacy default +// rather than preserving a smaller configured limit. +void test_smaller_positive_limit_is_preserved() { + check_equal(ds4cpp::EffectiveGenerationLimit(64, 4096, 100), 64, + "smaller positive is preserved"); +} + +// Mutation caught: consuming the final context slot instead of reserving it as +// required by DS4's generation loop. +void test_no_usable_room_returns_zero() { + check_equal(ds4cpp::EffectiveGenerationLimit(32, 100, 99), 0, + "one remaining context slot is not usable"); +} + +// Mutation caught: sending the original generation limit to a later +// speculative cycle instead of subtracting tokens already produced. +void test_remaining_budget_accounts_for_produced_tokens() { + check_equal(ds4cpp::RemainingGenerationBudget(10, 4), 6, + "remaining budget subtracts produced tokens"); + check_equal(ds4cpp::RemainingGenerationBudget(10, 12), 0, + "remaining budget never becomes negative"); +} + +// Mutation caught: giving speculative evaluation capacity beyond either the +// output budget, the draft allowance plus its first target token, or the fixed +// accepted-token buffer. +void test_speculative_capacity_obeys_all_bounds() { + check_equal(ds4cpp::SpeculativeAcceptedCapacity(3, 8, 8), 3, + "capacity respects remaining output budget"); + check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 4, 8), 5, + "capacity includes one target token beyond draft allowance"); + check_equal(ds4cpp::SpeculativeAcceptedCapacity(20, 8, 6), 6, + "capacity respects fixed buffer"); +} + +} // namespace + +int main() { + test_nonpositive_uses_legacy_default_when_space_permits(); + test_legacy_default_is_clamped_by_context(); + test_large_positive_limit_is_clamped_to_context(); + test_smaller_positive_limit_is_preserved(); + test_no_usable_room_returns_zero(); + test_remaining_budget_accounts_for_produced_tokens(); + test_speculative_capacity_obeys_all_bounds(); + + if (failures == 0) { + std::fprintf(stderr, "all generation limit checks passed\n"); + return 0; + } + std::fprintf(stderr, "%d check(s) failed\n", failures); + return 1; +} diff --git a/backend/cpp/ds4/grpc-server.cpp b/backend/cpp/ds4/grpc-server.cpp index 2118fd1cb270..68ebdd3e3551 100644 --- a/backend/cpp/ds4/grpc-server.cpp +++ b/backend/cpp/ds4/grpc-server.cpp @@ -10,6 +10,7 @@ #include "dsml_parser.h" // populated in Task 12 #include "dsml_renderer.h" // populated in Task 16 +#include "generation_limits.h" #include "kv_cache.h" // populated in Task 17 extern "C" { @@ -769,7 +770,6 @@ class DS4Backend final : public backend::Backend::Service { } ds4_tokens prompt = {}; build_prompt(g_engine, request, &prompt); - int n_predict = request->tokens() > 0 ? request->tokens() : 256; const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request)); const bool starts_in_thinking = think_enabled && @@ -792,6 +792,9 @@ class DS4Backend final : public backend::Backend::Service { int prompt_len = prompt.len; ds4_tokens_free(&prompt); if (rc == 0) { + const int n_predict = ds4cpp::EffectiveGenerationLimit( + request->tokens(), ds4_session_ctx(g_session), + ds4_session_pos(g_session)); const int eos = ds4_token_eos(g_engine); const int draft_max = ds4_engine_mtp_draft_tokens(g_engine); int produced = 0; @@ -810,9 +813,12 @@ class DS4Backend final : public backend::Backend::Service { if (draft_max > 0 && sp.temperature <= 0.0f) { constexpr int kAcceptedMax = 8; int accepted[kAcceptedMax]; - int cap = std::min(kAcceptedMax, draft_max + 1); + const int remaining = ds4cpp::RemainingGenerationBudget( + n_predict, produced); + const int cap = ds4cpp::SpeculativeAcceptedCapacity( + remaining, draft_max, kAcceptedMax); int n = ds4_session_eval_speculative_argmax( - g_session, first, draft_max, eos, + g_session, first, remaining, eos, accepted, cap, err, sizeof(err)); if (n < 0) { rc = -1; break; } bool stop = false; @@ -873,7 +879,6 @@ class DS4Backend final : public backend::Backend::Service { } ds4_tokens prompt = {}; build_prompt(g_engine, request, &prompt); - int n_predict = request->tokens() > 0 ? request->tokens() : 256; const bool think_enabled = ds4_think_mode_enabled(parse_think_mode(request)); const bool starts_in_thinking = think_enabled && @@ -891,6 +896,9 @@ class DS4Backend final : public backend::Backend::Service { int rc = ds4_session_sync(g_session, &prompt, err, sizeof(err)); ds4_tokens_free(&prompt); if (rc == 0) { + const int n_predict = ds4cpp::EffectiveGenerationLimit( + request->tokens(), ds4_session_ctx(g_session), + ds4_session_pos(g_session)); const int eos = ds4_token_eos(g_engine); const int draft_max = ds4_engine_mtp_draft_tokens(g_engine); int produced = 0; @@ -908,9 +916,12 @@ class DS4Backend final : public backend::Backend::Service { if (draft_max > 0 && sp.temperature <= 0.0f) { constexpr int kAcceptedMax = 8; int accepted[kAcceptedMax]; - int cap = std::min(kAcceptedMax, draft_max + 1); + const int remaining = ds4cpp::RemainingGenerationBudget( + n_predict, produced); + const int cap = ds4cpp::SpeculativeAcceptedCapacity( + remaining, draft_max, kAcceptedMax); int n = ds4_session_eval_speculative_argmax( - g_session, first, draft_max, eos, + g_session, first, remaining, eos, accepted, cap, err, sizeof(err)); if (n < 0) { rc = -1; break; } bool stop = false; diff --git a/docs/content/advanced/model-configuration.md b/docs/content/advanced/model-configuration.md index a162a6225319..8e0ca855b0a1 100644 --- a/docs/content/advanced/model-configuration.md +++ b/docs/content/advanced/model-configuration.md @@ -173,6 +173,14 @@ These settings will be used as defaults for all the API calls to the model. | `tfz` | float | `1.0` | Tail free z parameter | | `keep` | int | `0` | Number of tokens to keep from the prompt | +{{% notice note %}} +The DS4 backend preserves its legacy behavior for omitted or non-positive +`max_tokens` values by generating at most 256 tokens. Set `max_tokens` to a +positive value when you need a specific DS4 output limit. After processing the +prompt, DS4 clamps that limit to the available context space and reserves one +context slot for safe generation. +{{% /notice %}} + ### Language and Translation | Field | Type | Description |