flashinfer.gdn_decode.ugpu_state_pool

flashinfer.gdn_decode.ugpu_state_pool(pool_size: int, num_v_heads: int, head_size: int = 128, device: device | None = None) → Tensor

Allocate a bfloat16 GDN state pool for uGPU localization on Rubin.

Returns an uninitialized pool of shape [pool_size, HV, 2, V // 2, K] (V = K = head_size): element [s, h, t, r, k] is element [s, h, t * V // 2 + r, k] of the usual [pool_size, HV, V, K] pool. Half t of every state lives in locality domain (die) t’s memory. Pass it as initial_state to gated_delta_rule_decode_pretranspose(): on Rubin a T=1 call then runs half t of every state on die t’s SMs, so no state byte crosses the die-to-die link, whatever the batch order and slot indices. That is for CUDA-graph captured decode (make one eager call before capture); eager calls with B * HV < 8192 are host-bound and run both halves back to back on the current stream.

The pool comes from flashinfer.green_ctx.empty_per_locality_domain(), outside PyTorch’s caching allocator.

Parameters:

  • pool_size – Number of state slots.

  • num_v_heads – Number of value heads (HV).

  • head_size – K = V (128).

  • device – CUDA device with two locality domains (default: the current device).