flashinfer.green_ctx.split_device_green_ctx_by_sm_count

flashinfer.green_ctx.split_device_green_ctx_by_sm_count(dev: device, sm_counts: List[int]) → Tuple[List[Stream], List[CUdevResource]]

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

参数:

• dev – 要划分的设备。

• sm_counts – 每个分区 SM 数量的列表。每个数量都将向上取整以满足最小和对齐要求。

返回值:

返回的流对象列表对应于绿色上下文。resources：与绿色上下文对应的 CUdevResource 对象列表。

返回值类型:

streams

抛出：

• RuntimeError – 如果请求的 SM 分配超过设备容量：- 当 sum(rounded_sm_counts) > total_device_sms - 当由于无效资源类型导致 CUDA 操作失败 - 当设备未正确初始化时

• ValueError – 如果 sm_counts 为空或包含无效值（例如，负值）。

示例

>>> from flashinfer.green_ctx import split_device_green_ctx_by_sm_count
>>> import torch
>>> dev = torch.device("cuda:0")
>>>
>>> # Create three partitions with specific SM counts
>>> streams, resources = split_device_green_ctx_by_sm_count(dev, [8, 16, 24])
>>> print([r.sm.smCount for r in resources])
[8, 16, 24, 84]  # Last value is remaining SMs
>>>
>>> # Execute kernels on different partitions
>>> with torch.cuda.stream(streams[0]):
...     x = torch.randn(4096, 4096, device=dev, dtype=torch.bfloat16)
...     y = torch.randn(4096, 4096, device=dev, dtype=torch.bfloat16)
...     z = x @ y
...     print(f"Partition 0 result: {z.shape}")
...
>>> with torch.cuda.stream(streams[1]):
...     # Different computation on partition 1
...     a = torch.randn(2048, 2048, device=dev, dtype=torch.bfloat16)
...     b = torch.randn(2048, 2048, device=dev, dtype=torch.bfloat16)
...     c = a @ b
...     print(f"Partition 1 result: {c.shape}")

注意

返回的流和资源的长度为 len(sm_counts) + 1，其中最后一个包含未分配的剩余 SM。

对于 Compute Capability 9.0+ 的 SM 计数对齐示例：- 请求 7 个 SM → 分配 8 个 SM（向上取整到最小值）- 请求 10 个 SM → 分配 16 个 SM（向上取整到 8 的倍数）- 请求 16 个 SM → 分配 16 个 SM（无需取整）- 请求 17 个 SM → 分配 24 个 SM（向上取整到 8 的倍数）

可以从返回资源的 .sm.smCount 字段中获取实际的 SM 计数。

有关更多详细信息，请参阅 CUDA 绿色上下文。