Phase 3, Section 1: Scope, Milestones, and Exit Criteria

Phase 3: Sharding (the distribution axis). Section 1 of 4: Scope, Milestones, and Exit Criteria.

1. What this section commits to

This document outlines the scope, milestones, abandonment triggers, and exit criteria for Phase 3 of the Kina research program. The overarching goal of this phase is to produce a formally typed, unverified implementation of GSPMD-style sharding propagation, and the runtime foundation to execute it.

2. Scope Boundaries

In Scope

• Device Meshes as Kinds: Modeling logical device grids (1D or N-D) in the core type system.
• Sharding Refinements: S-expression syntax to annotate tensors with sharding strategies (e.g., @sharded, @replicated).
• GSPMD Type Inference: A compiler pass to propagate partial sharding annotations across operations like matmul and add.
• Type-Changing Coercions: Automatically inserting collective operations (all_reduce, all_gather, reduce_scatter, all_to_all) to satisfy sharding conflicts.
• Multi-Device Simulator: A local, single-process multi-threading simulator to execute the distributed MLIR runtime.

Deferred (Out of Scope)

• Lean 4 Mechanized Proofs: Verified sharding propagation is deferred to Phase 6.
• Sparse Sharding: Sharding of unstructured sparse tensors or MoE experts (Phase 4).
• Heterogeneous Hardware: Targeting multiple distinct hardware capabilities simultaneously (Phase 5).
• Physical Multi-Node Runtime: Production distributed runtime over network interconnects (slated for post-Phase 7).

3. Milestones

Milestone 1: Sharding Syntax & Core Calculus

Goal: Integrate device meshes and sharding refinements into the parser, AST, and typed IR.

• Define a 1D logical device mesh abstraction (MVP).
• Extend Ast_surface.expr and Ir_typed.tensor_type with sharding annotations.
• Update kina_parser.mly to support the new annotations (e.g., @replicated, @sharded(axis=0)).

Milestone 2: GSPMD Propagation Engine

Goal: Build the middle-end type inference pass that propagates sharding constraints.

• Implement the forward and backward propagation rules for elementwise, reduction, and contraction operations.
• Integrate the propagation solver into the Type_formation pipeline.
• Produce user-friendly type errors for irreconcilable sharding constraints.

Milestone 3: Collective Lowering & Simulator

Goal: Emit MLIR with collective operations and execute it on a simulated mesh.

• Extend the backend (emitter.ml) to lower type coercions into standard MLIR collective ops (or Kina-specific custom dialect ops if needed).
• Build a single-process threaded simulator in the Rust runtime to mock a 4-device mesh and execute the MLIR using shared memory barriers.

4. Risks and Abandonment Criteria

To keep the research falsifiable, we establish the following triggers:

1. Intractable Propagation Complexity

• Risk: GSPMD rules become undecidable or computationally explosive in the bidirectional type-checker.
• Trigger: Propagation takes > 10 seconds for a standard transformer layer, or requires arbitrary user annotations to converge.
• Action: Abandon automatic type-inferred sharding; fall back to mandatory explicit sharding annotations on every tensor.

2. Simulator Deadlocks

• Risk: The single-process threaded simulator fails to faithfully model NCCL semantics, leading to deadlocks.
• Trigger: Cannot reliably run a 4-device training loop without race conditions or deadlocks after 3 weeks of debugging.
• Action: Skip the simulator; use an existing framework (like XLA/PJRT) as the backend execution engine.

5. Exit Criteria

Phase 3 is complete when a release candidate satisfies all the following criteria:

Compiler & Type-Checker

• Compiler parses and type-checks @sharded and @replicated annotations.
• Sharding inference pass correctly deduces optimal layouts for a Transformer block without full annotations.
• Conflicting sharding annotations produce a typed compilation error, not a runtime crash.

Lowering & Runtime

• Sharding mismatches are successfully coerced by lowering to collective communication ops in MLIR.
• The threaded simulator runs the MLIR artifacts simulating at least 4 devices.
• A demo transformer (~100M parameters) successfully compiles and executes a forward/backward pass across the simulated mesh.