fit: also take into account n_streams - #27496
Conversation
|
@ServeurpersoCom thanks for testing! feel free to push your fixes directly here |
With a non-unified KV cache the target context now holds n_ctx_train tokens per sequence, while the draft context was still created with n_ctx = 0 and fell back to n_ctx_train / n_streams per sequence. A slot filled beyond that point makes the draft batch fail to decode, and the server answers 500 on the request. The draft context now takes its size from the target context, so both hold the same number of tokens per sequence. Contexts that share their cells with the target no longer need the kv_size override. The memory reserved for the draft model before fitting is measured at the largest context the target can take, since the draft context grows with the target and a fixed byte margin cannot express that.
|
Follow-up on the reservation side. The draft memory is currently folded into fit_params_target as a fixed byte margin. That is right for the mmproj, whose size does not depend on the context, but not for the draft context, which now scales with the target n_ctx. My commit works around it by probing the target n_ctx_train and measuring the draft at the upper bound, so it over-reserves whenever the fit ends up reducing the context. Would you rather have common_fit_params take an optional draft model and measure both at each candidate context? That drops the probe and the whole reservation block in server-context.cpp, and makes the reduce path exact instead of conservative, at the cost of one more parameter in the common API. |
Illustrates the alternative discussed on the draft context fix. The memory of a draft or MTP context is currently handed to the fit as a fixed byte margin, which cannot express a memory that grows with the context the fit is still deciding on. common_fit_params now takes an optional second model that shares the devices of the main one. Its context follows the main context and its memory is measured again whenever that context changes, so the reduce path stays exact instead of conservative. A model that cannot be measured on its own, such as a shared cell MTP context, is skipped with a warning and the main model is fitted alone. This drops the reservation block in the server, which no longer has to probe the trained context size of the target to guess an upper bound.
|
The second commit is the refactor, re-tested: |
2 participants
Overview
This PR fixes this particular setup on llama-server:
--fit -no-kvu -np 4If the machine has plenty of RAM, the max context size will be capped to model's trained context size. This is OK in kv-unified because one single request can still use the full trained context window
However, with no-kvu and for example -np 4, the slot ctx size will be model's trained context divided by 4, which is not something we want. That means we cannot use the full model trained context even if we have plenty of RAM to store 4*n_ctx_trained
Example: testing with
tinygemma3:--fit -no-kvu -np 4setsn_ctx = 131072, and so each slot only get32768tokens--fit -no-kvu -np 4setsn_ctx = 131072*4 = 524288, and so each slot get131072Note: one edge case that I couldn't test: what if we require
-np 4but physical RAM can only hold-np 3. Logically this will be handled, but just haven't testedRequirements