flashinfer.green_ctx

This module provides utilities for creating CUDA green contexts.

Green context creation

split_device_green_ctx(dev, num_groups, ...)	Split the device into multiple green contexts, return the corresponding streams and CUdevResource for each group and the remaining SMs.
split_device_green_ctx_by_sm_count(dev, ...)	Split the device into multiple green contexts, each with a fixed number of SMs, return the corresponding streams and CUdevResource for each group and the remaining SMs.

uGPU localization (multi-die GPUs)

On GPUs built from several dies (Rubin), each die is a locality domain with its own SMs and memory. A kernel on a domain’s green-context stream that reads memory placed in that domain keeps its DRAM traffic off the die-to-die link.

get_locality_domain_count(dev)	Number of locality domains (uGPUs, one per die) of dev.
split_device_green_ctx_by_locality_domain(dev)	Create one green context per locality domain (uGPU) of dev and return its stream and CUdevResource, in domain order.
empty_per_locality_domain(part_shape, dtype, dev)	Allocate an uninitialized tensor of shape [num_domains, *part_shape] whose part i lives in locality domain i's memory (one virtual address range over per-domain physical memory).