`core.inference.shards_spec`#

Parser and typed representation (:class:InferenceShardSpec) for the --inference-shards shard-layout string.

Module Contents#

Classes#

InferenceShardSpec

One inference shard’s parallelism – the canonical shard-layout type.

Functions#

`normalize_shard_specs`	Coerce the public shard-spec input (a spec string, a list of
`spec_declares_disaggregation`	Whether a shard spec tags any shard with a `role=` (prefill/decode).
`parse_inference_shards_spec`	Parse + validate the `--inference-shards` string.
`_finalize_and_validate`	Assert the shards partition the world exactly.

Data#

`VALID_INT_KEYS`
`VALID_ROLES`
`VALID_KEYS`

API#

core.inference.shards_spec.VALID_INT_KEYS#: (‘tp’, ‘pp’, ‘ep’, ‘expt_tp’, ‘dp’, ‘cp’)

core.inference.shards_spec.VALID_ROLES#: (‘prefill’, ‘decode’)

core.inference.shards_spec.VALID_KEYS#: ()

class core.inference.shards_spec.InferenceShardSpec#

One inference shard’s parallelism – the canonical shard-layout type.

Frozen and self-validating: expt_tp defaults to tp (resolved at construction), the expert decomposition is checked to tile within the shard, and role ("prefill" / "decode" for disaggregation, else None) is validated. A list of these, produced by :func:normalize_shard_specs, is what every inference-shard consumer (the PG builder, the disaggregation setup) operates on.

tp: int#: 1

pp: int#: 1

ep: int#: 1

dp: int#: 1

cp: int#: 1

expt_tp: Optional[int]#: None

role: Optional[str]#: None

__post_init__()#

property world_size: int#: Number of ranks this shard occupies (tp * pp * dp).

to_dict() → dict#: Plain-dict form (e.g. for serialization or external consumers).

core.inference.shards_spec.normalize_shard_specs( shards: Union[str, Sequence[core.inference.shards_spec.InferenceShardSpec], Sequence[dict]], world_size: int, ) → List[core.inference.shards_spec.InferenceShardSpec]#

Coerce the public shard-spec input (a spec string, a list of

Class:: InferenceShardSpec, or a list of raw dicts) into the validated list of :class:InferenceShardSpec the shard builders consume.

core.inference.shards_spec.spec_declares_disaggregation(spec_str: str) → bool#

Whether a shard spec tags any shard with a role= (prefill/decode).

A role tag is what marks the layout as a prefill->decode handoff rather than plain multi-shard / data-parallel inference. Cheap and world_size- free, so it can be checked at arg-validation time; full parsing + validation is :func:parse_inference_shards_spec.

core.inference.shards_spec.parse_inference_shards_spec( spec_str: str, world_size: int, ) → List[dict]#

Parse + validate the --inference-shards string.

Parameters:

spec_str – Raw CLI value, e.g. "tp=2,dp=1+tp=1,dp=2" or with disaggregation roles "tp=2,role=prefill+tp=1,role=decode".
world_size – Total number of ranks. Specs must partition it exactly (no idle ranks; see note below).

Returns:

List of :class:InferenceShardSpec, one per shard. Order matches the input (left-to-right corresponds to ascending rank_offset).

Raises:

AssertionError – on syntax errors, unknown keys, or a rank-count mismatch with world_size. Idle ranks are rejected to keep the partition explicit — any world-collective consumer must be able to enumerate every rank’s shard membership from the parsed list alone.
ValueError –
on an expert-grid mismatch within a shard (raised by

class:

InferenceShardSpec).

core.inference.shards_spec._finalize_and_validate( specs: List[core.inference.shards_spec.InferenceShardSpec], world_size: int, ) → List[core.inference.shards_spec.InferenceShardSpec]#

Assert the shards partition the world exactly.

Shared by the string parser and the object path (:func:normalize_shard_specs). Per-shard defaults and expert-grid validation live in

Class:: InferenceShardSpec; this only enforces the cross-shard total. Idle ranks are rejected so any world-collective consumer can enumerate every rank’s shard membership from the list alone.

core.inference.shards_spec#