1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
//! # Hibachi
//!
//! A **hi**gh-performance **batch** tensor processing library for efficiently batching
//! autoregressive model inference across multiple concurrent requests.
//!
//! ## Overview
//!
//! This library provides a flexible and extensible framework for processing
//! autoregressive model inference requests in batches. It's designed to maximize
//! throughput when handling multiple concurrent generation requests by dynamically
//! batching them and efficiently managing computational resources.
//!
//! Key components include:
//!
//! - A tensor abstraction layer supporting various backends
//! - A batched inference engine for autoregressive models
//! - Utilities for tensor manipulation and sequence management
//! - Asynchronous streaming of generated outputs
//!
//! ## Architecture
//!
//! The library is built around several key abstractions:
//!
//! ### Assumptions
//! Regardless of backend used, hibachi reserves two dimensions with special meanings:
//! - The `0th` dimension is reserved as the batch dimension
//! - The `1st` dimension is reserved as the sequence dimension
//! - Tensors may fill in other dimensions
//!
//! Given that we implement over backends that may panic, as of now, Hibachi reserves the right to panic
//! on invalid tensor operations.
//!
//! ### Backend Traits
//!
//! The [`backend::Backend`] and [`backend::Unsqueezable`] traits define the interface that any tensor
//! implementation must satisfy to work with the library. This allows the core
//! batching logic to remain independent of the specific tensor implementation.
//!
//! ### Autoregressive Processing
//!
//! The [`autoregressive::Autoregressive`] trait defines the interface for models that generate
//! outputs sequentially, while the [`autoregressive::AutoregressiveBatcher`] trait encapsulates
//! the logic for efficiently batching multiple generation requests. We provide a
//! [`autoregressive::AutoregressiveBatchInference`] which implements the [`autoregressive::AutoregressiveBatcher`]
//! trait out of the box.
//!
//! ### Feedforward Processing
//!
//! The [`feedforward::Feedforward`] trait defines the interface for models that generate
//! outputs in a single pass, while the [`feedforward::FeedforwardBatcher`] trait encapsulates
//! the logic for efficiently batching multiple generation requests. We provide a
//! [`feedforward::FeedforwardBatchInference`] which implements the [`feedforward::FeedforwardBatcher`]
//! trait out of the box.
//!
//! ## Features
//!
//! - **autoregressive** - Enables autoregressive model batching functionality
//! - **feedforward** - Enables feedforward model batching functionality
//! - **candle** - Enables candle backend
//! - **burn** - Enables burn backend
//!
//! ## Implementation Details
//!
//! The library maintains internal state to track active sequences and efficiently
//! allocate batch slots. When sequences complete (by generating a stop token),
//! they are automatically removed from the batch to make room for waiting requests.
//!
//! Tensor operations are abstracted through the [`backend::Backend`] trait, allowing for
//! different tensor implementations to be used without changing the core batching logic.
//!
/// Constants for client reference
pub use constant;