Fix according to upstream changes

b121ca4d · lazymio · 26f7b4af · b121ca4d
Unverified Commit b121ca4d authored Feb 27, 2025 by lazymio
Hide whitespace changes
Inline Side-by-side

Showing with 1 addition and 2 deletions

ktransformers/server/backend/interfaces/ktransformers.py ktransformers/server/backend/interfaces/ktransformers.py +1 -2

No files found.
--- a/ktransformers/server/backend/interfaces/ktransformers.py
+++ b/ktransformers/server/backend/interfaces/ktransformers.py
@@ -201,10 +201,9 @@ class KTransformersInterface(TransformersInterface):
        else:
            logits = self.model(inputs_embeds=inputs_embeds, return_dict=False)[0]

-        self.prepare_logits_wrapper(input_ids, device, temperature, top_p)
        if flashinfer_enabled:
            MLAWrapperSingleton.reset_buffer()
-        self.prepare_logits_wrapper(input_ids, device)
+        self.prepare_logits_wrapper(input_ids, device, temperature, top_p)
        next_token = self.logits_to_token(logits[0, -1, :])
        yield self.append_new_tokens(next_token)