The computational core of the generative model server. Engine defines an API that models must adhere to as they plug into the JetStream efficient serving infrastructure.
| 99 | |
| 100 | |
| 101 | class MaxEngine(engine_api.Engine): |
| 102 | """The computational core of the generative model server. |
| 103 | |
| 104 | Engine defines an API that models must adhere to as they plug into the |
| 105 | JetStream efficient serving infrastructure. |
| 106 | """ |
| 107 | |
| 108 | def __init__(self, config: Any, devices: config_lib.Devices | None = None): |
| 109 | self.config = config |
| 110 | |
| 111 | # Mesh definition |
| 112 | devices_array = maxtext_utils.create_device_mesh(config=config, devices=devices) |
| 113 | self._mesh = jax.sharding.Mesh(devices_array, config.mesh_axes) |
| 114 | |
| 115 | # Model and Optimizer definition |
| 116 | quant = quantizations.configure_quantization(config) |
| 117 | self.model = models.transformer_as_linen(config, mesh=self._mesh, quant=quant, model_mode=MODEL_MODE_PREFILL) |
| 118 | self.replicated_sharding = jax.sharding.NamedSharding(self._mesh, P(None)) |
| 119 | |
| 120 | self.abstract_params = None |
| 121 | self.prefill_kv_cache_annotations = None |
| 122 | self.kv_cache_annotations = None |
| 123 | self.kv_cache_annotations_named = None |
| 124 | self.prefill_kv_cache_shardings = None |
| 125 | self.kv_cache_shardings = None |
| 126 | self.state_mesh_annotations = None |
| 127 | self.decode_state_shapes = None |
| 128 | self.decode_state_layouts = None |
| 129 | self.param_layouts = None |
| 130 | self.rng = None |
| 131 | |
| 132 | # Initialize page manager and page state |
| 133 | self.page_manager = None |
| 134 | self.page_state = None |
| 135 | if self.config.attention == "paged": |
| 136 | self.page_manager = PageManager(self.config) |
| 137 | self.page_state = self.page_manager.get_initial_page_state() |
| 138 | |
| 139 | def print_stats(self, label: str): |
| 140 | max_utils.print_mem_stats(label) |
| 141 | max_utils.print_cpu_ram_stats(label) |
| 142 | |
| 143 | def generate_aot( |
| 144 | self, params: Params, decode_state: DecodeState, rng: PRNGKeyType | None = None |
| 145 | ) -> tuple[DecodeState, engine_api.ResultTokens]: |
| 146 | """Wrapper to generate for ahead of time compilation.""" |
| 147 | |
| 148 | return self.generate(params=params, decode_state=decode_state, rng=rng) |
| 149 | |
| 150 | def _compile_generate_and_get_layouts( |
| 151 | self, params: Any, decode_state: Any, rng_shape: Any, xla_flags: dict[str, Any] | None = None |
| 152 | ) -> tuple[Any, Any, Any, Any]: |
| 153 | """Optimal memory layout for params and decode_state.""" |
| 154 | |
| 155 | param_layout = Format(DLL.AUTO) |
| 156 | decode_state_layout = Format(DLL.AUTO) |
| 157 | # Keyword arguments are not yet supported in JAX for specifying shardings. Therefore, all AOT |
| 158 | # compiled functions use arguments instead. |
no outgoing calls