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 struct
Struct_fused_token_sampling_with_dist
struct Struct_fused_token_sampling_with_dist
Registers the sampler.fused_token_sampling_with_dist graph op.
Samples one token per row under joint top-k/top-p with temperature, and
also returns the masked, renormalized distribution it drew from.
Speculative decoding subtracts that distribution to build its rejection
residual, and reads the sampled token's probability out of it. Op arity
is fixed, so this is a second registration over the same kernel as
sampler.fused_token_sampling rather than an optional output on it.
Implemented traits
Methods
execute
static def execute[dtype: DType, dist_dtype: DType, target: StringSpan[ImmStaticOrigin], _trace_name: StringSpan[ImmStaticOrigin]](out_tokens: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_tokens.static_spec], out_dist: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_dist.static_spec], K: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=K.static_spec], max_k: Scalar, temperature: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=temperature.static_spec], top_p: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=top_p.static_spec], seed: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=seed.static_spec], input: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=input.static_spec], ctx: DeviceContext)