作者： 吃果冻不吐果冻皮 来源： 吃果冻不吐果冻皮

####**【点击】加入大模型技术交流群**

在实际的模型部署场景中，我们一般会先优化模型的性能，这也是最直接提升模型服务性能的方式。但如果从更全局方面考虑的话，除了模型的性能，整体的调度和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

• Continuous batching: https://zhuanlan.zhihu.com/p/676109470

• How continuous batching enables 23x throughput in LLM inference while reducing p50 latency: https://www.anyscale.com/blog/continuous-batching-llm-inference

• nougat模型: https://facebookresearch.github.io/nougat/

更多AI工具，参考Github-AiBard123，国内AiBard123