Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of same dense type. Since inputs are dense, the resulting tensor will b

Buffer associated with the batch metadata tensor that can be populated in preparation for passing a new input to the device.

Backwards-compatible way of querying the tensor model parallel rank from an ``mpu`` object. *Tensor* model parallelism means that tensors are

Traverse model and detect MoE layers. Returns list of MoELayerSpec.

Derive TP, dense-DP, EP, and EDP rank tables without process groups.

Return the number of processes along the given axis. For example: >>> X = ProcessTopology(axes=['x', 'y'], dims=[2,3])

Helper to determine the number of local heads of a given shard. Args: shard_rank (int): The rank of the shard. num_shards (i

Return a list of all parameters that would have LoRA for the module.

Buffer associated with the metadata of each sequence in the ragged batch. If the device Tensor is requested, the Tensor is truncated

Query whether the current micro-batch is at the boundary of gradient accumulation, and thus will trigger gradient reductions and

True if this process is in the last stage in the pipeline.

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

Utility method for sharding a QKV parameter. Both biases and weights are supported. It is assumed that the layout of the parameter is such th

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

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

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

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

Stores a fragment to memory

Fused SwiGLU activation: ``silu(gate) * up``. Args: gate: Gate projection output, any shape, float16/bfloat16/float32. up: Up pro

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