flashinfer.green_ctx.split_device_green_ctx_by_locality_domain

flashinfer.green_ctx.split_device_green_ctx_by_locality_domain(dev: device, backfill: bool = True) → Tuple[List[Stream], List[CUdevResource]]

Create one green context per locality domain (uGPU) of dev and return its stream and CUdevResource, in domain order. A kernel launched on stream i runs only on domain i’s SMs, so together with memory placed in that domain (empty_per_locality_domain()) its DRAM traffic stays on the die.

Parameters:

  • dev – The device to split.

  • backfill – If True (default), each group gets num_SMs // num_domains SMs and the driver completes a domain’s group with SMs that belong to no domain (on Rubin VR200: 100 native SMs + 6 of the 12 remaining ones per die, so all 212 SMs work). If False, each group holds its domain’s native SMs only.

Returns:

One torch.Stream per locality domain. resources: The matching CUdevResource objects (.sm.smCount is the SM count).

Return type:

streams

Raises:

RuntimeError – if the device has fewer than two locality domains.