RBLN Sampler¶

Experimental feature

Some parts of this feature are experimental. Full support is coming soon.

Overview¶

Sampling selects the next token from a model's output distribution at each generation step. For LLMs with large vocabularies, this can account for a significant portion of inference time. Top-k and top-p sampling can add filtering and sorting overhead when selecting candidate tokens.

To reduce this overhead, RBLN vLLM supports custom operators (custom ops) for sampling on RBLN NPUs, which are used by RBLNSampler. In particular, the top-k / top-p sampling custom op implements the FlashInfer sampling algorithm, using rejection sampling to select tokens without sorting the entire vocabulary.

Sampling support¶

Greedy and top-k / top-p sampling are stable. All remaining sampling methods are experimental.

Configuration¶

Sampler selection¶

RBLNSampler is enabled by default. To disable it and use the upstream vLLM sampler, pass --no-rbln-use-custom-sampler.

Warm-up¶

Sampling graphs are compiled during initialization by default (VLLM_RBLN_ENABLE_WARM_UP=1). Setting it to 0 skips warm-up for all graphs, including the sampler, and can slow down initial inference.

Caching¶

Sampling graphs are not stored in the compile cache.

Speculative sampling¶

Sampling is also supported with speculative decoding. See Speculative Decoding for the supported methods and configuration.