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_argmax_with_bias
struct Struct_argmax_with_bias
Registers the sampler.argmax_with_bias graph op.
Takes, per row, the argmax of logits[:, position, :] + bias, with both
inputs upcast and added in float32 in registers: the same index as an
argmax of the materialized float32 sum. The DSpark Markov head runs it
once per draft position and reads the draft logits in place.
Implemented traits
Methods
execute
static def execute[dtype: DType, bias_dtype: DType, target: StaticString, _trace_name: StaticString](out_tokens: ManagedTensorSlice[IOSpec[_, _].Output, static_spec=out_tokens.static_spec], position: Scalar, logits: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=logits.static_spec], bias: ManagedTensorSlice[IOSpec[_, _].Input, static_spec=bias.static_spec], ctx: DeviceContext)