feat: Phase 43 — Sparse/Grouped MoE Dispatch (v4.0.0-alpha.3)

[aarambh-darshan]

Summary

Resolves the "documented future optimisation" carried forward unresolved
since v2 §35 and v3's out-of-scope list: real sparse dispatch,
where each token's forward pass only computes its assigned top-k experts,
rather than every expert computing on every token and being masked
afterward (the dense-masked-matmul trade-off v2 Phase 22 and v3 Phase 31
deliberately shipped).

The output is numerically equivalent to DenseMasked (same tokens, same
weights, same per-expert reduction order) — just faster, because the
masked-away matmuls never run. This is the milestone tagged
v4.0.0-alpha.3.

What's New

Core (aarambh-studio-core)

• DispatchKind enum (src/config.rs, alongside AttentionKind):
  • DenseMasked — v2/v3 behaviour, default, CPU fallback, correctness reference
  • Sparse — new, CUDA-only throughput win
  • Serialized as dense_masked / sparse (snake_case)
• MoeConfig.dispatch field defaulting to DenseMasked for exact
  backward compatibility with every existing MoE checkpoint — old
  TOML/JSON without the field deserialises to the dense path,
  byte-identical to v2/v3.
• Exported DispatchKind from lib.rs.
• New tests: moe_config_dispatch_defaults_to_dense_masked,
  moe_config_dispatch_serializes_as_snake_case, + assertion in
  old_moe_json_defaults_to_phase22_behavior.

Neural Network (aarambh-studio-nn)

• sparse_grouped_dispatch (src/dispatch.rs): tokens are grouped by
  router assignment into per-expert contiguous batches via arg_sort
  (no-grad permutation), then each expert's SwiGLU matmul executes only
  on its assigned token group via index_select → expert.forward →
  index_add scatter. Fully differentiable through candle's
  gather / index_select / index_add — router logits, expert
  parameters, and input activations all receive correct gradients.
• effective_dispatch_kind: Sparse only activates on a CUDA device;
  the CPU path keeps DenseMasked regardless of configuration, documented
  plainly as "GPU only pays off" (the honesty discipline v2 §29 applied to
  speculative decoding's speed claim).
• MoeFfn::dispatch_kind(): exposes the configured kind. The forward
  path selects dense vs sparse by the effective kind. QAT calibration
  (forward_with_capture) always uses the dense reference to observe
  full per-expert activation distributions.
• Load-balancing auxiliary loss unchanged — computed in
  top_k_gating before dispatch, so identical for both kinds. Sparse
  changes the compute path only, not the loss the router is trained
  against.
• Re-exported sparse_grouped_dispatch, effective_dispatch_kind from
  lib.rs.

Why no custom .cu kernel

candle's index_select / index_add / matmul already use cuBLAS on
CUDA, so sparse_grouped_dispatch is the CUDA grouped path and truly
skips non-routed experts. A hand-written grouped-GEMM kernel can't be
tested in CPU CI and the release audit forbids empty kernel bodies —
documented as a future optimisation.

Tests (18 new)

Test	Gate
sparse_dispatch_output_matches_dense_masked_reference_within_tolerance	correctness (max abs diff < 1e-5)
dispatch_kind_dense_masked_is_bit_identical_to_v2_v3_behaviour	backward compat (diff == 0.0)
sparse_dispatch_supports_top_k_greater_than_one	top_k > 1 equivalence
sparse_dispatch_backward_reaches_router_and_expert_weights	differentiability
sparse_dispatch_empty_expert_group_is_skipped	empty expert group handled
sparse_dispatch_matches_dense_with_shared_expert_summed_separately	shared expert interaction
load_balancing_loss_value_is_unaffected_by_dispatch_kind	aux loss dispatch-independent
sparse_configured_moe_falls_back_to_dense_masked_on_cpu	CPU fallback policy
effective_dispatch_kind_uses_sparse_on_cuda	CUDA selection (skips on CPU)
sparse_dispatch_cuda_throughput_exceeds_dense_masked_at_kaggle_gpu_scale	wall-clock (CUDA-gated)
moe_config_dispatch_defaults_to_dense_masked	config default
moe_config_dispatch_serializes_as_snake_case	serde round-trip
+ existing dispatch tests retained	regression

CUDA-only tests use Device::cuda_if_available(0) and early-return on
CPU — they compile cleanly in CPU CI and exercise the real CUDA path
when GPU hardware is present.

Configs & Scripts

• configs/moe_sparse_smoke.toml — tiny CPU smoke (dispatch = "sparse",
  runs through the dense fallback).
• configs/large_sparse_moe.toml — Kaggle GPU config mirroring
  large_finegrained_moe.toml with dispatch = "sparse".
• scripts/phase43_smoke.sh — unit tests + 2-step CPU train + checkpoint
  verification + scorecard.

Docs

• docs/phase43_sparse_moe.md — full phase doc (mechanism, differentiability,
  CPU/CUDA policy, test matrix).
• ROADMAP_V4.md — Phase 43 tasks marked [x], tests annotated as
  implemented.
• CHANGELOG.md — added [4.0.0-alpha.3] section.
• README.md — bumped alpha.2 → alpha.3, added Phase 43 to v4 arc +
  capabilities table + phase-docs table.
• ARCHITECTURE_V4.md §57 — added "Status: shipped" + implementation
  subsection.
• SELF_LEARNING_V4.md §44 — added "Status: shipped" note.
• Version bumped 4.0.0-alpha.2 → 4.0.0-alpha.3 across all 20
  packages (workspace inheritance).

CI Gates (all green ✅)

Gate	Result
cargo fmt --all --check	✅
cargo check --workspace --all-targets --locked	✅
cargo clippy --workspace --all-targets --locked -- -D warnings -D clippy::undocumented_unsafe_blocks	✅
cargo test --workspace --no-fail-fast --locked	✅ 488 passed, 0 failed
RUSTDOCFLAGS="-D warnings -D missing_docs" cargo doc --workspace --no-deps --locked	✅
scripts/phase28_release_audit.sh	✅ (v4.0.0-alpha.3)
cargo build --release -p aarambh-studio --locked	✅ (26MB binary)
CLI --help smoke (all subcommands)	✅

Smoke Test

Phase 43 smoke (scripts/phase43_smoke.sh substance) ran a 2-step CPU
training on moe_sparse_smoke.toml (dispatch="sparse" → dense
fallback on CPU): loss 4.87 → 4.61, moe_aux computed correctly
(dispatch-independent), routed_experts=8 active_routed=2 shared_experts=1 dead_experts=0. Checkpoint saved with 45 tensors
including all 8 routed expert down-projections + 1 shared expert.
Scorecard: artifacts/phase43_sparse_moe_smoke.json.

Backward Compatibility

• Default dispatch = DenseMasked → byte-identical to v2/v3 behaviour
  (verified by dispatch_kind_dense_masked_is_bit_identical_to_v2_v3_behaviour,
  diff == 0.0).
• Old configs without the dispatch field deserialize unchanged.
• No new crate (still 20 packages).
• No new dependencies.
• Aux loss path unchanged — existing checkpoints and training recipes
  are unaffected.