{
  "schemaVersion": 1,
  "id": "atlas-lesson-sharded-matmul-ownership",
  "title": "Sharded matrix multiplication and result ownership",
  "category": "learning",
  "format": "Atlas lesson",
  "description": "Derive the communication operation from the partitioned matrix axis and the next consumer’s required result layout.",
  "canonicalUrl": "https://semiaifoundry.com/learn/atlas/sharded-matmul-ownership/",
  "publishedDate": null,
  "revisionDate": "2026-10-07",
  "author": null,
  "topic": "Executable software",
  "series": null,
  "license": null,
  "rightsUrl": "https://semiaifoundry.com/ai-atlas/licenses.html",
  "version": "1.3.0",
  "status": "Published educational lesson",
  "evidenceScope": "The sources establish sharding and collective semantics. All matrices, the two-row exchange schedule and byte counts here are synthetic. This is not an observed NCCL implementation, a topology model or a performance comparison.",
  "contentCoverage": "published-lesson",
  "artifactUrl": "https://semiaifoundry.com/ai-atlas/lesson-data/sharded-matmul-ownership-cdb4f7ad8c2ad536.json",
  "interactiveUrl": "https://semiaifoundry.com/ai-atlas/#sharded-matmul-ownership",
  "related": [],
  "supersedes": [],
  "supersededBy": [],
  "contentHash": "d67325120a946e3bf935348c66795a1336ca7d911e8c531980b4516ffe21ed28",
  "contentLength": 9004,
  "jsonUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-sharded-matmul-ownership.json",
  "markdownUrl": "https://semiaifoundry.com/agents/publications/atlas-lesson-sharded-matmul-ownership.md",
  "sections": [
    {
      "id": "section-1",
      "title": "Sharded matrix multiplication and result ownership",
      "start": 36,
      "end": 532
    },
    {
      "id": "section-2",
      "title": "The mechanism",
      "start": 532,
      "end": 2583
    },
    {
      "id": "section-3",
      "title": "Science and assumptions",
      "start": 2583,
      "end": 2611
    },
    {
      "id": "section-4",
      "title": "Inner-dimension partition preserves the sum",
      "start": 2611,
      "end": 3414
    },
    {
      "id": "section-5",
      "title": "Two-rank payload ledger",
      "start": 3414,
      "end": 4318
    },
    {
      "id": "section-6",
      "title": "Engineering example",
      "start": 4318,
      "end": 5189
    },
    {
      "id": "section-7",
      "title": "Distinctions and tradeoffs",
      "start": 5189,
      "end": 5738
    },
    {
      "id": "section-8",
      "title": "Specifications and interfaces",
      "start": 5738,
      "end": 6116
    },
    {
      "id": "section-9",
      "title": "Interfaces",
      "start": 6116,
      "end": 6305
    },
    {
      "id": "section-10",
      "title": "Check your understanding",
      "start": 6305,
      "end": 6334
    },
    {
      "id": "section-11",
      "title": "Each rank computed all output coordinates using a different half of the contracted dimension. Is concatenating those local arrays enough?",
      "start": 6334,
      "end": 6822
    },
    {
      "id": "section-12",
      "title": "After the stated reduce-scatter, which row belongs to rank 1?",
      "start": 6822,
      "end": 7035
    },
    {
      "id": "section-13",
      "title": "How many payload bytes are sent across both ranks in the two stated phases?",
      "start": 7035,
      "end": 7285
    },
    {
      "id": "section-14",
      "title": "Evidence and scope",
      "start": 7285,
      "end": 7748
    },
    {
      "id": "section-15",
      "title": "Sources",
      "start": 7748,
      "end": 8192
    },
    {
      "id": "section-16",
      "title": "Continue exploring",
      "start": 8192,
      "end": 8215
    },
    {
      "id": "section-17",
      "title": "Foundations",
      "start": 8215,
      "end": 8538
    },
    {
      "id": "section-18",
      "title": "Related lessons",
      "start": 8538,
      "end": 9004
    }
  ],
  "citations": [
    {
      "title": "Explore in Atlas",
      "url": "https://semiaifoundry.com/ai-atlas/#sharded-matmul-ownership"
    },
    {
      "title": "JAX — Manual parallelism with shard_map",
      "url": "https://docs.jax.dev/en/latest/notebooks/shard_map.html"
    },
    {
      "title": "NVIDIA NCCL — Collective Operations",
      "url": "https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html"
    },
    {
      "title": "Published lesson data",
      "url": "https://semiaifoundry.com/ai-atlas/lesson-data/sharded-matmul-ownership-cdb4f7ad8c2ad536.json"
    },
    {
      "title": "Atlas releases and licensing",
      "url": "https://semiaifoundry.com/ai-atlas/licenses.html"
    },
    {
      "title": "Compiler layout, lowering and buffer planning",
      "url": "https://semiaifoundry.com/learn/atlas/compiler-layout-lowering/"
    },
    {
      "title": "Collective communication library",
      "url": "https://semiaifoundry.com/learn/atlas/nccl/"
    },
    {
      "title": "Data, tensor, pipeline, and state parallelism",
      "url": "https://semiaifoundry.com/learn/atlas/parallelism/"
    },
    {
      "title": "Systolic arrays and tensor dataflow",
      "url": "https://semiaifoundry.com/learn/atlas/systolic-array-dataflow/"
    },
    {
      "title": "Collective traffic and congestion",
      "url": "https://semiaifoundry.com/learn/atlas/fabric-collective-traffic/"
    },
    {
      "title": "Runtime, driver and device dispatch",
      "url": "https://semiaifoundry.com/learn/atlas/runtime-device-dispatch/"
    },
    {
      "title": "Floating-point reductions and compiler permissions",
      "url": "https://semiaifoundry.com/learn/atlas/floating-point-reduction-order/"
    }
  ],
  "recordHash": "4bad5c886f904ef3f561ae10f23ca12b74c6dfefc5f5bdeb6153690955bfc51b",
  "contentFormat": "text/markdown",
  "content": "Executable software · Atlas v1.3.0\n\n# Sharded matrix multiplication and result ownership\n\nDerive the communication operation from the partitioned matrix axis and the next consumer’s required result layout.\n\n[Explore in Atlas](https://semiaifoundry.com/ai-atlas/#sharded-matmul-ownership)Logical values and distributed ownership; no physical scale or vendor timing\n\nLearning objectiveReconstruct the same 2×2 product from inner-dimension shards, choose a collective from output ownership, and reconcile the two-rank payload ledger.\n\n## The mechanism\n\nA compiler must preserve both a tensor’s values and who owns them. Splitting matrix output rows gives different ranks independent output rows; splitting the contracted inner dimension gives each rank only part of every output sum. Those are different meanings of a local matrix product, even if the local arrays have similar shapes. JAX’s explicit sharding example forms local products and then reduces their partial sums. A sum all-reduce makes the completed result available at every rank. Reduce-scatter instead finishes the sum while distributing different output pieces, and all-gather can later assemble those pieces when a consumer needs a replica. Concatenating unfinished partial sums does not perform their missing arithmetic. This distinction connects tensor-axis information in a compiler to the collective scheduled over the network. It also prevents allocating replicas merely because the preceding stage happened to use all-reduce. A valid next operator may consume distributed rows directly, while another needs a different layout and an explicit conversion. The example below uses exactly representable small integers; its equality does not imply arbitrary floating-point reduction trees are bitwise identical.\n\n1.  Declare C = AB, the complete tensor shapes, stored element size and required consumer ownership before partitioning.\n2.  Split the inner dimension into disjoint sets Kᵣ; each rank receives matching columns of A and rows of B.\n3.  Compute Cᵣ locally. Each Cᵣ covers all output coordinates but contains only the products indexed by Kᵣ.\n4.  Add contributions with matching output coordinates. Use all-reduce for replicated C or reduce-scatter for completed disjoint output shards.\n5.  Retain row shards when the next operation can use them; all-gather only when the consumer requires the whole tensor.\n6.  Check every output against a separate dense reference, and count a sent payload once at its chosen link boundary. Input distribution, temporary buffers and protocol overhead are additional costs.\n\n## Science and assumptions\n\n### Inner-dimension partition preserves the sum\n\n$$\nC^{(r)}_{ij}=\\sum_{k\\in K_r}A_{ik}B_{kj},\\qquad C_{ij}=\\sum_{r=0}^{P-1}C^{(r)}_{ij}\n$$\n\nRank r sums products for its contracted indices $K_{r}$; completed C sums the rank contributions at matching coordinates.\n\nSymbols and units\n\nCij\n\nCompleted matrix output element\n\nvalue\n\nC(r)ij\n\nPartial contribution from rank r\n\nvalue\n\nAik, Bkj\n\nInput matrix elements\n\nvalue\n\nKr\n\nDisjoint contracted indices assigned to rank r\n\nindex set\n\nP\n\nNumber of participating ranks\n\ncount\n\nThe disjoint Kᵣ sets cover the contracted dimension exactly once. Each rank contributes to the same output coordinates.\n\nAssumptions Exact real arithmetic or the small integer example; floating-point summation order can change rounding. No duplicated or omitted contracted indices.\n\n### Two-rank payload ledger\n\n$$\nQ_{\\mathrm{RS}}=2\\times4=8\\ \\mathrm{bytes/rank},\\quad Q_{\\mathrm{AG}}=8\\ \\mathrm{bytes/rank},\\quad Q_{\\mathrm{sent}}=2(8+8)=32\\ \\mathrm{bytes}\n$$\n\nThis two-rank schedule sends two float32 values per rank per phase: eight bytes for reduce-scatter and eight for all-gather, giving thirty-two bytes across both senders.\n\nSymbols and units\n\nQRS\n\nPayload sent by each rank in the chosen reduce-scatter exchange\n\nbytes/rank\n\nQAG\n\nPayload sent by each rank in the chosen all-gather exchange\n\nbytes/rank\n\nQsent\n\nPayload sent across both ranks and both phases\n\nbytes\n\nFor this explicitly described two-rank, 2×2 float32-output schedule, each phase sends one two-element row per rank.\n\nAssumptions Logical payload sent by each rank in this chosen exchange; excludes initial input placement, headers, retransmission, local reads/writes and actual collective implementation details.\n\n## Engineering example\n\nUse A = \\[\\[1,2\\],\\[3,4\\]\\] and B = \\[\\[5,6\\],\\[7,8\\]\\]. Rank 0 owns k=0 and forms C₀ = \\[\\[5,6\\],\\[15,18\\]\\]; rank 1 owns k=1 and forms C₁ = \\[\\[14,16\\],\\[28,32\\]\\]. Reduce-scatter with row ownership sends C₀ row 1 \\[15,18\\] from rank 0 to rank 1, and C₁ row 0 \\[14,16\\] in the opposite direction. Rank 0 adds \\[5,6\\]+\\[14,16\\]=\\[19,22\\]; rank 1 adds \\[15,18\\]+\\[28,32\\]=\\[43,50\\]. Each sends 8 payload bytes if elements are float32. If both consumers need all of C, all-gather exchanges these completed rows for another 8 bytes sent per rank. Both then hold \\[\\[19,22\\],\\[43,50\\]\\]. Total sent payload is 16 bytes per rank, 32 across the two senders; counting the matching receives again would double-count this network boundary. Concatenating the two original 2×2 partial matrices would create four rows and would leave their sums unfinished.\n\n## Distinctions and tradeoffs\n\nA shard can be an unfinished contribution to every output, not necessarily a disjoint piece of the finished output.\n\nA shard-preserving collective may save result replication, but a later consumer can reintroduce an all-gather or resharding cost. Logical bytes alone do not predict network latency.\n\nDistributed linear layers become correct only when the compiler and runtime finish the same tensor operation and deliver its expected ownership. Avoiding an unnecessary replica can also save memory and communication.\n\n## Specifications and interfaces\n\nContracted-axis split\n\nPartial sums overlap in output coordinates\n\nRequires addition, not concatenation.\n\nOutput-row split\n\nDisjoint completed rows\n\nConcatenation can reconstruct the global output without summing corresponding rows.\n\nConsumer contract\n\nReplicated or sharded result\n\nThe useful collective is determined by the required layout.\n\n### Interfaces\n\nCompiler sharding plan Input\n\nTensor shapes, partitioned axes, dtype and input ownership\n\nConsumer-ready output Output\n\nCompleted values with explicit per-rank coordinates\n\n## Check your understanding\n\n### Each rank computed all output coordinates using a different half of the contracted dimension. Is concatenating those local arrays enough?\n\n1.  No: matching coordinates still require a sum\n2.  Yes: more rows complete the dot products\n3.  Yes: any collective performs the missing multiplication\n\nAnswer and reasoning\n\nNo: matching coordinates still require a sum\n\nThe local arrays are partial sums over different k indices. Reduce them before choosing replicated or sharded ownership.\n\n### After the stated reduce-scatter, which row belongs to rank 1?\n\n1.  \\[14,16\\]\n2.  \\[43,50\\]\n3.  \\[19,22\\]\n\nAnswer and reasoning\n\n\\[43,50\\]\n\nRank 1 combines both partial contributions to the second output row.\n\n### How many payload bytes are sent across both ranks in the two stated phases?\n\n1.  16 bytes\n2.  32 bytes\n3.  64 bytes\n\nAnswer and reasoning\n\n32 bytes\n\nTwo ranks each send 8 bytes in each of two phases. The same received bytes are not added again.\n\n## Evidence and scope\n\nThe sources establish sharding and collective semantics. All matrices, the two-row exchange schedule and byte counts here are synthetic. This is not an observed NCCL implementation, a topology model or a performance comparison.\n\nThe ideal ring all-reduce lab varies communication payload, ranks, startup and bandwidth. It does not execute these matrix shards or alter the optimizer; the exact computation is worked through in this lesson.\n\n## Sources\n\n-   [JAX — Manual parallelism with shard\\_map](https://docs.jax.dev/en/latest/notebooks/shard_map.html)\n-   [NVIDIA NCCL — Collective Operations](https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/collectives.html)\n\n[Published lesson data](https://semiaifoundry.com/ai-atlas/lesson-data/sharded-matmul-ownership-cdb4f7ad8c2ad536.json) · [Atlas releases and licensing](https://semiaifoundry.com/ai-atlas/licenses.html)\n\n## Continue exploring\n\n### Foundations\n\n-   [Compiler layout, lowering and buffer planning](https://semiaifoundry.com/learn/atlas/compiler-layout-lowering/)\n-   [Collective communication library](https://semiaifoundry.com/learn/atlas/nccl/)\n-   [Data, tensor, pipeline, and state parallelism](https://semiaifoundry.com/learn/atlas/parallelism/)\n\n### Related lessons\n\n-   [Systolic arrays and tensor dataflow](https://semiaifoundry.com/learn/atlas/systolic-array-dataflow/)\n-   [Collective traffic and congestion](https://semiaifoundry.com/learn/atlas/fabric-collective-traffic/)\n-   [Runtime, driver and device dispatch](https://semiaifoundry.com/learn/atlas/runtime-device-dispatch/)\n-   [Floating-point reductions and compiler permissions](https://semiaifoundry.com/learn/atlas/floating-point-reduction-order/)"
}