flashinfer.fused_moe.alloc_scratchpad

flashinfer.fused_moe.alloc_scratchpad(device: device) → Tensor

Allocate a zero-initialized scratchpad on device for the kernel.

Returns a 1-D uint8 tensor sized to get_scratchpad_size_bytes(). The buffer is reusable across calls (the kernel self-maintains its barrier counters via the ping-pong reset discipline), so callers should allocate once and pass the same tensor to every mono_moe() invocation to avoid per-call allocation and the one-time counter zero-init on the C++ side.

Parameters:

device (torch.device) – The CUDA device on which to allocate the scratchpad buffer.

Returns:

A 1-D uint8 tensor of get_scratchpad_size_bytes() bytes, zero-initialised and placed on device.

Return type:

torch.Tensor