mixle.engines.heterogeneous module

Precision-aware planning for distributed EM across heterogeneous compute.

Large worker pools are rarely uniform: some workers may have GPU tensor cores, while others are CPU-only or accuracy-oriented. This module chooses, per worker, how many E-step rows to assign and which precision band to run. The selected precision is the fastest supported band that still satisfies the requested error budget. The plan also sizes the k-way reduction depth so fixed-size sufficient-statistic payloads fold in O(log W) instead of a single-root fan-in.

This module is the pure-Python planning layer. Spark, MPI, or other distributed dispatchers consume the returned plan from the inference layer.

class Worker(name, device, precisions, base_throughput=1.0)[source]

Bases: object

A compute worker: its device, the precisions it can run (any order), and a base throughput.

Parameters:
class WorkerAssignment(name, rows, precision, effective_throughput)[source]

Bases: object

One worker’s row allocation, precision, and effective throughput.

Parameters:
class HeterogeneousPlan(assignments, reduce_depth)[source]

Bases: object

Assignments and reduction depth for heterogeneous execution.

Parameters:
  • assignments (tuple[WorkerAssignment, ...])

  • reduce_depth (int)

total_rows()[source]

Return total rows assigned across workers.

Return type:

int

plan_heterogeneous(workers, n_rows, allowed_precisions=('fp8', 'bfloat16', 'float16', 'float32', 'float64', 'dd'), target_rel_error=None, op_count=1000)[source]

Assign rows + a precision band to each worker, balanced by precision-adjusted throughput.

Each worker runs the fastest precision its hardware supports that stays within target_rel_error (None = no accuracy constraint); rows are split proportionally to the resulting throughput so all workers finish together. reduce_depth is the k-way tree depth for folding the sufficient-statistic payloads (~ceil(log2(W)/2)), avoiding the single-root fan-in.

Parameters:
  • workers (list[Worker])

  • n_rows (int)

  • allowed_precisions (tuple[str, ...])

  • target_rel_error (float | None)

  • op_count (int)

Return type:

HeterogeneousPlan