[Pallas][Mosaic GPU] Add GPU pipelining docs

Author: justinjfu

docs/_static/pallas/gpu/pipeline_matmul.svg

@@ -134,6 +134,7 @@ def _do_not_evaluate_in_jax(
     'notebooks/*.md',
     'pallas/quickstart.md',
     'pallas/pipelining.md',
+    'pallas/gpu/pipelining.md',
     'pallas/tpu/pipelining.md',
     'pallas/tpu/distributed.md',
     'pallas/tpu/sparse.md',
@@ -230,6 +231,7 @@ def _do_not_evaluate_in_jax(
     # Requires accelerators
     'pallas/quickstart.*',
     'pallas/pipelining.*',
+    'pallas/gpu/pipelining.*',
     'pallas/tpu/pipelining.*',
     'pallas/tpu/distributed.*',
     'pallas/tpu/sparse.*',

docs/_static/pallas/gpu/pipeline_matmul_ws.svg

@@ -7,6 +7,7 @@ Backend specific documentation for the Mosaic GPU backend.
    :maxdepth: 2
 
    reference
+   pipelining
 
 .. toctree::
    :caption: Guides

docs/_static/pallas/gpu/warp_specialization.svg

@@ -10,7 +10,7 @@
     "\n",
     "Software pipelining is an important technique in performance optimization by overlapping multiple asynchronous operations even if there are data dependencies between them. In the context of kernel writing, the most common form of pipelining involves overlapping communication and memory transfers with compute such that the hardware accelerator never stalls while waiting for data to arrive. Therefore, we will solely focus on the problem of communication-compute pipelining in this tutorial. We will begin by covering the problem conceptually, outlining the Pallas API for writing pipelines, and going over some realistic examples using the API.\n",
     "\n",
-    "This tutorial only covers the conceptual foundations of pipelining. For platform-specific references, please see the [TPU](https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html), or GPU (coming soon!) specific pipelining references.\n"
+    "This tutorial only covers the conceptual foundations of pipelining. For platform-specific references, please see the [TPU](https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html), or [GPU](https://docs.jax.dev/en/latest/pallas/gpu/pipelining.html) specific pipelining references.\n"
    ]
   },
   {

docs/conf.py

@@ -17,7 +17,7 @@ jupyter:
 
 Software pipelining is an important technique in performance optimization by overlapping multiple asynchronous operations even if there are data dependencies between them. In the context of kernel writing, the most common form of pipelining involves overlapping communication and memory transfers with compute such that the hardware accelerator never stalls while waiting for data to arrive. Therefore, we will solely focus on the problem of communication-compute pipelining in this tutorial. We will begin by covering the problem conceptually, outlining the Pallas API for writing pipelines, and going over some realistic examples using the API.
 
-This tutorial only covers the conceptual foundations of pipelining. For platform-specific references, please see the [TPU](https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html), or GPU (coming soon!) specific pipelining references.
+This tutorial only covers the conceptual foundations of pipelining. For platform-specific references, please see the [TPU](https://docs.jax.dev/en/latest/pallas/tpu/pipelining.html), or [GPU](https://docs.jax.dev/en/latest/pallas/gpu/pipelining.html) specific pipelining references.
 
 <!-- #endregion -->