{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "40530272",
   "metadata": {},
   "source": [
    "# Course 4 Lab — Infrastructure Mechanics\n",
    "\n",
    "Five deterministic teaching miniatures aligned with Kimi K3 Section 5. They execute **affine composition, balancing arithmetic, activation lifetimes, hybrid-cache boundary validation, and feasible fleet admission**. They do not execute distributed kernels, MoonEP, remote offload, AgentENV, or a production serving fleet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e0193789",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:27.455606Z",
     "iopub.status.busy": "2026-08-11T03:31:27.455332Z",
     "iopub.status.idle": "2026-08-11T03:31:28.648171Z",
     "shell.execute_reply": "2026-08-11T03:31:28.646678Z"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import subprocess, sys, tempfile, torch\n",
    "repo = next((p for p in (Path.cwd(), *Path.cwd().parents) if (p / 'src').is_dir()), None)\n",
    "if repo is None:\n",
    "    repo = (Path('/content') if Path('/content').is_dir() else Path(tempfile.gettempdir())) / 'build-Kimi-K3-architecture'\n",
    "    if not (repo / 'src').is_dir():\n",
    "        subprocess.run(['git', 'clone', '--depth', '1', 'https://github.com/mailtotanvir/build-Kimi-K3-architecture.git', str(repo)], check=True)\n",
    "sys.path.insert(0, str(repo))\n",
    "from src.infrastructure.learning_miniatures import (\n",
    "    activation_memory_timeline, affine_prefix_scan, apply_affine,\n",
    "    balanced_expert_plan, choose_feasible_node, hybrid_cache_hit,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94db6767",
   "metadata": {},
   "source": [
    "## 5.1.2 — Affine KDA segment composition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9a6a8deb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:28.650319Z",
     "iopub.status.busy": "2026-08-11T03:31:28.650063Z",
     "iopub.status.idle": "2026-08-11T03:31:28.658914Z",
     "shell.execute_reply": "2026-08-11T03:31:28.657895Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "sequential terminal state: [2.602, 0.966]\n",
      "prefix-scan terminal state: [2.602, 0.966]\n",
      "exact miniature equality: True\n"
     ]
    }
   ],
   "source": [
    "transforms = [\n",
    "    (torch.tensor([[1.0,.2],[0.0,.9]]), torch.tensor([1.0,0.0])),\n",
    "    (torch.tensor([[.8,0.0],[.1,1.0]]), torch.tensor([0.0,2.0])),\n",
    "    (torch.tensor([[1.0,-.1],[0.0,.7]]), torch.tensor([.5,0.0])),\n",
    "]\n",
    "initial = torch.tensor([2.0,-1.0])\n",
    "state = initial\n",
    "for transform in transforms: state = apply_affine(transform, state)\n",
    "scan_state = apply_affine(affine_prefix_scan(transforms)[-1], initial)\n",
    "assert torch.allclose(state, scan_state)\n",
    "print('sequential terminal state:', [round(x,4) for x in state.tolist()])\n",
    "print('prefix-scan terminal state:', [round(x,4) for x in scan_state.tolist()])\n",
    "print('exact miniature equality:', torch.allclose(state, scan_state))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ad1ade0f",
   "metadata": {},
   "source": [
    "## 5.2.1 — Rank-balancing arithmetic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "7a07ed7d",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:28.660713Z",
     "iopub.status.busy": "2026-08-11T03:31:28.660494Z",
     "iopub.status.idle": "2026-08-11T03:31:28.664797Z",
     "shell.execute_reply": "2026-08-11T03:31:28.663889Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "rank loads: [120, 50, 175, 75]\n",
      "equal capacity: 105\n",
      "overflow by rank: (15, 0, 70, 0)\n",
      "spare capacity by rank: (0, 55, 0, 30)\n",
      "minimum tokens moved to remove overflow: 85\n",
      "Boundary: this arithmetic does not choose MoonEP expert replicas or perform network dispatch.\n"
     ]
    }
   ],
   "source": [
    "loads = [120, 50, 175, 75]\n",
    "plan = balanced_expert_plan(loads)\n",
    "assert plan.capacity == 105 and plan.moved_tokens == 85\n",
    "print('rank loads:', loads)\n",
    "print('equal capacity:', plan.capacity)\n",
    "print('overflow by rank:', plan.overflow)\n",
    "print('spare capacity by rank:', plan.spare)\n",
    "print('minimum tokens moved to remove overflow:', plan.moved_tokens)\n",
    "print('Boundary: this arithmetic does not choose MoonEP expert replicas or perform network dispatch.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c70d0398",
   "metadata": {},
   "source": [
    "## 5.2.2 — Activation lifetime ledger"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "5eaa2aba",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:28.666733Z",
     "iopub.status.busy": "2026-08-11T03:31:28.666554Z",
     "iopub.status.idle": "2026-08-11T03:31:28.672062Z",
     "shell.execute_reply": "2026-08-11T03:31:28.670928Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step 0: 40 bytes live — activation\n",
      "step 1: 65 bytes live — activation, checkpoint\n",
      "step 2: 65 bytes live — activation, checkpoint\n",
      "step 3: 60 bytes live — checkpoint, gradient\n",
      "step 4: 35 bytes live — gradient\n",
      "peak live bytes: 65\n"
     ]
    }
   ],
   "source": [
    "records = [('activation',40,0,3), ('checkpoint',25,1,4), ('gradient',35,3,5)]\n",
    "timeline, peak = activation_memory_timeline(records, 5)\n",
    "assert [row['bytes'] for row in timeline] == [40,65,65,60,35]\n",
    "for row in timeline: print(f\"step {row['step']}: {row['bytes']:2d} bytes live — {', '.join(row['live'])}\")\n",
    "print('peak live bytes:', peak)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba4cdb26",
   "metadata": {},
   "source": [
    "## 5.4.1 — Hybrid cache boundary consistency"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "8d8461cb",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:28.673738Z",
     "iopub.status.busy": "2026-08-11T03:31:28.673562Z",
     "iopub.status.idle": "2026-08-11T03:31:28.677962Z",
     "shell.execute_reply": "2026-08-11T03:31:28.676824Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MLA=2560, KDA=2560, request=2560 -> valid=True\n",
      "MLA=2560, KDA=2048, request=2560 -> valid=False\n",
      "MLA=2048, KDA=2048, request=2560 -> valid=False\n"
     ]
    }
   ],
   "source": [
    "cases = [(2560,2560,2560), (2560,2048,2560), (2048,2048,2560)]\n",
    "for mla_end, kda_checkpoint, requested in cases:\n",
    "    valid = hybrid_cache_hit(mla_end, kda_checkpoint, requested)\n",
    "    print(f'MLA={mla_end}, KDA={kda_checkpoint}, request={requested} -> valid={valid}')\n",
    "assert [hybrid_cache_hit(*case) for case in cases] == [True,False,False]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf2fce0d",
   "metadata": {},
   "source": [
    "## 5.4.3 — Cache affinity subject to feasibility"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "cfacede5",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2026-08-11T03:31:28.679777Z",
     "iopub.status.busy": "2026-08-11T03:31:28.679598Z",
     "iopub.status.idle": "2026-08-11T03:31:28.684411Z",
     "shell.execute_reply": "2026-08-11T03:31:28.683172Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "highest-affinity node: cached-hot\n",
      "admitted node: feasible\n",
      "reason: cached-hot exceeds the pressure limit; locality is not permission.\n"
     ]
    }
   ],
   "source": [
    "nodes = [\n",
    " {'name':'cached-hot','affinity':1.0,'queue':.1,'pressure':.96,'budget_risk':.2},\n",
    " {'name':'feasible','affinity':.6,'queue':.1,'pressure':.4,'budget_risk':.1},\n",
    " {'name':'slow','affinity':.2,'queue':.8,'pressure':.3,'budget_risk':.4},\n",
    "]\n",
    "chosen = choose_feasible_node(nodes)\n",
    "assert chosen['name'] == 'feasible'\n",
    "print('highest-affinity node:', max(nodes,key=lambda n:n['affinity'])['name'])\n",
    "print('admitted node:', chosen['name'])\n",
    "print('reason: cached-hot exceeds the pressure limit; locality is not permission.')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b1dfb0f",
   "metadata": {},
   "source": [
    "## Evidence boundary\n",
    "\n",
    "Executed here: algebraic affine-prefix equivalence, equal-capacity overflow arithmetic, interval-based peak memory, exact hybrid-cache boundary checks, and feasibility-filtered node choice. Paper reported only: FlashKDA performance, KCP distributed implementation, MoonEP replica planning and zero-copy transfer, remote activation offload, AgentENV, production cache layout, and fleet behavior."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}