Extract tp_plan from HuggingFace model. Prefer base_model_tp_plan (from model config) over _tp_plan (runtime attribute) because _tp_plan ofte

Returns attention qkv and dense parameters weight: (3*hidden, hidden) and (hidden, hidden) bias: (3*hidden) and (hidden)

Shutdown the worker process gracefully. Sends termination signal to worker and waits for clean shutdown. If the process does

Load a tiny LlamaForCausalLM, tag inputs with AutoSP keys, mark the sequence dimension dynamic, and capture the torch-fx GraphModule via a cu

Release GPU memory consumed by offloaded model parameters.

Build a config dict. force_fp32 keeps the engine in fp32: the frozen-param non-reentrant CheckpointError only reproduces in fp32 (bf16 recomp

Calculates the learning rate at batch index. This function treats `self.last_batch_iteration` as the last batch index.

Return a registered AutoEP preset adapter by name.

Partition routed assignments across TP peers by stable per-expert ordinal. Each peer keeps only ``assignment_index % tp_size == tp_rank`` of the

use dynamic programming solve `The Linear Partition Problem`. see https://www8.cs.umu.se/kurser/TDBAfl/VT06/algorithms/BOOK/BOOK2/NODE45.HTM

Determine the number of tokens and KV blocks to reserve for a given request. Given a UID (this UID may not be recognized by the model

Update the model parameters in a single subgroup (by index).

Stores a fragment to memory

Given a QKV tensor and a KV cache, validate that the cache contains the correct values.

Helper function for testing blocked flash attention. This implementation is based on the implemnentation in ``unit.inference.kernels.ragged_o

Linear layer plus an fp32 buffer that mimics the rotary inv_freq.

Adds a pointer offset in units of Element

Allocate a list of blocks from the associated KV-caches. This will return `num_blocks` blocks from the KV-cache if they are available

input tensor t is of shape [seq_length, ..., dim] rotary positional embeding tensor freqs is of shape [seq_length, ..., dim] check https:

Classify fused gate/up expert weights from raw tensor shapes.

Clear the ragged batch. This will reset the number of tokens and sequences to 0.

Clears the predicate set efficiently

Collect the model-level conversion schema for AutoTP universal checkpoints. The returned `UNIVERSAL_CHECKPOINT_INFO` is intentionally limited to

Return size in bytes of the given dtype.

Clears the predicate set efficiently

Gets the mask

Gets the mask

Given a dimension to shard with size dim_size, return the start and end indices of the slice that belong to the given rank. The typical

Create the inference logger. NOTE: Logging is not cost free. On a 3960X, there is a cost of about 6 us per call to a no-op logger, so this sh

Returns mlp intermediate and output weight: (intermediate, hidden) and (hidden, intermediate) bias: (intermediate) and (hidde

r""" Context manager to disable gradient reduction during backward pass. This context manager has the following effects on oth

Replace a single MoE module with AutoEPMoELayer in-place on the model.

Sets the predicate mask, overriding value stored in predicate iterator

Sets the predicate mask, overriding value stored in predicate iterator

Utility method for getting the size of the intermediate dimension of a sharded MLP. Args: intermediate_size (int): The size of the i

Compare two tensors or non-tensor numbers for their closeness. Add msg=blah to add an additional comment to when assert fails. For defa

Helper to create a subclass of ParametrizedList with the desired `count_attr`. In this manner, we can annotate the type of a Parameter depen

Initialize the Sparsity Pattern Config. For usage example please see, TODO DeepSpeed Sparse Transformer Tutorial Arguments:

Helper function for testing blocked flash attention. Used to enable parametrize to only set up a subset of parameters before being passed to

Ensure that numerator is divisible by the denominator.