Testing Guide¶

mlxcel uses a layered testing strategy: inline unit tests for pure logic, inline _tests.rs modules for module-level behavior, integration tests in tests/ for cross-module correctness, and distributed end-to-end tests for the full TP, PP, and disaggregated inference stacks.

Table of Contents¶

• Unit Tests
• Integration Tests
• Distributed End-to-End Tests
• Benchmarking
• Running Tests
• Test Organization Conventions

Unit Tests¶

Inline #[cfg(test)] Modules¶

Small, focused tests live directly in the module they test, guarded by #[cfg(test)]. These are the most common kind of test in the codebase.

// src/distributed/tensor_parallel/shard_strategy.rs
#[cfg(test)]
#[path = "shard_strategy_tests.rs"]
mod tests;

The #[path = ...] pattern keeps the test file as a sibling of the source file with an explicit _tests.rs suffix, making test discovery easy:

src/distributed/tensor_parallel/
  shard_strategy.rs         # implementation
  shard_strategy_tests.rs   # tests for shard_strategy.rs
  collective.rs
  collective_tests.rs
  ...

This pattern is used consistently throughout the codebase. Every non-trivial module has a companion _tests.rs file.

Example Unit Test¶

// shard_strategy_tests.rs
use super::*;

#[test]
fn test_column_parallel_shard_plan() {
    let plan = LayerShardPlan {
        weight_pattern: "model.layers.{}.self_attn.q_proj.weight".to_string(),
        strategy: ShardStrategy::ColumnParallel,
        shard_axis: 0,
        comm_pattern: CommPattern::AllGather,
    };
    assert_eq!(plan.strategy, ShardStrategy::ColumnParallel);
    assert_eq!(plan.comm_pattern, CommPattern::AllGather);
}

Where to Add Tests¶

Integration Tests¶

Integration tests are in the tests/ directory. They test cross-module behavior and do not require actual model weights for the distributed tests.

Test Files¶

What the Distributed Tests Cover¶

The distributed E2E tests exercise the full stack without real hardware by using in-process simulations:

TP E2E (tp_e2e.rs): - Attention head assignment for GQA-aware sharding - FFN intermediate dimension coverage across ranks - MoE expert assignment and coverage verification - Shard plan generation for standard Transformer, MoE, and SSM architectures - Collective operation correctness (all-reduce, all-gather, reduce-scatter) - Scheduler/executor lockstep simulation with in-process message passing - Scaling analysis and crossover benchmarks

PP E2E (pipeline_e2e.rs): - Layer partition correctness (auto and manual) - GPipe schedule drive-to-completion - Micro-batch splitting and size recommendations - Pipeline coordinator request lifecycle - Chunked prefill pipeline - Stage health tracking - Per-stage cache management - Scaling benchmarks

Disaggregated E2E (disaggregated_e2e.rs): - Prefill scheduler admission control and priority ordering - Decode scheduler ingestion queue and batch management - Request router lifecycle across prefill → decode handoff - Load balancing strategy correctness - Stream bridge state transitions and token ordering - KV cache serialization / deserialization round-trips - Transfer strategy selection - End-to-end latency benchmarks and crossover analysis

Distributed End-to-End Tests¶

The distributed E2E tests use in-process simulations that do not require real GPUs or network connections. This means they run in CI on any machine.

In-Process Collective Simulation¶

Collective operations are tested by creating simulated rank groups connected by std::sync::mpsc channels:

// Create loopback collective groups for 4 ranks
fn create_loopback_groups(world_size: usize) -> Vec<CollectiveGroup> {
    // Shared mailboxes: mailbox[dest_rank] holds data sent to that rank
    let mailboxes: Vec<Mutex<Option<Vec<u8>>>> =
        (0..world_size).map(|_| Mutex::new(None)).collect();
    let mailboxes = Arc::new(mailboxes);

    (0..world_size).map(|rank| {
        let mb = Arc::clone(&mailboxes);
        let config = CollectiveConfig { rank, world_size, chunk_size: 1024 };
        CollectiveGroup::new(config, Arc::new(move |dest, data| {
            *mb[dest].lock().unwrap() = Some(data);
            // Poll for incoming data...
        })).unwrap()
    }).collect()
}

Running E2E Tests¶

# All tests including distributed E2E
cargo test -- --test-threads=1

# Specific test file
cargo test --test tp_e2e -- --test-threads=1
cargo test --test pipeline_e2e -- --test-threads=1
cargo test --test disaggregated_e2e -- --test-threads=1

# With verbose output
cargo test --test tp_e2e -- --nocapture --test-threads=1

Note: --test-threads=1 is required because many tests share global state (e.g., model registry singletons, file system fixtures).

Benchmarking¶

mlxcel includes internal benchmark suites for distributed components, runnable from the Makefile or directly with cargo test.

Makefile Benchmark Targets¶

# Run all model benchmarks (text + VLM): requires local model weights
make bench

# Text models only
make bench-text

# VLM models only
make bench-vlm

# Single model
make bench-model MODEL=./models/qwen3-8b

# Show last benchmark results summary
make bench-report

The benchmark harness runs the model with a standard prompt and reports tokens per second. Results are written to benchmark_results.log.

Engine Parity and B=1 Decode Bench¶

mlxcel generate, mlxcel run, and the server all decode on the same engine (mlxcel_core::engine), so their token streams for one request are expected to match. mlxcel-engine-parity checks that: it runs one prompt through the CLI raw-completion client (a:cli), the server BatchScheduler in-process at B=1 with dense (b:dense) and paged (c:paged) decode storage, a direct Engine run (d:engine), mlxcel run -p (e:run), and the mlxcel-server default configuration (f:server, greedy case only). For every pair it prints identical or the first divergent token. On the dense arm it also compares a prompt-cache miss against a hit and reports how each side's prefill was split.

cargo build --release --bin mlxcel-engine-parity
MLXCEL_SDPA_DETERMINISTIC=1 ./target/release/mlxcel-engine-parity \
    --model models/mlx/qwen3-1.7b-4bit --expect-identical

Without --expect-identical the run records a baseline and exits 0; with it, any divergence exits 1. --expect-identical refuses to run unless MLXCEL_SDPA_DETERMINISTIC=1 is set, because on CUDA the default decode SDPA reduction order is not reproducible. --cases greedy,seeded, --no-chat-template, --server-prefill-chunk N, --no-prompt-cache-case, and --json PATH narrow or record a run.

mlxcel-bench-engine measures single-stream decode on both front ends with the same synthesized prompt: cli (the DirectEngine loop mlxcel generate runs, the same measured call as mlxcel-bench-decode) and server (the model worker and BatchScheduler driven in-process with pre-tokenized ids, so HTTP and tokenization stay outside the timed region). The default contexts are 256 and 8192 tokens.

cargo build --release --bin mlxcel-bench-engine
./target/release/mlxcel-bench-engine --model models/mlx/qwen3-1.7b-4bit

--prefill-chunk sets the chunk on both paths and --decode-storage selects the server storage backend. One pass is a functional check; for numbers, run several interleaved processes. mlxcel-bench-decode reports Decode path: engine and accepts --temperature, --top-p, and --top-k so a profile can include the cost of a sampled draw.

Distributed Component Benchmarks¶

The distributed E2E test files include benchmark tests tagged with #[ignore] that measure throughput and latency for each component:

# TP collective benchmarks
cargo test --test tp_e2e -- --ignored --nocapture --test-threads=1

# Pipeline benchmarks
cargo test --test pipeline_e2e -- --ignored --nocapture --test-threads=1

# Disaggregated inference benchmarks
cargo test --test disaggregated_e2e -- --ignored --nocapture --test-threads=1

These benchmarks measure:

TP benchmarks: - All-reduce throughput vs tensor size and world size - Lockstep scheduler/executor throughput (sequences per second) - Crossover analysis: when TP overhead exceeds single-GPU cost

PP benchmarks: - Pipeline throughput vs number of stages and micro-batches - Schedule efficiency (bubble fraction) - Activation transfer bandwidth

DI benchmarks: - End-to-end latency: time from request submission to first token - Prefill scheduler throughput - Decode scheduler throughput - KV cache transfer time vs cache size and strategy - Crossover analysis: when disaggregation overhead exceeds the gain

Running Tests¶

Full Test Suite¶

# Using Make (recommended for development)
make test

# Equivalent cargo command
cargo test -- --test-threads=1

Library Tests Only¶

make test-lib
# or
cargo test --lib -- --test-threads=1

With test-utils Feature¶

Some tests require the test-utils feature flag to enable additional test helpers (mock transports, in-process cluster simulation):

cargo test --features test-utils -- --test-threads=1

Model Loading Tests (requires local weights)¶

cargo test --test model_loading -- --ignored --test-threads=1

TurboQuant KV Cache Quality Gate (requires local weights)¶

tests/turbo_kv_e2e.rs validates the KVCacheMode::Turbo4Asym quality regression budget for the TurboQuant KV cache compression path.

Three independent gates run per model:

1. Wikitext-2 perplexity gate. Aggregates negative log-likelihood over non-overlapping 4K-token chunks of tests/fixtures/wikitext2_excerpt.txt, then asserts (ppl_turbo - ppl_baseline) / ppl_baseline ≤ 0.01.
2. NIAH single-needle retrieval gate. Plants a magic word at 9 (depth, length) cells (depths 0%/50%/100% × lengths 1K/2K/4K) at 4K context plus 3 cells at 8K, generates ≤ 32 tokens at temperature 0, and asserts that the turbo retrieval count is at least the baseline retrieval count.
3. Rotation kurtosis sanity test. Loads a non-quantized K-projection weight tensor, applies the WHT + sign-flip rotation via mlxcel_core::cache::turbo::turbo4_v_rotate, and asserts the post-rotation non-excess kurtosis stays below 5.

Decode tok/s, prefill tok/s, and wall-clock latency are recorded (not gated) during each PPL run and appended to benchmarks/turbo_kv/<YYYY-MM-DD>_<machine>.csv for documentation updates and trend analysis.

All four tests are gated with #[ignore] and skip gracefully when the required model checkouts are absent. Run each gate individually to keep wall time predictable:

# PPL + NIAH gate: Qwen2.5-1.5B (base variant)
cargo test --test turbo_kv_e2e --release -- --ignored test_qwen25_15b_quality_gate --nocapture

# PPL + NIAH gate: Meta-Llama-3.1-8B-Instruct-4bit
cargo test --test turbo_kv_e2e --release -- --ignored test_llama31_8b_quality_gate --nocapture

# PPL + NIAH gate: gemma-3-4b-it-4bit
cargo test --test turbo_kv_e2e --release -- --ignored test_gemma3_4b_quality_gate --nocapture

# Rotation kurtosis sanity (uses any non-quantized model checkout)
cargo test --test turbo_kv_e2e --release -- --ignored test_rotation_kurtosis_sanity --nocapture

Prerequisite model checkouts (any subset enables the matching tests):

./target/release/mlxcel download mlx-community/Qwen2.5-1.5B-4bit
./target/release/mlxcel download mlx-community/Meta-Llama-3.1-8B-Instruct-4bit
./target/release/mlxcel download mlx-community/gemma-3-4b-it-4bit

The kurtosis sanity test requires a non-quantized model (e.g., qwen2.5-0.5b-bf16) because 4-bit quantized K-projection weights are stored as packed integers and cannot be reinterpreted as floats without dequantization.

VLM quality gates are not part of the current quality gate.

Quality Gate¶

The full quality gate script runs format check, clippy, and tests:

scripts/run_quality_gate.sh

This is equivalent to:

make ci
# which runs: fmt-check + check + clippy + test

CPU-only Link Check¶

A Linux build with no GPU feature compiles MLX without any GPU backend, so MLX's GPU-side helpers do not exist in it. Bridge C++ under src/lib/mlxcelverse/kernels/ and src/lib/mlxcel-core/cpp/ that calls one must sit behind the MLXCEL_BRIDGE_GPU_BACKEND preprocessor define, which mlxcel-core/build.rs sets only when MLX builds Metal, CUDA, or ROCm. A missing guard shows up only as a link error in this configuration, since cargo check does not link. Check it with a separate target directory so it does not evict a GPU build cache:

CARGO_TARGET_DIR=target/cpu cargo test -p mlxcel-core --lib --no-run

Documentation Tests¶

make test-doc
# or
cargo test --doc

Test Organization Conventions¶

1. Co-locate tests with code: use #[path = "<module>_tests.rs"] to keep the test file next to its source file.

2. One test file per module: avoid mixing tests for different modules in a single file.

3. Integration tests in tests/: anything that crosses module boundaries or tests the full stack belongs in tests/.

4. Mark hardware-dependent tests as #[ignore]: any test that requires real GPU hardware, network, or local model weights must be marked #[ignore]. These tests should be documentable and runnable manually.

5. Use descriptive test names: test function names should describe what they test and what the expected outcome is, e.g., test_column_parallel_shard_plan_produces_correct_axis().

6. Keep unit tests fast: unit tests must complete in milliseconds. If a test takes more than a few seconds, consider marking it #[ignore] and moving it to the integration test suite.