Performs a single, JIT-compiled autoregressive decoding step. This function takes the current decoding state, which includes the KV cache for the sequence generated so far, and generates the next token. It uses the model to predict logits for the next token based on the previous token a
(
self,
params: Params,
decode_state: DecodeState,
*,
sampler: Callable[[Any], Any] | None = None, # pylint: disable=unused-argument
rng: PRNGKeyType | None = None,
page_state: PageState | None = None,
algorithm: str | None = None,
topk: int | None = None,
nucleus_topp: float | None = None,
temperature: float | None = None,
)
| 954 | jax.jit, static_argnums=(0,), donate_argnums=(2,), static_argnames=("algorithm", "topk", "nucleus_topp") |
| 955 | ) |
| 956 | def _generate_jit( |
| 957 | self, |
| 958 | params: Params, |
| 959 | decode_state: DecodeState, |
| 960 | *, |
| 961 | sampler: Callable[[Any], Any] | None = None, # pylint: disable=unused-argument |
| 962 | rng: PRNGKeyType | None = None, |
| 963 | page_state: PageState | None = None, |
| 964 | algorithm: str | None = None, |
| 965 | topk: int | None = None, |
| 966 | nucleus_topp: float | None = None, |
| 967 | temperature: float | None = None, |
| 968 | ) -> tuple[DecodeState, engine_api.ResultTokens]: |
| 969 | """Performs a single, JIT-compiled autoregressive decoding step. |
| 970 | |
| 971 | This function takes the current decoding state, which includes the KV cache |
| 972 | for the sequence generated so far, and generates the next token. It uses the |
| 973 | model to predict logits for the next token based on the previous token and |
| 974 | then samples from that distribution. |
| 975 | |
| 976 | Args: |
| 977 | params: The model parameters. |
| 978 | decode_state: The current state of the decoding process, containing the |
| 979 | KV cache, the previously generated token, the current position, etc. |
| 980 | This argument is donated to save memory. |
| 981 | sampler: A callable for custom sampling logic (currently unused). |
| 982 | rng: JAX random number generator key for sampling. |
| 983 | page_state: The current state of the paged attention manager. |
| 984 | algorithm: The sampling algorithm to use (e.g., 'greedy', 'composite'). |
| 985 | Overrides the engine's default. |
| 986 | topk: The value for top-k sampling. Overrides the engine's default. |
| 987 | nucleus_topp: The value for top-p (nucleus) sampling. Overrides the |
| 988 | engine's default. |
| 989 | temperature: The sampling temperature. Overrides the engine's default. |
| 990 | |
| 991 | Returns: |
| 992 | A tuple containing: |
| 993 | - The updated `DecodeState` with the new KV cache, new token, and |
| 994 | incremented position, ready for the next decoding step. |
| 995 | - An `engine_api.ResultTokens` object containing the newly generated |
| 996 | token and its metadata. |
| 997 | """ |
| 998 | |
| 999 | previous_token = decode_state["tokens"] |
| 1000 | rng, new_rng = jax.random.split(rng) |
| 1001 | # run one step generation |
| 1002 | with self._mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules): |
| 1003 | out_logits, new_vars = self.model.apply( |
| 1004 | params | {"cache": decode_state["cache"]}, |
| 1005 | previous_token, |
| 1006 | decode_state["next_pos"], |
| 1007 | enable_dropout=False, |
| 1008 | model_mode=MODEL_MODE_AUTOREGRESSIVE, |
| 1009 | rngs={"params": new_rng}, |
| 1010 | mutable=["cache"], |
| 1011 | page_state=page_state, |
| 1012 | ) |
| 1013 | out_logits = jax.lax.with_sharding_constraint(out_logits, self.replicated_sharding) |