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 zero fill establishes the kernel’s handoff invariants (the 0.0f activation-scale sentinel, launch parity counters, and readiness flags — see docs/design_docs/monomoe_kernel.md §2/§4); afterwards the kernel self-maintains them, so allocate once and reuse the same tensor for every mono_moe() invocation.

Parameters:

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

Returns:

A zero-initialized 1-D uint8 tensor containing get_scratchpad_size_bytes() bytes on device.

Return type:

torch.Tensor