For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
gumbel_sampling_fused_gpu
def gumbel_sampling_fused_gpu[dtype: DType, out_idx_type: DType, //, TemperatureLayoutType: TensorLayout = Layout[*(), *()], SeedLayoutType: TensorLayout = Layout[*(), *()], from_probs: Bool = False, TemperatureStorageType: TensorStorage = PointerStorage, SeedStorageType: TensorStorage = PointerStorage](ctx: DeviceContext, input: TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], out_idxs: TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type], temperature: Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]] = None, seed: Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]] = None)
Fused Gumbel sampling: applies Gumbel(0,1) noise and selects the argmax in a single GPU kernel launch (no intermediate noised-logits HBM buffer).
Mathematically equivalent to gumbel_sampling_gpu and produces bit-identical
results for the same seed, but saves one full [batch, vocab] HBM
round-trip by fusing noise generation and argmax.
With from_probs the input rows are unnormalized probabilities and the
draw is proportional to them; see _gumbel_argmax_fused_kernel.
Args:
- ctx (
DeviceContext): Device context for GPU operations. - input (
TileTensor[dtype, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): Input logits tensor [batch, vocab_size]. - out_idxs (
TileTensor[out_idx_type, Storage=out_idxs.Storage, address_space=out_idxs.address_space, linear_idx_type=out_idxs.linear_idx_type]): Output tensor for sampled indices [batch, 1]. - temperature (
Optional[TileTensor[DType.float32, TemperatureLayoutType, ImmutAnyOrigin, Storage=TemperatureStorageType]]): Optional per-token temperature scaling [batch]. - seed (
Optional[TileTensor[DType.uint64, SeedLayoutType, ImmutAnyOrigin, Storage=SeedStorageType]]): Optional per-token random seeds [batch] for reproducibility.