mstar.engine.resources.sampler.config

What a model declares about sampling: its spec, its per-request config, its step.

Kept free of the resource and its Triton kernels so a submodule can declare a step without pulling them in behind it.

Classes

SamplerSpec(resource_key, nodes, vocab_size)
SamplerStep(segments, ...] | None = None, ...)
SamplingReqConfig([temperature, top_k, ...])

class mstar.engine.resources.sampler.config.SamplerSpec(resource_key: str, nodes: set[str], vocab_size: int | None, enable_repetion_penalty: bool = True)

Bases: NodeResourceSpec

Parameters:

• resource_key (str)

• nodes (set[str])

• vocab_size (int | None)

• enable_repetion_penalty (bool)

enable_repetion_penalty: bool = True

property resource_class: type[Resource]

What builds this spec. Imported inside the property, so declaring a resource stays free of the manager and its kernels.

The builder, not necessarily the class built: an attention spec names AttentionManager, whose build picks a backend subclass.

vocab_size: int | None

class mstar.engine.resources.sampler.config.SamplerStep(segments: tuple[mstar.engine.resources.step.Segment, ...] | None=None, apply_penalty: bool = True, prefill_tracked_tokens: dict[str, torch.Tensor]=<factory>)

Bases: ResourceStep

Parameters:

• segments (tuple[Segment, ...] | None)

• apply_penalty (bool)

• prefill_tracked_tokens (dict[str, Tensor])

apply_penalty: bool = True

prefill_tracked_tokens: dict[str, Tensor]

class mstar.engine.resources.sampler.config.SamplingReqConfig(temperature: float = 0.6, top_k: int = 0, top_p: float = 1, ignore_eos: bool = False, repetition_penalty: float = 1, _seed: int = 0)

Bases: ResourceReqConfig

Parameters:

• temperature (float)

• top_k (int)

• top_p (float)

• ignore_eos (bool)

• repetition_penalty (float)

• _seed (int)

apply_conductor_config(seed=0, **kwargs)

Parameters:

seed (int)

ignore_eos: bool = False

repetition_penalty: float = 1

property seed

temperature: float = 0.6

top_k: int = 0

top_p: float = 1