在实际的模型部署场景中，我们一般会先优化模型的性能，这也是最直接提升模型服务性能的方式。但如果从更全局方面考虑的话，除了模型的性能，整体的调度和pipeline优化对服务的性能影响也是很大。
比如LLM中提的很多的Continuous batching[1]，对整体LLM推理的性能影响就很大，这个不光光是提升kernel性能能够解决的问题。
这里总结下各种batching策略，以及各种batch策略对整体性能的影响，可能不够全面，也希望能够抛砖引玉，一起交流。
单batch
单batch就是不组batch，也就是一个或者一个sentence传过来，直接送入模型进行推理。
对于普通CV模型来说，我们的输入tensor大小可以是[1,3,512,512]，以NCHW维度举例子，这里的N是1，即batch=1。对于LLM来说，可能是一个input_ids，维度是[1,1]，比如：
input_ids
tensor([[ 0, 376, 1366, 338, 263, 3017, 775, 6160]], device=’cuda:0′)
input_ids.shape
torch.Size([1, 8])
这种情况比较简单，在搭建服务的时候不需要额外处理什么，单batch比较适合dymamic shape的场景，模型尺寸是否是dynamic也就是NCHW中的HW是否是可变的。如果你的模型的推理场景需要dynamic shape，那么一般无法组batch（不过可以用ragged batch，后续介绍），只能设置为batch=1，一般我们batch的时候HW都需要固定，比如[8,3,256,256]。
对于dynamic shape，不管是优化还是转模型都需要额外注意，比如tensorrt中转换dynamic shape需要设置动态范围：
polygraphy run model_sim.onnx –trt –onnxrt –fp16 \
–trt-min-shapes input:[1,64,64,3] \
–trt-opt-shapes input:[1,1024,1024,3] \
–trt-max-shapes input:[1,1920,1920,3]
dynamic shape的好处就是不需要padding了，避免了额外计算，适合那种请求shape变化特别剧烈的场景，不管是传输还是推理，少了无效的padding像素，自然变快了不少。
Static Batching
上述batch=1的情况虽然简单而且灵活，不过因为每次只能处理一张图，对GPU资源的利用率还是不如大batch，增大batch一般可以提升模型的FLOPS，同时也可以更多地利用tensor core，实际场景中表现一般就是这样：

centernet res50

[1,3,1024,1024] # trt inference 需要10ms
[8,3,1024,1024] # trt inference 需要50ms
很显然增大batch可以提升throughput，所以某些场景，对于时延（latency）要求没有那么高的时候，可以尝试大batch来提升吞吐。
静态batch比较简单，对于图像来说，我们可以在模型推理前，输入将图像cat到一起，比如[4,3,1024,1024]，送给模型推理，这时的batch=4。
对于非图像场景也一样，将输入tensor合并后就相当利用了batch。不过比较尴尬的是，LLM场景因为decode阶段batch中每个case最终结束时输出的长度不一样，所以会有有些case吐完字儿了，有些还在继续吐，早吐完的需要等还没吐完的。这种情况不管是速度还是GPU利用率都是比较低的：

LLM-static 来自https://www.anyscale.com/blog/continuous-batching-llm-inference
不过当然有解决方案：
https://www.anyscale.com/blog/continuous-batching-llm-inference
Dynamic Batching
除了在客户端直接送给模型大batch数据，也可以在triton服务端组batch，
Dynamic batching, i.e., server-side batching of incoming queries, significantly improves both latency and performance.
动态组batch在图像推理场景很常用，也可以配合static batch使用。
比如你的模型支持最大batch是16，然后客户端可以每次发送batch=x过来，服务端在收到这些请求后，可以根据规则去组batch，组成更大的batch送给模型：

image
当模型实际拿到的tensor batch越大，模型的性能就越强，找个例子压测看下：
$ perf_analyzer -m inception_graphdef –percentile=95 –concurrency-range 1:8
…
Inferences/Second vs. Client p95 Batch Latency
Concurrency: 1, throughput: 66.8 infer/sec, latency 19785 usec
Concurrency: 2, throughput: 80.8 infer/sec, latency 30732 usec
Concurrency: 3, throughput: 118 infer/sec, latency 32968 usec
Concurrency: 4, throughput: 165.2 infer/sec, latency 32974 usec
Concurrency: 5, throughput: 194.4 infer/sec, latency 33035 usec
Concurrency: 6, throughput: 217.6 infer/sec, latency 34258 usec
Concurrency: 7, throughput: 249.8 infer/sec, latency 34522 usec
Concurrency: 8, throughput: 272 infer/sec, latency 35988 usec
随着请求并行度的增加，triton服务端模型组的batch越来越大，服务的QPS也随之增加，不过时延也会有所增加，这个时候就要trade off了。
Continuing Batching
Continuing Batching，也可以叫做inflight batching或者Iteration batching。不同于static batching，当一个输入的生成结束了，就将新的输入插进来，所以batch size是动态的。下图第三个输入先生成完，新的输入S5就插入进来了，一直到输出最长的S2的输出结束的时候，batch size由4变成了7，任意case跑完就能返回：

Continuing Batching
当然实际情况更为复杂，这里不进行详细讨论。另外Continuing Batching也有人称为dynamic batching，和上述不是一回事儿哈
之前static batching的时候也提过，对于LLM场景，假如batch中某个case已经提前结束了，但是其余batch还在decode，那么这个case只能等其他case而无法直接返回，导致某些请求时延变长了。
另外，假如遇到特别长的case，为了等这个case跑完，其他case都需要等它，其他请求进来也得等他，等来等去待请求队列就长了，服务也就崩了。
所以Continuing Batching这个调度策略很重要，相比模型本身的kernel性能，调度对整体性能的影响是很大的，夸张点，拿23倍的吞吐量轻而易举：
How continuous batching enables 23x throughput in LLM inference while reducing p50 latency[2]
很多优秀的LLM推理框架都用到了这个技术，几乎是必用的：TensorRT-LLM、vLLM、Imdeploy等。
Ragged Batching
上文我们提到Triton有动态批处理功能，它可以将多个相同模型执行的请求合并以提供更大的吞吐量。
默认情况下，只有在每个请求中的输入具有相同shape时，才能进行动态批处理。如果我们想同时用上dynamic batching和dynamic shape，一般则需要客户端将请求中的输入tensor填充到相同的形状。
举个例子，比如输入[1,3,768,932]和[1,3,1024,768]，合并为batch的时候需要将两个tensor都padding到同样尺寸，比如[1,3,1024,1024]，然后再cat起来为[2,3,1024,1024]才可以传给服务端模型；又或者客户端分别发了两个padding后的[1,3,1024,1024]过来，服务端自动组batch为[2,3,1024,1024]送给模型。
咱们之前也说过，padding会带来不必要的传输和计算量，所以我们可以稍微修改下调度逻辑从而实现dynamic shape + dynamic batching，也就是所谓的ragged batching。
在triton中，ragged batching是一种避免显式padding的功能，它允许用户指定哪些输入不需要进行形状检查。用户可以通过在模型配置中设置allow_ragged_batch字段来指定这样的输入（不规则输入）：
input [
{
name: “input0”
data_type: TYPE_FP32
dims: [ 16 ]
allow_ragged_batch: true
}
]
举个triton官方的例子。
如果我们有一个接受变长输入tensor INPUT 的模型，INPUT 的形状为 [ -1, -1 ]。第一个维度是 batch 维度，第二个维度是变长内容。当客户端发送三个形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ] 的请求时，为了利用动态 batching，最直接的实现方法是期望 INPUT 的形状为 [ -1, -1 ] 并假设所有输入都被填充到相同的长度，这样所有请求都变成形状为 [ 1, 5 ]，因此 Triton 可以将它们 batch 并作为一个 [ 3, 5 ] 张量发送到模型中。在这种情况下，会有padding的开销和对padding进行额外模型计算的开销。以下是输入配置：
max_batch_size: 16
input [
{
name: “INPUT”
data_type: TYPE_FP32
dims: [ -1 ]
}
]
那如果使用 Triton 的 ragged batching，模型将被实现为期望 INPUT 形状为 [ -1 ]，并且有一个额外的 batch 输入 INDEX，形状为 [ -1 ]，模型应该使用它来解释 INPUT 中的 batch 元素。对于这种模型，客户端请求不需要padding，可以按原样发送（形状为 [ 1, 3 ]、[ 1, 4 ]、[ 1, 5 ]）。上述后端将把输入 batch 成一个形状为 [ 12 ] 的张量，其中包含 3 + 4 + 5 个请求的连接。Triton 还创建了一个形状为 [ 3 ] 的 batch 输入张量，值为 [ 3, 7, 12 ]，它给出了每个 batch 元素在输入张量中的结束位置，即索引。以下是输入配置：
max_batch_size: 16
input [
{
name: “INPUT”
data_type: TYPE_FP32
dims: [ -1 ]
allow_ragged_batch: true
}
]
batch_input [
{
kind: BATCH_ACCUMULATED_ELEMENT_COUNT
target_name: “INDEX”
data_type: TYPE_FP32
source_input: “INPUT”
}
]
使用ragged batch需要实际的模型支持才行（如何处理ragged batch和index）。实际LLM推理中用到的比较多，比如TensorRT-LLM中kernel实现的时候已经考虑到了这种情况：
// TensorRT-LLM/cpp/tensorrt_llm/kernels/gptKernels.cu
// This kernel also computes the padding offsets: Given the index (idx) of a token in a ragged tensor,
// we need the index of the token in the corresponding tensor with padding. We compute an array
// of numTokens elements, called the paddingOffsets, such that the position in the padded tensor
// of the token “idx” in the ragged tensor is given by idx + paddingOffset[idx].
//
// That kernel uses a grid of batchSize blocks.
在图像中，ragged batch一般用不上。不过我们也可以自行设计一个使用ragged batch的策略，简单来说，我们可以把padding操作放在服务端（只是举个栗子，实际场景padding最好在客户端做）。
当客户端分别请求[1,3,768,932]和[1,3,1024,768]这两个shape的时候，我们可以在服务端将这两个请求组batch，然后进行某种预处理（padding或者resize），将input处理好再传入模型，需要我们设计模块处理这种ragged input并且利用起来：
def execute(self, requests):

    responses = []
    num_request = len(requests)
    input_tensors = []
    time_start = time.time()

    for idx, request in enumerate(requests):

        input = pb_utils.get_input_tensor_by_name(request, "input")
        input = input.as_numpy()
        input = input.squeeze(0)  # NHWC -> HWC
        input = Image.fromarray(input)
        # 这里可以将不同shape的input 搞成一个shape
        input_tensor = self.prepare_input(input).unsqueeze(0)
        input_tensors.append(input_tensor)

    input_tensor = torch.cat(input_tensors, dim=0).to("cuda")    
    self.logger.log_info("input_tensor: {}".format(input_tensor.shape))

    output_texts = self.model.run(input_tensor, num_request, 4096)

    for output_text in output_texts:

        output = pb_utils.Tensor(
            'output',
            np.array(output_text).astype(self.output_dtype))

        inference_response = pb_utils.InferenceResponse(
            output_tensors=[output]
        )

        responses.append(inference_response)

    # You should return a list of pb_utils.InferenceResponse. Length
    # of this list must match the length of `requests` list.
    return responses

使用场景还有很多，需要我们自行探索了。
Custom batching
自定义的一种batch策略，一般就是有具体使用场景才会有目的性的去设计。
举个实际的例子，比如LLM中多模态推理中，有个nougat模型[3]，会识别一幅图中所有的单词并且一个一个吐出来，这个模型是由两部分组成的：
encoder（普通的cv模型，传入图像传出特征）
decoder（可以理解为和llama一样的decoder模型，带有cross attention结构）
当decoder暂时不支持inflight batching的时候，我们只能使用static batching，但是显然在组batching的时候，字儿少的要等字儿多的都吐完才能一起返回。
怎么办，我们可以利用nougat的特性，提前用一个检测模型检测当前中字儿的数量，然后根据数量将数量相近的请求组成batch传入模型，并且将那种数量特别多的请求单独推理，这样就可以避免一些等待的情况。
这种根据数量进行组batch就是一种custom batching策略。
参考
https://github.com/triton-inference-server/server/blob/main/docs/user_guide/ragged_batching.md
https://github.com/triton-inference-server/server/issues/6458
https://github.com/NVIDIA/TensorRT-LLM/discussions/1469
参考资料
[1]
Continuous batching: https://zhuanlan.zhihu.com/p/676109470
[2]
How continuous batching enables 23x throughput in LLM inference while reducing p50 latency: https://www.anyscale.com/blog/continuous-batching-llm-inference
[3]
nougat模型: https://facebookresearch.github.io/nougat/