1. 简介
面向对象设计（OOP）是传统后台开发领域非常主流的设计思想，但在大模型推理领域，极致的性能往往藏在 DOD（面向数据）的细节里。 本文通过对 mini-sglang 项目 _make_2d_indices 函数的迭代优化，展示了如何通过重构内存布局，消除对象转换开销，实现性能飞跃。核心结论：架构层用 OOP，计算层用 DOD。
注：本文仅用于学习探索，实际上 _make_2d_indices 函数当前的实现已经能满足通常的 batch size 场景，没有改造的诉求。
2. 场景需求
在做批量推理时，我们可能会遇到这样的需求：有一批二维数组，每一行代表一个请求，要求读取每一行指定区间的值，譬如：
有几种思路：
- 最笨的方法是采用 for 循环逐行读取，既没有使用 GPU，又发起了多次 IO 读取，性能较差，通常没人这么写。
- 好一点的做法是用 CPU 先建平坦索引（一维索引），然后发起一次 IO 读取（当前 mini-sglang
_make_2d_indices的实现）。 - 更好的做法是既考虑到用 GPU 构建平坦索引，又只发起一次 IO 读取（后文初步优化的实现）。
- 最好的做法：系统的考虑面向数据的设计：从输入参数到内部实现再到结果应用，都采用面向数据的设计。
下面介绍后三种实现。
3. 面向对象的实现
这是当前 mini-sglang _make_2d_indices 函数的实现，用于将二维数组的索引转换为一维索引，以便一次性读取关联数组。
代码：
def _make_2d_indices(table_2d: torch.Tensor, ranges: List[Tuple[int, int, int]]) -> torch.Tensor:
assert table_2d.dim() == 2 and table_2d.is_contiguous()
STRIDE = table_2d.stride(0)
needed_size = sum(end - begin for _, begin, end in ranges)
indices_host = torch.empty(needed_size, dtype=torch.int32, pin_memory=True)
offset = 0
for entry, begin, end in ranges:
length = end - begin
offset += length
torch.arange(
begin + entry * STRIDE,
end + entry * STRIDE,
dtype=torch.int32,
out=indices_host[offset - length : offset],
)
return indices_host.to(table_2d.device, non_blocking=True)
......
load_indices = _make_2d_indices(
self.token_pool, [(r.table_idx, r.cached_len, r.device_len) for r in batch.padded_reqs]
)
self.page_table.view(-1)[load_indices] = batch.out_loc
该实现虽用到了并行读取，降低了 IO 耗时，但核心逻辑仍依赖 CPU 循环逐个处理请求对象，保留了面向对象设计中‘以单个请求（对象）为单位处理’的思维模式，未能充分利用 GPU 的并行计算特性。随着 ranges 长度增加，耗时呈线性增长，无法利用 GPU 的并发优势。此外，indices_host 先在 CPU 创建再迁移到 GPU，也引入了额外的 CPU/GPU 数据传输开销，进一步降低了性能。
4. 初步优化: 批量计算
_make_2d_indices 函数可以通过 GPU 批量计算和掩码过滤做优化：
- 核心逻辑：构建一个
[num_ranges, max_range_length]的矩阵，通过广播机制一次性计算所有候选索引。 - 掩码过滤：利用长度掩码对填充位置进行裁剪，通过
indices_2d[mask]得到一维索引。
def _make_2d_indices_gpu(
table_2d: torch.Tensor,
ranges: List[Tuple[int, int, int]] # (entry, begin, end)
) -> torch.Tensor:
if not ranges:
return torch.empty(0, dtype=torch.int32, device=table_2d.device)
device = table_2d.device
stride = table_2d.stride(0)
# Convert to tensors
entries = torch.tensor([r[0] for r in ranges], dtype=torch.int32, device=device)
begins = torch.tensor([r[1] for r in ranges], dtype=torch.int32, device=device)
ends = torch.tensor([r[2] for r in ranges], dtype=torch.int32, device=device)
# Calculate max range length
lengths = ends - begins
max_length = lengths.max().item()
# Calculate 1D indices: (entry * stride + begin + position)
# Use broadcasting: [num_ranges, 1] + [max_length] -> [num_ranges, max_length]
positions = torch.arange(max_length, dtype=torch.int32, device=device) # [max_length]
indices_2d = entries.view(-1, 1) * stride + begins.view(-1, 1) + positions # [num_ranges, max_length]
# Create mask for valid positions
mask = positions < lengths.unsqueeze(1)
# Filter and return
return indices_2d[mask]
这版实现，对于 256 batch size 有 4 倍速的提升。它的核心优化是将 CPU 循环的 “逐一遍历” 转为 GPU 的 “广播批量计算”，消除了 CPU 循环的线性耗时，但仍需将 Python 列表（对象化的 ranges）转换为张量，存在少量数据转换开销。
5. 极致优化: 面向数据的设计
如果要达到极致性能，需要去掉“请求对象化”的结构，进行面向数据的设计（Data-Oriented Design, DOD）：
- 重构数据结构：将
batch.padded_reqs从“对象列表”改造为“张量集合”。将table_id、cached_len、device_len等属性直接存储为 GPU 上的连续张量（如table_id_array）。 - 符合批量计算/并行计算原则：这种改造使整个推理引擎从“以请求为中心（Request-centric）”转向“以算子为中心（Operator-centric）”，是实现极致吞吐量的关键。
最终代码：
def method6_batch_read_filter(
token_pool: torch.Tensor,
pool_indices: torch.Tensor,
begin_lens: torch.Tensor,
end_lens: torch.Tensor,
) -> torch.Tensor:
max_range_len = (end_lens - begin_lens).max().item()
# Create position index matrix using broadcasting: [Num_reqs, 1] + [max_range_len] -> [num_reqs, max_range_len]
positions = torch.arange(max_range_len, dtype=torch.int64, device=token_pool.device) # [max_range_len]
position_indices = positions + begin_lens.unsqueeze(1) # [num_reqs, max_range_len]
# Create mask for valid positions using broadcasting
mask = position_indices < end_lens.unsqueeze(1) # [num_reqs, max_range_len]
# Batch indexing
batch_values = token_pool[pool_indices.unsqueeze(1), position_indices]
# Filter with mask
return batch_values[mask]
这版实现，对于 256 batch size 有 8 倍速的提升。这里彻底消除了 “对象列表→张量” 的转换开销，所有数据均以连续张量的形式存储在 GPU 上，完全适配 GPU 的内存访问模式和并行计算模型。
除了这三种实现之外，还有其他小调整的测试，我共做了 6 个评测函数，细节见：https://github.com/cswuyg/tools/blob/main/benchmark_data_oriented_design.py
6. 优化总结
（1）核心差异：OOP vs DOD
面向对象设计（OOP）：以“请求对象”为中心，内存布局离散，优势是可读性、可扩展性强，适合业务逻辑层；
面向数据的设计（DOD）：以“连续张量”为中心，内存布局紧凑，优势是最大化GPU并行效率和内存带宽，适合高性能计算层。
（2）优化路径
原始OOP思维（CPU循环+数据传输）→ GPU批量计算（消除CPU循环）→ DOD（消除对象→张量转换），性能从基线提升至4倍→8倍。
（3）实践原则：
“分层设计”——架构层/业务层用 OOP 保证可维护性，计算层用 DOD 保证性能，兼顾可读性与高性能。
两者融合时，以 OOP 为主，再外挂 DOD 的 Tensor 结构。
7. 附：真实的应用
面向数据的设计，我在给 SGLang 项目添加 Beam Search 功能时有所应用。在涉及一批数据的计算时，使用面向对象封装管理，同时外挂维护一些偏平的 Tensor 结构用于批量运算加速。相关代码：https://github.com/sgl-project/sglang/pull/15645/ BeamSearchList 的 Tensor 成员。
浙公网安备 33010602011771号