flashinfer.green_ctx.split_device_green_ctx

flashinfer.green_ctx.split_device_green_ctx(dev: device, num_groups: int, min_count: int) → Tuple[List[Stream], List[CUdevResource]]

将设备划分为多个绿色上下文，返回每个组对应的流和 CUdevResource，以及剩余的 SM。绿色上下文允许在不同的 SM 分区上并发执行多个内核。

参数:

• dev – 要划分的设备。

• num_groups – 将设备划分为的组数。

• min_count – 每个组所需的最小 SM 数量，将根据对齐和粒度要求进行调整。

返回值:

与绿色上下文对应的 torch.Streams 对象列表。 resources: 与绿色上下文对应的 CUdevResource 对象列表。

返回值类型:

streams

示例

>>> from flashinfer.green_ctx import split_device_green_ctx
>>> import torch
>>> dev = torch.device("cuda:0")
>>> streams, resources = split_device_green_ctx(dev, 2, 16)
>>> print([r.sm.smCount for r in resources])
[16, 16, 100]
>>> with torch.cuda.stream(streams[0]):
...     x = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16)
...     y = torch.randn(8192, 8192, device=dev, dtype=torch.bfloat16)
...     z = x @ y
...     print(z.shape)
...
torch.Size([8192, 8192])

注意

返回的流和资源的长度为 num_groups + 1，其中最后一个是剩余的 SM。

以下示例显示了如何四舍五入 SM 计数以满足对齐和粒度要求：- 请求 7 个 SM → 分配 8 个 SM（四舍五入到最小值）- 请求 10 个 SM → 分配 16 个 SM（四舍五入到 8 的倍数）- 请求 16 个 SM → 分配 16 个 SM（无需四舍五入）- 请求 17 个 SM → 分配 24 个 SM（四舍五入到 8 的倍数）

抛出：

• RuntimeError – 请求的 SM 分配超过设备容量时

• num_groups * rounded_min_count > total_device_sms –