Medusa multi-head speculative prediction (Kamath Ch 8)

k heads produce k future tokens (0-BASED ids); with verify the run of accepted tokens is counted and reported as the estimate.

Usage

morie_kamath_medusa_heads(hidden_state, medusa_heads, k, verify = NULL)

Arguments

hidden_state

Hidden vector.

medusa_heads

Matrices or functions.

k

Heads used.

verify

Optional (position, token) -> logical.

Value

List with tokens, probabilities, accepted.

Examples

morie_kamath_medusa_heads(c(1, 2), list(diag(2), matrix(c(0,
    1, 1, 0), 2)), 2)
#> $tokens
#> [1] 1 0
#> 
#> $probabilities
#> [1] 0.7310586 0.7310586
#> 
#> $distributions
#> $distributions[[1]]
#> [1] 0.2689414 0.7310586
#> 
#> $distributions[[2]]
#> [1] 0.7310586 0.2689414
#> 
#> 
#> $n_heads_available
#> [1] 2
#> 
#> $n_heads_used
#> [1] 2
#> 
#> $accepted
#> NULL
#> 
#> $estimate
#> [1] 0.7310586
#> 
#> $n
#> [1] 2
#> 
#> $method
#> [1] "Medusa multi-head speculative prediction"
#>