model.py Source File

model.py Source File#

SDK qb Runtime Library: model.py Source File
SDK qb Runtime Library v1.5
MCS001-KR
model.py
Go to the documentation of this file.
1
4
5from typing import List, Optional, Tuple, Union
6
7import numpy as np
8
9import qbruntime.qbruntime as _cQbRuntime
10from .accelerator import Accelerator
11from .future import *
12from .model_variant_handle import *
13from .npu_data import NPUData
14from .pinned_memory import PinnedMemory
15from .type import *
16
17_Shape = Tuple[int, ...]
18
19__all__ = ["Model", "load"]
20
21
24
25
26# input ndarray의 shape이 유효한 shape인지 판별한다.
27def _is_valid_shape(input_shape: _Shape, shape: _Shape) -> bool:
28 if (len(input_shape) < len(shape)) or (len(input_shape) > len(shape) + 1):
29 return False
30 # input을 batch일 경우도 고려하여 [h, w, c] 및 [batch, h, w, c] 모두 고려한다
31 offset = 1 if len(input_shape) > len(shape) else 0
32 for s1, s2 in zip(input_shape[offset:], shape):
33 # Dimensions that allow variable lengths are represented by negative values.
34 # A variable-length dimension only permits multiples of the original value.
35 if s1 % s2 != 0 or (s2 > 0 and s1 != s2):
36 return False
37 return True
38
39
40# input ndarray의 shape를 검사하여 HWC인지 CHW인지 판별한다. HWC/CHW의
41# shape이 동일한 경우, `is_hwc`와 `is_chw`를 모두 true로 반환한다.
42def _find_memory_format(
43 inputs: List[np.ndarray], shapes: List[_Shape]
44) -> Optional[Tuple[bool, bool]]:
45 if len(inputs) != len(shapes):
46 return None
47
48 is_hwc = True
49 is_chw = True
50 for arr, shape in zip(inputs, shapes):
51 shape_hwc = (shape[0], shape[1], shape[2])
52 shape_chw = (shape[2], shape[0], shape[1])
53 is_hwc = is_hwc and _is_valid_shape(arr.shape, shape_hwc)
54 is_chw = is_chw and _is_valid_shape(arr.shape, shape_chw)
55
56 if not is_hwc and not is_chw:
57 return None
58 return is_hwc, is_chw
59
60
61# input ndarray에 맞는 model variant index와 shape를 판별한다.
62def _find_matching_variant_idx_and_memory_format(
63 model, inputs: List[np.ndarray]
64) -> Tuple[int, Tuple[bool, bool]]:
65 variant_idx = None
66 is_hwc = None
67 is_chw = None
68 for i in range(model.get_num_model_variants()):
69 res = _find_memory_format(
70 inputs, model.get_model_variant_handle(i).get_model_input_shape()
71 )
72 if res is not None:
73 variant_idx = i
74 is_hwc, is_chw = res
75 break
76
77 if variant_idx is None:
78 raise ValueError("Input shape is invalid.")
79 return variant_idx, (is_hwc, is_chw)
80
81
82# shape에 맞게 numpy ndarray를 생성한다.
83def _build_outputs(
84 shapes: List[_Shape], is_hwc: bool, dtype: np.dtype
85) -> List[np.ndarray]:
86 outputs = []
87 for shape in shapes:
88 if is_hwc:
89 shape = (shape[0], shape[1], shape[2])
90 else:
91 shape = (shape[2], shape[0], shape[1])
92 outputs.append(np.empty(shape, dtype=dtype))
93 return outputs
94
95
96# output에 들어있는 numpy ndarray의 shape가 올바른지 검사한다.
97def _check_output_shapes(
98 outputs: List[np.ndarray], shapes: List[_Shape], is_hwc: bool, dtype: np.dtype
99) -> None:
100 if len(outputs) != len(shapes):
101 raise ValueError("The number of outputs is different.")
102
103 for output, shape in zip(outputs, shapes):
104 if output.dtype != dtype:
105 raise ValueError("Output dtype mismatch.")
106
107 if is_hwc:
108 shape = (shape[0], shape[1], shape[2])
109 else:
110 shape = (shape[2], shape[0], shape[1])
111 if output.shape != shape:
112 raise ValueError("Output shape mismatch.")
113
114
115class Model:
116 """
117 @brief Represents an AI model loaded from an MXQ file.
118
119 This class loads an AI model from an MXQ file and provides functions to launch it
120 on the NPU and perform inference.
121 """
122
123 def __init__(self, path: str, model_config: Optional[ModelConfig] = None):
124 """
125 @brief Creates a Model object from the specified MXQ model file and configuration.
126
127 Parses the MXQ file and constructs a Model object using the provided configuration,
128 initializing the model with the given settings.
129
130 @note The created Model object must be launched before performing inference.
131 See Model.launch for more details.
132
133 @param[in] path The path to the MXQ model file.
134 @param[in] model_config The configuration settings to initialize the Model.
135 """
136 if model_config is None:
137 self._model = _cQbRuntime.Model(path)
138 else:
139 self._model = _cQbRuntime.Model(path, model_config._model_config)
140
141 # 기존 BufferInfo 대신에 ModelShape를 사용한다.
142 # Model {input,output} shape는 batch를 포함한 4D이다.
145
146 def launch(self, acc: Accelerator) -> None:
147 """
148 @brief Launches the model on the specified Accelerator, which represents
149 the actual NPU.
150
151 @param[in] acc The accelerator on which to launch the model.
152 """
153 self._model.launch(acc._accelerator)
154 self._acc = acc
155
156 def dispose(self) -> None:
157 """
158 @brief Disposes of the model loaded onto the NPU.
159
160 Releases any resources associated with the model on the NPU.
161 """
162 self._model.dispose()
163 self._acc = None
164
165 def is_target(self, core_id: CoreId) -> bool:
166 """
167 @brief Checks if the NPU core specified by CoreId is the target of the model.
168 In other words, whether the model is configured to use the given NPU core.
169
170 @param[in] core_id The CoreId to check.
171 @return True if the model is configured to use the specified CoreId, false
172 otherwise.
173 """
174 return self._model.is_target(core_id._core_id)
175
176 def get_core_mode(self) -> CoreMode:
177 """
178 @brief Retrieves the core mode of the model.
179
180 @return The CoreMode of the model.
181 """
182 return CoreMode(self._model.get_core_mode())
183
184 def get_device_names(self) -> List[str]:
185 """
186 @brief Returns the supported target device name(s) this model can run on.
187
188 Returns device names from the model's target device in the MXQ file. A single
189 model may map to multiple names (e.g. a REGULUS model returns both SoC and
190 USB devices). Any of the returned names can be passed to `Accelerator(device_name=)`.
191
192 @return A list of supported target device names.
193 """
194 return self._model.get_device_names()
195
196 def get_target_cores(self) -> List[CoreId]:
197 """
198 @brief Returns the NPU cores the model is configured to use.
199
200 @return A list of CoreIds representing the target NPU cores.
201 """
202 return [CoreId.from_cpp(target) for target in self._model.target_cores]
203
204 @property
205 def target_cores(self) -> List[CoreId]:
206 """@deprecated"""
207 return [CoreId.from_cpp(target) for target in self._model.target_cores]
208
209 def infer(
210 self,
211 inputs: Union[np.ndarray, List[np.ndarray]],
212 outputs: Optional[List[np.ndarray]] = None,
213 cache_size: int = 0,
214 params: Optional[List[BatchParam]] = None,
215 ) -> Optional[List[np.ndarray]]:
216 """
217 @brief Performs inference.
218
219 Fowllowing types of inference supported.
220 1. infer(in:List[numpy]) -> List[numpy] (float / int)
221 2. infer(in:numpy) -> List[numpy] (float / int)
222 3. infer(in:List[numpy], out:List[numpy]) (float / int)
223 4. infer(in:List[numpy], out:List[]) (float / int)
224 5. infer(in:numpy, out:List[numpy]) (float / int)
225 6. infer(in:numpy, out:List[]) (float / int)
226
227 @param[in] inputs Input data as a single numpy.ndarray or a list
228 of numpy.ndarray's.
229 @param[out] outputs Optional pre-allocated list of numpy.ndarray's
230 to store inference results.
231 @param[in] cache_size The number of tokens accumulated in the KV cache so far.
232 @param[in] params A List of `BatchParam`, specifying each batch's information
233 for BatchLLM inference. If `params` is specified,
234 `cache_size` is ignored.
235 @return Inference results as a list of numpy.ndarray.
236 """
237 return self._infer(inputs, outputs, cache_size, params=params)
238
239 def infer_hwc(
240 self,
241 inputs: Union[np.ndarray, List[np.ndarray]],
242 outputs: Optional[List[np.ndarray]] = None,
243 cache_size: int = 0,
244 params: Optional[List[BatchParam]] = None,
245 ) -> Optional[List[np.ndarray]]:
246 return self._infer(inputs, outputs, cache_size, True, params)
247
248 def infer_chw(
249 self,
250 inputs: Union[np.ndarray, List[np.ndarray]],
251 outputs: Optional[List[np.ndarray]] = None,
252 cache_size: int = 0,
253 params: Optional[List[BatchParam]] = None,
254 ) -> Optional[List[np.ndarray]]:
255 return self._infer(inputs, outputs, cache_size, False, params)
256
257 def _infer(
258 self,
259 inputs: Union[np.ndarray, List[np.ndarray]],
260 outputs: Optional[List[np.ndarray]],
261 cache_size: int,
262 is_target_hwc: Optional[bool] = None,
263 params: Optional[List[BatchParam]] = None,
264 ) -> Optional[List[np.ndarray]]:
265 if not isinstance(inputs, list):
266 inputs = [inputs]
267
268 variant_idx, (is_hwc, is_chw) = _find_matching_variant_idx_and_memory_format(
269 self, inputs
270 )
271 if (is_target_hwc is not None) and (
272 (is_target_hwc != is_hwc) and (is_target_hwc == is_chw)
273 ):
274 raise ValueError("Input shape is invalid.")
275 elif is_target_hwc is None:
276 is_target_hwc = is_hwc
277 inputs = [np.ascontiguousarray(i) for i in inputs]
278
279 infer_func = self._model.infer if is_target_hwc else self._model.infer_chw
280 if outputs is None:
281 # No Output Parameter
282 if params == None:
283 return [np.asarray(o) for o in infer_func(inputs, cache_size)]
284 else:
285 return [
286 np.asarray(o)
287 for o in infer_func(
288 inputs, [param._batch_param for param in params]
289 )
290 ]
291 else:
292 if outputs:
293 _check_output_shapes(
294 outputs,
296 is_target_hwc,
297 inputs[0].dtype,
298 )
299 for oi in range(len(outputs)):
300 outputs[oi] = np.ascontiguousarray(outputs[oi])
301 else:
302 outputs[:] = _build_outputs(
304 is_target_hwc,
305 inputs[0].dtype,
306 )
307
308 if params == None:
309 infer_func(inputs, outputs, cache_size)
310 else:
311 infer_func(inputs, outputs, [param._batch_param for param in params])
312
314 self,
315 inputs: Union[
316 np.ndarray,
317 List[np.ndarray],
318 ],
319 cache_size: int = 0,
320 ) -> List[np.ndarray]:
321 """
322 @brief int8_t-to-float inference
323 Performs inference with input and output elements of type `int8_t`
324
325 Using these inference APIs requires manual scaling (quantization)
326 of float values to `int8_t` for input.
327
328 @note These APIs are intended for advanced use rather than typical usage.
329 """
330 return self._infer_to_float(inputs, cache_size)
331
332 def infer_hwc_to_float(
333 self,
334 inputs: Union[
335 np.ndarray,
336 List[np.ndarray],
337 ],
338 cache_size: int = 0,
339 ) -> List[np.ndarray]:
340 return self._infer_to_float(inputs, cache_size, True)
341
342 def infer_chw_to_float(
343 self,
344 inputs: Union[
345 np.ndarray,
346 List[np.ndarray],
347 ],
348 cache_size: int = 0,
349 ) -> List[np.ndarray]:
350 return self._infer_to_float(inputs, cache_size, False)
351
353 self,
354 inputs: Union[
355 np.ndarray,
356 List[np.ndarray],
357 ],
358 cache_size: int,
359 is_target_hwc: Optional[bool] = None,
360 ) -> List[np.ndarray]:
361 """
362 @brief int8_t-to-float inference
363 Performs inference with input and output elements of type `int8_t`
364
365 Using these inference APIs requires manual scaling (quantization)
366 of float values to `int8_t` for input.
367
368 @note These APIs are intended for advanced use rather than typical usage.
369 """
370 if not isinstance(inputs, list):
371 inputs = [inputs]
372
373 _, (is_hwc, is_chw) = _find_matching_variant_idx_and_memory_format(self, inputs)
374 if (is_target_hwc is not None) and (
375 (is_target_hwc != is_hwc) and (is_target_hwc == is_chw)
376 ):
377 raise ValueError("Input shape is invalid.")
378 elif is_target_hwc is None:
379 is_target_hwc = is_hwc
380 inputs = [np.ascontiguousarray(i) for i in inputs]
381
382 if is_target_hwc:
383 outputs = self._model.infer_to_float(inputs, cache_size)
384 else:
385 outputs = self._model.infer_chw_to_float(inputs, cache_size)
386
387 return [np.asarray(o) for o in outputs]
388
390 self,
391 inputs: List[Buffer],
392 outputs: List[Buffer],
393 shape: List[List[int]] = [],
394 cache_size: int = 0,
395 ) -> None:
396 """
397 @brief Buffer-to-Buffer inference
398
399 Performs inference using input and output elements in the NPU’s internal data type.
400 The inference operates on buffers allocated via the following APIs:
401
402 - `Model.acquire_input_buffer()`
403 - `Model.acquire_output_buffer()`
404 - `ModelVariantHandle.acquire_input_buffer()`
405 - `ModelVariantHandle.acquire_output_buffer()`
406
407 Additionally, `Model.reposition_inputs()`, `Model.reposition_outputs()`,
408 `ModelVariantHandle.reposition_inputs()`, `ModelVariantHandle.reposition_outputs()`
409 must be used properly.
410
411 @note These APIs are intended for advanced use rather than typical usage.
412 """
413 self._model.infer_buffer(
414 [i._buffer for i in inputs], [o._buffer for o in outputs], shape, cache_size
415 )
416
417 def infer_speedrun(self) -> None:
418 """
419 @brief Development-only API for measuring pure NPU inference speed.
420
421 Runs NPU inference without uploading inputs and without retrieving outputs.
422 """
424
426 self,
427 inputs: List[PinnedMemory],
428 outputs: Optional[List[PinnedMemory]] = None,
429 cache_size: int = 0,
430 ) -> Optional[List[PinnedMemory]]:
431 """
432 @brief Performs inference directly on pinned memory buffers (zero-copy).
433
434 The NPU reads the inputs from and writes the outputs into the pinned
435 buffers in place, avoiding host-to-device copies of the I/O tensors. Write
436 the input data into each input buffer via indexing (e.g. `pm[...] = data`)
437 before calling, and read the results from the output buffers afterwards.
438
439 @note This is an experimental API and is only supported on REGULUS device.
440 It is not supported for models that use CPU offload.
441
442 @param inputs A list of PinnedMemory buffers holding the input data. All
443 inputs must share the same dtype (`numpy.float32` or
444 `numpy.uint8`).
445 @param outputs A list of pre-allocated PinnedMemory buffers (dtype
446 `numpy.float32`) that will receive the output data. If `None`
447 (the default), the output buffers are allocated internally and
448 returned by this method.
449 @param cache_size The size of the cache to use for inference.
450 @return If `outputs` is `None`, a list of newly allocated PinnedMemory
451 buffers (dtype `numpy.float32`) holding the output data; otherwise
452 `None` (the given `outputs` are filled in place).
453 """
454 if outputs is None:
455 _pms = self._model.infer_pinned_memory(
456 [pm._pinned_memory for pm in inputs],
457 [],
458 cache_size,
459 )
460 return [PinnedMemory(_pm) for _pm in _pms]
461
463 [pm._pinned_memory for pm in inputs],
464 [pm._pinned_memory for pm in outputs],
465 cache_size,
466 )
467 return None
468
470 self,
471 shape: List[int],
472 idx: int = 0,
473 upload: bool = False,
474 dtype: Optional[DataType] = None,
475 ) -> NPUData:
476 """
477 @brief Acquires an NPUData for the model input at the given index.
478
479 @note This is an advanced API rather than a typical usage. (Running inference on
480 the acquired NPUData via `infer_npu_data` additionally requires a
481 single-NPU-op, non-CPU-offload, relocatable (MXQv7+) model.)
482 @warning This API is in beta: it may still contain bugs, and its behavior may
483 change in a future release.
484
485 @param shape The shape of the input tensor, in NHWC/HWC or NCHW/CHW layout.
486 @param idx The index of the model input.
487 @param upload If True, NPU memory is allocated and the data is placed on the
488 NPU; otherwise the data is kept on the host (CPU).
489 @param dtype The element type of the data. `None` (the default) uses the
490 model's user-facing input type. Pass `DataType.Int8` to hand the
491 model data that is already quantized to the NPU-native type, which
492 skips the host-side quantization.
493 @return The acquired NPUData.
494 """
495 shape = [int(s) for s in shape]
496 if dtype is None:
497 return NPUData(self._model.acquire_input_npu_data(shape, idx, upload))
498 return NPUData(
499 self._model.acquire_input_npu_data(shape, idx, upload, dtype.value)
500 )
501
503 self,
504 shape: List[int],
505 idx: int = 0,
506 upload: bool = False,
507 dtype: Optional[DataType] = None,
508 ) -> NPUData:
509 """
510 @brief Acquires an NPUData for the model output at the given index.
511
512 @note This is an advanced API rather than a typical usage. (Running inference on
513 the acquired NPUData via `infer_npu_data` additionally requires a
514 single-NPU-op, non-CPU-offload, relocatable (MXQv7+) model.)
515 @warning This API is in beta: it may still contain bugs, and its behavior may
516 change in a future release.
517
518 @param shape The shape of the output tensor, in NHWC/HWC or NCHW/CHW layout.
519 @param idx The index of the model output.
520 @param upload If True, NPU memory is allocated and the data is placed on the
521 NPU; otherwise the data is kept on the host (CPU).
522 @param dtype The element type of the data. `None` (the default) uses the
523 model's user-facing output type. Pass `DataType.Int8` to receive
524 the raw NPU-native output, which skips the host-side
525 dequantization.
526 @return The acquired NPUData.
527 """
528 shape = [int(s) for s in shape]
529 if dtype is None:
530 return NPUData(self._model.acquire_output_npu_data(shape, idx, upload))
531 return NPUData(
532 self._model.acquire_output_npu_data(shape, idx, upload, dtype.value)
533 )
534
536 self,
537 inputs: List[NPUData],
538 outputs: Optional[List[NPUData]] = None,
539 cache_size: int = 0,
540 ) -> Optional[List[NPUData]]:
541 """
542 @brief Performs inference using NPUData for both inputs and outputs.
543
544 Every input and output NPUData must share the same residency: either all of
545 them are on the host (acquired with `upload=False`) or all of them are already
546 on the NPU (acquired with `upload=True`). A mixed set results in an error.
547
548 @note This is an advanced API rather than a typical usage. Only single-NPU-op
549 (non-CPU-offload), relocatable (MXQv7+) models are supported.
550 @note Every NPUData must have been acquired with an NHWC/HWC shape; use
551 `infer_npu_data_chw` for NCHW/CHW-shaped NPUData. A mismatch raises
552 `Model_ShapeMismatched`.
553 @warning This API is in beta: it may still contain bugs, and its behavior may
554 change in a future release.
555
556 @param inputs A list of input NPUData.
557 @param outputs An optional list of pre-allocated output NPUData. If None (the
558 default), the output NPUData are acquired internally and returned.
559 @param cache_size The number of tokens accumulated in the KV cache so far.
560 @return If `outputs` is None, a list of output NPUData holding the inference
561 results; otherwise None (the given `outputs` are filled in place).
562 """
563 if outputs is None:
564 _outs = self._model.infer_npu_data(
565 [d._npu_data for d in inputs],
566 [],
567 cache_size,
568 )
569 return [NPUData(_out) for _out in _outs]
570
572 [d._npu_data for d in inputs],
573 [d._npu_data for d in outputs],
574 cache_size,
575 )
576 return None
577
579 self,
580 inputs: List[NPUData],
581 outputs: Optional[List[NPUData]] = None,
582 cache_size: int = 0,
583 ) -> Optional[List[NPUData]]:
584 """
585 @brief Performs inference using NPUData in NCHW/CHW layout.
586
587 Same as `infer_npu_data` except that every input and output NPUData must have
588 been acquired with a NCHW/CHW shape; a mismatch raises `Model_ShapeMismatched`.
589
590 @note This is an advanced API rather than a typical usage. Only single-NPU-op
591 (non-CPU-offload), relocatable (MXQv7+) models are supported.
592 @warning This API is in beta: it may still contain bugs, and its behavior may
593 change in a future release.
594
595 @param inputs A list of input NPUData.
596 @param outputs An optional list of pre-allocated output NPUData. If None (the
597 default), the output NPUData are acquired internally and returned.
598 @param cache_size The number of tokens accumulated in the KV cache so far.
599 @return If `outputs` is None, a list of output NPUData holding the inference
600 results; otherwise None (the given `outputs` are filled in place).
601 """
602 if outputs is None:
603 _outs = self._model.infer_npu_data_chw(
604 [d._npu_data for d in inputs],
605 [],
606 cache_size,
607 )
608 return [NPUData(_out) for _out in _outs]
609
611 [d._npu_data for d in inputs],
612 [d._npu_data for d in outputs],
613 cache_size,
614 )
615 return None
616
618 self,
619 inputs: Union[np.ndarray, List[np.ndarray]],
620 ) -> Future:
621 """
622 @brief Asynchronous Inference
623
624 Performs inference asynchronously.
625
626 To use asynchronous inference, the model must be created using a `ModelConfig`
627 object with the async pipeline configured to be enabled. This is done by calling
628 @ref ModelConfig.set_async_pipeline_enabled
629 "ModelConfig.set_async_pipeline_enabled(True)" before passing the configuration to
630 `Model()`.
631
632 Example:
633 @code
634 import qbruntime
635
636 mc = qbruntime.ModelConfig()
637 mc.set_async_pipeline_enabled(True)
638
639 model = qbruntime.Model(MXQ_PATH, mc)
640 acc = qbruntime.Accelerator()
641
642 model.launch(acc)
643
644 future = model.infer_async(inputs)
645
646 ret = future.get()
647 @endcode
648
649 @note Currently, only CNN-based models are supported, as asynchronous execution is
650 particularly effective for this type of workload.
651
652 @note Limitations:
653 - RNN/LSTM and LLM models are not supported yet.
654 - Models requiring CPU offloading are not supported yet.
655 - Currently, only single-batch inference is supported (i.e., N = 1).
656 - Currently, Buffer inference is not supported. The following types
657 are supported in the synchronous API for advanced use cases, but are not
658 yet available for asynchronous inference:
659 - Buffer to Buffer
660 - Buffer to float
661 """
662 if not isinstance(inputs, list):
663 inputs = [inputs]
664 _, (is_hwc, _) = _find_matching_variant_idx_and_memory_format(self, inputs)
665 inputs = [np.ascontiguousarray(i) for i in inputs]
666 infer_async_func = (
667 self._model.infer_async if is_hwc else self._model.infer_async_chw
668 )
669 return Future.from_cpp(infer_async_func(inputs), inputs)
670
672 self,
673 inputs: Union[np.ndarray, List[np.ndarray]],
674 ) -> Future:
675 """
676 @brief This method supports int8_t-to-float asynchronous inference.
677
678 @param[in] inputs Input data as a single numpy.ndarray or a list
679 of numpy.ndarray's.
680
681 @return A future that can be used to retrieve the inference result.
682 """
683 if not isinstance(inputs, list):
684 inputs = [inputs]
685 _, (is_hwc, _) = _find_matching_variant_idx_and_memory_format(self, inputs)
686 inputs = [np.ascontiguousarray(i) for i in inputs]
687 infer_async_func = (
688 self._model.infer_async_to_float
689 if is_hwc
690 else self._model.infer_async_chw_to_float
691 )
692 return Future.from_cpp(infer_async_func(inputs), inputs)
693
695 self,
696 inputs: List[np.ndarray],
697 input_bufs: List[Buffer],
698 seqlens: List[List[int]] = [],
699 ) -> None:
700 """Reposition input"""
701 inputs = [np.ascontiguousarray(i) for i in inputs]
703 inputs, [buf._buffer for buf in input_bufs], seqlens
704 )
705
707 self,
708 output_bufs: List[Buffer],
709 outputs: List[np.ndarray],
710 seqlens: List[List[int]] = [],
711 ) -> None:
712 """Reposition output"""
713 if len(outputs) != len(self._output_shape):
714 outputs.clear()
715 for shape in self._output_shape:
716 outputs.append(np.empty(shape=shape, dtype=np.float32))
717 else:
718 for oi in range(len(outputs)):
719 outputs[oi] = np.ascontiguousarray(outputs[oi])
721 [buf._buffer for buf in output_bufs], outputs, seqlens
722 )
723
724 def get_num_model_variants(self) -> int:
725 """
726 @brief Returns the total number of model variants available in this model.
727
728 The `variant_idx` parameter passed to `Model.get_model_variant_handle()` must be
729 in the range [0, return value of this function).
730
731 @return The total number of model variants.
732 """
733 return self._model.get_num_model_variants()
734
735 def get_model_variant_handle(self, variant_idx) -> ModelVariantHandle:
736 """
737 @brief Retrieves a handle to the specified model variant.
738
739 Use the returned `ModelVariantHandle` to query details such as input and output
740 shapes for the selected variant.
741
742 @param[in] variant_idx Index of the model variant to retrieve.
743 Must be in the range [0, getNumModelVariants()).
744
745 @return A `ModelVariantHandle` object if successful;
746 otherwise, raise qbruntime.QbRuntimeError "Model_InvalidVariantIdx".
747 """
748 return ModelVariantHandle.from_cpp(
749 self._model.get_model_variant_handle(variant_idx)
750 )
751
752 def get_model_input_shape(self) -> List[_Shape]:
753 """
754 @brief Returns the input shape of the model.
755
756 @return A list of input shape of the model.
757 """
758 return self._model.get_model_input_shape()
759
760 def get_model_output_shape(self) -> List[_Shape]:
761 """
762 @brief Returns the output shape of the model.
763
764 @return A list of output shape of the model.
765 """
766 return self._model.get_model_output_shape()
767
768 def get_input_scale(self) -> List[Scale]:
769 """
770 @brief Returns the input quantization scale(s) of the model.
771
772 @return A list of input scales.
773 """
774 return [Scale.from_cpp(s) for s in self._model.get_input_scale()]
775
776 def get_output_scale(self) -> List[Scale]:
777 """
778 @brief Returns the output quantization scale(s) of the model.
779
780 @return A list of output scales.
781 """
782 return [Scale.from_cpp(s) for s in self._model.get_output_scale()]
783
784 def get_input_buffer_info(self) -> List[BufferInfo]:
785 """
786 @brief Returns the input buffer information for the model.
787
788 @return A list of input buffer information.
789 """
790 return [BufferInfo.from_cpp(bi) for bi in self._model.get_input_buffer_info()]
791
792 def get_output_buffer_info(self) -> List[BufferInfo]:
793 """
794 @brief Returns the output buffer information of the model.
795
796 @return A list of output buffer information.
797 """
798 return [BufferInfo.from_cpp(bi) for bi in self._model.get_output_buffer_info()]
799
800 def get_model_input_data_type(self) -> DataType:
801 """
802 @brief Returns a data type for model inputs.
803
804 @return An input data type.
805 """
807
808 def get_model_output_data_type(self) -> DataType:
809 """
810 @brief Returns a data type for model outputs.
811
812 @return An output data type.
813 """
815
816 def acquire_input_buffer(self, seqlens: List[List[int]] = []) -> List[Buffer]:
817 """
818 @brief Buffer Management API
819
820 Acquires list of `Buffer` for input.
821 These API is required when calling `Model.infer_buffer()`.
822
823 @note These APIs are intended for advanced use rather than typical usage.
824 """
825 return [Buffer(b) for b in self._model.acquire_input_buffer(seqlens)]
826
827 def acquire_output_buffer(self, seqlens: List[List[int]] = []) -> List[Buffer]:
828 """
829 @brief Buffer Management API
830
831 Acquires list of `Buffer` for output.
832 These API is required when calling `Model.infer_buffer()`.
833
834 @note These APIs are intended for advanced use rather than typical usage.
835 """
836 return [Buffer(b) for b in self._model.acquire_output_buffer(seqlens)]
837
838 def release_buffer(self, buffer: List[Buffer]) -> None:
839 """
840 @brief Buffer Management API
841
842 Deallocate acquired Input/Output buffer
843
844 @note These APIs are intended for advanced use rather than typical usage.
845 """
846 self._model.release_buffer([b._buffer for b in buffer])
847
848 def get_identifier(self) -> int:
849 """
850 @brief Returns the model's unique identifier.
851
852 This identifier distinguishes multiple models within a single user program.
853 It is assigned incrementally, starting from 0 (e.g., 0, 1, 2, 3, ...).
854
855 @return The model identifier.
856 """
857 return self._model.get_identifier()
858
859 def get_model_path(self) -> str:
860 """
861 @brief Returns the path to the MXQ model file associated with the Model.
862
863 @return The MXQ file path.
864 """
865 return self._model.get_model_path()
866
867 def get_cache_infos(self) -> List[CacheInfo]:
868 """
869 @brief Returns informations of KV-cache of the model.
870
871 @return A list of CacheInfo objects.
872 """
873 return [CacheInfo.from_cpp(c) for c in self._model.get_cache_infos()]
874
875 def dump_cache_memory(self, cache_id: int = 0) -> List[bytes]:
876 """
877 @brief Dumps the KV cache memory into buffers.
878
879 Writes the current KV cache data into provided buffers.
880
881 @param[in] cache_id Index of target cache.
882
883 @return A list of bytes containing the KV cache data.
884 """
885 bufs = self._model.dump_cache_memory(cache_id)
886 return [np.asarray(buf, np.int8).tobytes() for buf in bufs]
887
888 def load_cache_memory(self, bufs: List[bytes], cache_id: int = 0) -> None:
889 """
890 @brief Loads the KV cache memory from buffers.
891
892 Restores the KV cache from the provided buffers.
893
894 @param[in] bufs A list of bytes containing the KV cache
895 """
897 [np.frombuffer(buf, dtype=np.int8) for buf in bufs], cache_id
898 )
899
900 def dump_cache_memory_to(self, cache_dir: str, cache_id: int = 0) -> None:
901 """
902 @brief Dumps KV cache memory to files in the specified directory.
903
904 Writes the KV cache data to binary files within the given directory.
905 Each file is named using the format: `cache_<layer_hash>.bin`.
906
907 @param[in] cache_dir Path to the directory where KV cache files will be saved.
908 @param[in] cache_id Index of target cache.
909 """
910 self._model.dump_cache_memory(cache_dir, cache_id)
911
912 def load_cache_memory_from(self, cache_dir: str, cache_id: int = 0) -> None:
913 """
914 @brief Loads the KV cache memory from files in the specified directory.
915
916 Reads KV cache data from files within the given directory and restores them.
917 Each file is named using the format: `cache_<layer_hash>.bin`.
918
919 @param[in] cache_dir Path to the directory where KV cache files are saved.
920 """
921 self._model.load_cache_memory(cache_dir, cache_id)
922
923 def dump_cache_memory_by_name(self, name: str, cache_id: int = 0) -> bytes:
924 """
925 @brief Dumps a single KV cache memory into a buffer.
926
927 Writes the current data of the KV cache whose `CacheInfo.name` matches `name`
928 into a buffer and returns it.
929
930 @param[in] name Name of the target KV cache. See `CacheInfo.name`.
931 @param[in] cache_id Index of target cache.
932
933 @return Bytes containing the KV cache data.
934 """
935 buf = self._model.dump_cache_memory_by_name(name, cache_id)
936 return np.asarray(buf, np.int8).tobytes()
937
939 self, layer_hash: str, cache_id: int = 0
940 ) -> bytes:
941 """
942 @brief Dumps a single KV cache memory into a buffer.
943
944 Writes the current data of the KV cache whose `CacheInfo.layer_hash` matches
945 `layer_hash` into a buffer and returns it.
946
947 @param[in] layer_hash Layer hash of the target KV cache. See
948 `CacheInfo.layer_hash`.
949 @param[in] cache_id Index of target cache.
950
951 @return Bytes containing the KV cache data.
952 """
953 buf = self._model.dump_cache_memory_by_layer_hash(layer_hash, cache_id)
954 return np.asarray(buf, np.int8).tobytes()
955
957 self, buf: bytes, name: str, cache_id: int = 0
958 ) -> None:
959 """
960 @brief Loads a single KV cache memory from a buffer.
961
962 Restores only the KV cache whose `CacheInfo.name` matches `name` from the given
963 buffer. The other KV caches are left untouched.
964
965 @param[in] buf Bytes containing the KV cache data. Its size must be equal to
966 the size of the target KV cache.
967 @param[in] name Name of the target KV cache. See `CacheInfo.name`.
968 @param[in] cache_id Index of target cache.
969 """
971 np.frombuffer(buf, dtype=np.int8), name, cache_id
972 )
973
975 self, buf: bytes, layer_hash: str, cache_id: int = 0
976 ) -> None:
977 """
978 @brief Loads a single KV cache memory from a buffer.
979
980 Restores only the KV cache whose `CacheInfo.layer_hash` matches `layer_hash`
981 from the given buffer. The other KV caches are left untouched.
982
983 @param[in] buf Bytes containing the KV cache data. Its size must be equal to
984 the size of the target KV cache.
985 @param[in] layer_hash Layer hash of the target KV cache. See
986 `CacheInfo.layer_hash`.
987 @param[in] cache_id Index of target cache.
988 """
990 np.frombuffer(buf, dtype=np.int8), layer_hash, cache_id
991 )
992
993 def reset_fixed_tail_cache_memory(self, cache_id: int = 0) -> None:
994 """
995 @brief Reset the fixed tail KV cache memory.
996
997 Fills the entire memory region of every KV cache whose `CacheType` is
998 `CacheType.FixedTailCache` with zeros. The caches of the other types are left
999 untouched.
1000
1001 @note This does not update the cache size tracked by the caller;
1002 the caller must reset its own `cache_size` to 0 after this call.
1003
1004 @param[in] cache_id Index of target cache.
1005 """
1007
1009 self, cache_size: int, tail_size: int, mask: List[bool]
1010 ) -> int:
1011 """
1012 @brief Filter the tail of the KV cache memory
1013
1014 Retains the desired caches in the tail of the KV cache memory, excludes the others,
1015 and shifts the remaining caches forward.
1016
1017 @param[in] cache_size The number of tokens accumulated in the KV cache so far.
1018 @param[in] tail_size The tail size of the KV cache to filter (<=32).
1019 @param[in] mask A mask indicating tokens to retain or exclude at the tail of the KV
1020 cache.
1021
1022 @return New cache size after tail filtering.
1023 """
1024 return self._model.filter_cache_tail(cache_size, tail_size, mask)
1025
1026 def move_cache_tail(self, num_head: int, num_tail: int, cache_size: int) -> int:
1027 """
1028 @brief Moves the tail of the KV cache memory to the end of the head.
1029
1030 Slice the tail of the KV cache memory up to the specified size
1031 and moves it to the designated cache position.
1032
1033 @param[in] num_head The size of the KV cache head where the tail is appended.
1034 @param[in] num_tail The size of the KV cache tail to be moved.
1035 @param[in] cache_size The total number of tokens accumulated in the KV cache so
1036 far.
1037
1038 @return The updated cache size after moving the tail.
1039 """
1040 return self._model.move_cache_tail(num_head, num_tail, cache_size)
1041
1042
1043def load(path: str, model_config: Optional[ModelConfig] = None) -> Model:
1044 """
1045 @brief Single-step inference API. Creates model and uploads the model
1046 into NPU immediately.
1047
1048 This operation performs the Accelerator declaration, Model declaration,
1049 and launch in a single step.
1050 """
1051 acc = Accelerator()
1052 model = Model(path, model_config)
1053 model.launch(acc)
1054 return model
1055
1056
1057
Represents an accelerator, i.e., an NPU, used for executing models.
Represents an AI model loaded from an MXQ file.
Definition model.py:115
DataType get_model_input_data_type(self)
Returns a data type for model inputs.
Definition model.py:800
None launch(self, Accelerator acc)
Launches the model on the specified Accelerator, which represents the actual NPU.
Definition model.py:146
None reposition_outputs(self, List[Buffer] output_bufs, List[np.ndarray] outputs, List[List[int]] seqlens=[])
Reposition output.
Definition model.py:711
List[Scale] get_input_scale(self)
Returns the input quantization scale(s) of the model.
Definition model.py:768
CoreMode get_core_mode(self)
Retrieves the core mode of the model.
Definition model.py:176
List[CoreId] get_target_cores(self)
Returns the NPU cores the model is configured to use.
Definition model.py:196
List[_Shape] _input_shape
Definition model.py:143
bytes dump_cache_memory_by_layer_hash(self, str layer_hash, int cache_id=0)
Dumps a single KV cache memory into a buffer.
Definition model.py:940
bytes dump_cache_memory_by_name(self, str name, int cache_id=0)
Dumps a single KV cache memory into a buffer.
Definition model.py:923
List[CoreId] target_cores(self)
Definition model.py:205
List[_Shape] _output_shape
Definition model.py:144
List[Buffer] acquire_input_buffer(self, List[List[int]] seqlens=[])
Buffer Management API.
Definition model.py:816
Optional[List[PinnedMemory]] infer_pinned_memory(self, List[PinnedMemory] inputs, Optional[List[PinnedMemory]] outputs=None, int cache_size=0)
Performs inference directly on pinned memory buffers (zero-copy).
Definition model.py:430
List[BufferInfo] get_output_buffer_info(self)
Returns the output buffer information of the model.
Definition model.py:792
str get_model_path(self)
Returns the path to the MXQ model file associated with the Model.
Definition model.py:859
Optional[List[NPUData]] infer_npu_data_chw(self, List[NPUData] inputs, Optional[List[NPUData]] outputs=None, int cache_size=0)
Performs inference using NPUData in NCHW/CHW layout.
Definition model.py:583
Future infer_async(self, Union[np.ndarray, List[np.ndarray]] inputs)
Asynchronous Inference.
Definition model.py:620
DataType get_model_output_data_type(self)
Returns a data type for model outputs.
Definition model.py:808
NPUData acquire_output_npu_data(self, List[int] shape, int idx=0, bool upload=False, Optional[DataType] dtype=None)
Acquires an NPUData for the model output at the given index.
Definition model.py:508
Optional[List[NPUData]] infer_npu_data(self, List[NPUData] inputs, Optional[List[NPUData]] outputs=None, int cache_size=0)
Performs inference using NPUData for both inputs and outputs.
Definition model.py:540
List[Buffer] acquire_output_buffer(self, List[List[int]] seqlens=[])
Buffer Management API.
Definition model.py:827
int filter_cache_tail(self, int cache_size, int tail_size, List[bool] mask)
Filter the tail of the KV cache memory.
Definition model.py:1010
List[np.ndarray] _infer_to_float(self, Union[np.ndarray, List[np.ndarray],] inputs, int cache_size, Optional[bool] is_target_hwc=None)
int8_t-to-float inference Performs inference with input and output elements of type int8_t
Definition model.py:360
int get_num_model_variants(self)
Returns the total number of model variants available in this model.
Definition model.py:724
None load_cache_memory_from(self, str cache_dir, int cache_id=0)
Loads the KV cache memory from files in the specified directory.
Definition model.py:912
bool is_target(self, CoreId core_id)
Checks if the NPU core specified by CoreId is the target of the model.
Definition model.py:165
None infer_speedrun(self)
Development-only API for measuring pure NPU inference speed.
Definition model.py:417
None load_cache_memory_by_layer_hash(self, bytes buf, str layer_hash, int cache_id=0)
Loads a single KV cache memory from a buffer.
Definition model.py:976
None dump_cache_memory_to(self, str cache_dir, int cache_id=0)
Dumps KV cache memory to files in the specified directory.
Definition model.py:900
Future infer_async_to_float(self, Union[np.ndarray, List[np.ndarray]] inputs)
This method supports int8_t-to-float asynchronous inference.
Definition model.py:674
ModelVariantHandle get_model_variant_handle(self, variant_idx)
Retrieves a handle to the specified model variant.
Definition model.py:735
None load_cache_memory(self, List[bytes] bufs, int cache_id=0)
Loads the KV cache memory from buffers.
Definition model.py:888
List[_Shape] get_model_output_shape(self)
Returns the output shape of the model.
Definition model.py:760
List[CacheInfo] get_cache_infos(self)
Returns informations of KV-cache of the model.
Definition model.py:867
None release_buffer(self, List[Buffer] buffer)
Buffer Management API.
Definition model.py:838
List[bytes] dump_cache_memory(self, int cache_id=0)
Dumps the KV cache memory into buffers.
Definition model.py:875
None dispose(self)
Disposes of the model loaded onto the NPU.
Definition model.py:156
List[_Shape] get_model_input_shape(self)
Returns the input shape of the model.
Definition model.py:752
List[Scale] get_output_scale(self)
Returns the output quantization scale(s) of the model.
Definition model.py:776
None infer_buffer(self, List[Buffer] inputs, List[Buffer] outputs, List[List[int]] shape=[], int cache_size=0)
Buffer-to-Buffer inference.
Definition model.py:395
List[np.ndarray] infer_to_float(self, Union[np.ndarray, List[np.ndarray],] inputs, int cache_size=0)
int8_t-to-float inference Performs inference with input and output elements of type int8_t
Definition model.py:320
List[str] get_device_names(self)
Returns the supported target device name(s) this model can run on.
Definition model.py:184
__init__(self, str path, Optional[ModelConfig] model_config=None)
Creates a Model object from the specified MXQ model file and configuration.
Definition model.py:123
Optional[List[np.ndarray]] _infer(self, Union[np.ndarray, List[np.ndarray]] inputs, Optional[List[np.ndarray]] outputs, int cache_size, Optional[bool] is_target_hwc=None, Optional[List[BatchParam]] params=None)
Definition model.py:264
Optional[List[np.ndarray]] infer(self, Union[np.ndarray, List[np.ndarray]] inputs, Optional[List[np.ndarray]] outputs=None, int cache_size=0, Optional[List[BatchParam]] params=None)
Performs inference.
Definition model.py:215
NPUData acquire_input_npu_data(self, List[int] shape, int idx=0, bool upload=False, Optional[DataType] dtype=None)
Acquires an NPUData for the model input at the given index.
Definition model.py:475
None reposition_inputs(self, List[np.ndarray] inputs, List[Buffer] input_bufs, List[List[int]] seqlens=[])
Reposition input.
Definition model.py:699
List[BufferInfo] get_input_buffer_info(self)
Returns the input buffer information for the model.
Definition model.py:784
None load_cache_memory_by_name(self, bytes buf, str name, int cache_id=0)
Loads a single KV cache memory from a buffer.
Definition model.py:958
None reset_fixed_tail_cache_memory(self, int cache_id=0)
Reset the fixed tail KV cache memory.
Definition model.py:993
int get_identifier(self)
Returns the model's unique identifier.
Definition model.py:848
int move_cache_tail(self, int num_head, int num_tail, int cache_size)
Moves the tail of the KV cache memory to the end of the head.
Definition model.py:1026
A model input or output tensor that can reside on the host (CPU) or NPU.
Definition npu_data.py:29
An NPU-accessible pinned (physically contiguous) memory buffer.
A simple byte-sized buffer.
Definition type.py:219
Defines the core mode for NPU execution.
Definition type.py:244
Model load(str path, Optional[ModelConfig] model_config=None)
Single-step inference API.
Definition model.py:1043