diff --git a/docs/Makefile b/docs/Makefile
index 55c0a3c56fac57e573339a9c24c7881d681379a9..800b4020ab21a30f920105cee12613bc0a1749e7 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -21,10 +21,13 @@ clean:
 
 # Catch-all target: route all unknown targets to Sphinx using the new
 # "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
-%: Makefile
+# %: Makefile
+
+html:
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
 clean:
 	rm -rf source/reference/generated
+	rm -rf source/api/generated
 	rm -rf source/backend/generated
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
\ No newline at end of file
diff --git a/docs/source/api/codegen.rst b/docs/source/api/codegen.rst
new file mode 100644
index 0000000000000000000000000000000000000000..cf8db1700895b32c5e86291270190a0cf421dbde
--- /dev/null
+++ b/docs/source/api/codegen.rst
@@ -0,0 +1,57 @@
+pystencils.codegen
+==================
+
+.. module:: pystencils.codegen
+
+Kernel Creation Driver
+----------------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+
+  create_kernel
+  get_driver
+  
+Configuration
+-------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/entire_class.rst
+
+  CreateKernelConfig
+  CpuOptimConfig
+  OpenMpConfig
+  VectorizationConfig
+  GpuIndexingConfig
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+
+  AUTO
+
+Target Specification
+--------------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/recursive_class.rst
+
+  Target
+
+Output Code Objects
+-------------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/entire_class.rst
+
+  Kernel
+  GpuKernel
+  Parameter
+  GpuThreadsRange
diff --git a/docs/source/reference/api/field.rst b/docs/source/api/field.rst
similarity index 100%
rename from docs/source/reference/api/field.rst
rename to docs/source/api/field.rst
diff --git a/docs/source/api/jit.rst b/docs/source/api/jit.rst
new file mode 100644
index 0000000000000000000000000000000000000000..5a6bb2f4f60e0ba3ebebc8bb97652e1c92bd7701
--- /dev/null
+++ b/docs/source/api/jit.rst
@@ -0,0 +1,41 @@
+pystencils.jit
+==============
+
+.. module:: pystencils.jit
+
+Base Infrastructure
+-------------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/entire_class.rst
+
+    KernelWrapper
+    JitBase
+    NoJit
+
+.. autodata:: no_jit
+
+Legacy CPU JIT
+--------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/entire_class.rst
+
+  LegacyCpuJit
+  KernelWrapper
+
+CuPy-based GPU JIT
+------------------
+
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/entire_class.rst
+
+  CupyJit
+  CupyKernelWrapper
+  LaunchGrid
diff --git a/docs/source/reference/api/sympyextensions.rst b/docs/source/api/sympyextensions.rst
similarity index 100%
rename from docs/source/reference/api/sympyextensions.rst
rename to docs/source/api/sympyextensions.rst
diff --git a/docs/source/backend/index.rst b/docs/source/backend/index.rst
index 74b57e27b6f267c04a0a08e2eead88cb99726b4e..5ab8dbd34eb37fbc38230f3db0506c572d4b6964 100644
--- a/docs/source/backend/index.rst
+++ b/docs/source/backend/index.rst
@@ -18,7 +18,6 @@ who wish to customize or extend the behaviour of the code generator in their app
     platforms
     transformations
     errors
-    jit
     extensions
 
 Internal Representation
diff --git a/docs/source/backend/jit.rst b/docs/source/backend/jit.rst
deleted file mode 100644
index 79776eb67a93203fcbe9c1ef95e48cd8c9c9c888..0000000000000000000000000000000000000000
--- a/docs/source/backend/jit.rst
+++ /dev/null
@@ -1,6 +0,0 @@
-************************
-Just-In-Time Compilation
-************************
-
-.. automodule:: pystencils.jit
-    :members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 66582cb4b6e847fc6b3022c1726257d866b3ba85..6aa09bdbdf74e070f03d450d37ea501230d88c02 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -69,28 +69,36 @@ Topics
 ------
 
 .. toctree::
-   :maxdepth: 1
-   :caption: Getting Started
+  :maxdepth: 1
+  :caption: Getting Started
 
-   installation
-   tutorials/index
+  installation
+  tutorials/index
 
 .. toctree::
-   :maxdepth: 1
-   :caption: Reference Guides
+  :maxdepth: 1
+  :caption: Reference Guides
 
-   reference/symbolic_language
-   reference/kernelcreation
-   reference/gpu_kernels
-   reference/types
-   reference/api/index
+  reference/symbolic_language
+  reference/kernelcreation
+  reference/gpu_kernels
+  reference/types
 
 .. toctree::
-   :maxdepth: 1
-   :caption: Advanced
+  :maxdepth: 1
+  :caption: API
 
-   migration
-   backend/index
+  api/field
+  api/sympyextensions
+  api/codegen
+  api/jit
+
+.. toctree::
+  :maxdepth: 1
+  :caption: Advanced
+
+  migration
+  backend/index
 
 Projects using pystencils
 -------------------------
diff --git a/docs/source/reference/api/codegen.rst b/docs/source/reference/api/codegen.rst
deleted file mode 100644
index 6418f32f6fe6d78267a373150fe7a2257c5e0b97..0000000000000000000000000000000000000000
--- a/docs/source/reference/api/codegen.rst
+++ /dev/null
@@ -1,29 +0,0 @@
-Code Generator and Configuration
-================================
-
-.. module:: pystencils.kernelcreation
-
-.. autosummary::
-  :toctree: generated
-  :nosignatures:
-
-  create_kernel
-
-.. module:: pystencils.config
-
-.. autosummary::
-  :toctree: generated
-  :nosignatures:
-  :template: autosummary/entire_class.rst
-
-  CreateKernelConfig
-  CpuOptimConfig
-  OpenMpConfig
-  VectorizationConfig
-  GpuIndexingConfig
-
-.. autosummary::
-  :toctree: generated
-  :nosignatures:
-
-  AUTO
\ No newline at end of file
diff --git a/docs/source/reference/api/index.rst b/docs/source/reference/api/index.rst
deleted file mode 100644
index b19c6303eb578f4b621febe507519ae3822df20a..0000000000000000000000000000000000000000
--- a/docs/source/reference/api/index.rst
+++ /dev/null
@@ -1,13 +0,0 @@
-***
-API
-***
-
-Modules
-=======
-
-.. toctree::
-    :maxdepth: 1
-
-    field
-    sympyextensions
-    codegen
diff --git a/docs/source/reference/gpu_kernels.md b/docs/source/reference/gpu_kernels.md
index 7fb0febb46d59988787e1c506640e27b35e1ece9..786840d182b0e06d4e26085cf6a95dbcb31d16b2 100644
--- a/docs/source/reference/gpu_kernels.md
+++ b/docs/source/reference/gpu_kernels.md
@@ -49,9 +49,9 @@ ps.inspect(kernel)
 ```
 
 The `kernel` object returned by the code generator in above snippet is an instance
-of the {py:class}`GpuKernelFunction` class.
-It extends {py:class}`KernelFunction` with some GPU-specific information.
-In particular, it defines the {any}`threads_range <GpuKernelFunction.threads_range>`
+of the {py:class}`GpuKernel` class.
+It extends {py:class}`Kernel` with some GPU-specific information.
+In particular, it defines the {any}`threads_range <GpuKernel.threads_range>`
 property, which tells us how many threads the kernel is expecting to be executed with:
 
 ```{code-cell} ipython3
@@ -208,12 +208,10 @@ only a part of the triangle is being processed.
 
 ```{eval-rst}
 .. autosummary::
-  :toctree: generated
   :nosignatures:
-  :template: autosummary/recursive_class.rst
 
-  pystencils.backend.kernelfunction.GpuKernelFunction
-  pystencils.backend.jit.gpu_cupy.CupyKernelWrapper
+  pystencils.codegen.GpuKernel
+  pystencils.jit.gpu_cupy.CupyKernelWrapper
 ```
 
 :::{admonition} Developers To Do:
diff --git a/docs/source/reference/kernelcreation.md b/docs/source/reference/kernelcreation.md
index 2a504578535604e0f4f28746a3b9eb3d7765de21..248855fc1c755d9fee4c62c7db07d469d3ac84ed 100644
--- a/docs/source/reference/kernelcreation.md
+++ b/docs/source/reference/kernelcreation.md
@@ -34,17 +34,19 @@ and their effects on the generated kernel.
 
 ## Running the Code Generator
 
-The primary way to invoke the code generation engine is through the `create_kernel` function.
+The primary way to invoke the code generation engine is through the {any}`create_kernel` function.
 It takes two arguments:
 - the list of assignment that make up the kernel (optionally wrapped as an ``AssignmentCollection``),
 - and a configuration object, an instance of {any}`CreateKernelConfig <pystencils.codegen.config.CreateKernelConfig>`.
 
 ```{eval-rst}
+.. currentmodule:: pystencils.codegen
+
 .. autosummary::
   :nosignatures:
 
-  pystencils.kernelcreation.create_kernel
-  pystencils.config.CreateKernelConfig
+  create_kernel
+  CreateKernelConfig
 ```
 
 For a simple kernel, an invocation of the code generator might look like this:
@@ -110,21 +112,14 @@ their interaction and effects, use cases and caveats.
 Pystencils supports code generation for a variety of CPU and GPU hardware.
 
 ```{eval-rst}
-.. currentmodule:: pystencils.codegen.config
+.. currentmodule:: pystencils.codegen
 
 .. autosummary::
   :nosignatures:
 
   CreateKernelConfig.target
-
-.. module:: pystencils.codegen.target
-
-.. autosummary::
-  :toctree: generated
-  :nosignatures:
-  :template: autosummary/recursive_class.rst
-
   Target
+
 ```
 
 ### Data Types
@@ -176,7 +171,7 @@ are using the `int32` data type, as specified in {py:data}`index_dtype <CreateKe
 ```{code-cell} ipython3
 :tags: [remove-input]
 
-driver = ps.kernelcreation.get_driver(cfg, retain_intermediates=True)
+driver = ps.codegen.get_driver(cfg, retain_intermediates=True)
 kernel = driver(assignments)
 ps.inspect(driver.intermediates.materialized_ispace, show_cpp=False)
 ```
@@ -186,7 +181,7 @@ To learn more about inspecting code after different stages of the code generator
 :::
 
 ```{eval-rst}
-.. currentmodule:: pystencils.codegen.config
+.. currentmodule:: pystencils.codegen
 
 .. autosummary::
   :nosignatures:
@@ -220,7 +215,7 @@ only one of which can be specified at a time:
 :::
 
 ```{eval-rst}
-.. currentmodule:: pystencils.codegen.config
+.. currentmodule:: pystencils.codegen
 
 .. autosummary::
   :nosignatures:
@@ -281,7 +276,7 @@ To illustrate, the following kernel accesses neighbor nodes with a maximum offse
 ```{code-cell} ipython3
 ranged_update = ps.Assignment(u.center(), v[-2, -1] + v[2, 1])
 
-cfg = ps.CreateKernelConfig(ghost_layers=ps.config.AUTO)
+cfg = ps.CreateKernelConfig(ghost_layers=ps.AUTO)
 kernel = ps.create_kernel(ranged_update, cfg)
 ```
 
@@ -510,18 +505,3 @@ driver = ps.codegen.get_driver(cfg, retain_intermediates=True)
 kernel = driver(assignments)
 ps.inspect(driver.intermediates)
 ```
-
-## API: Kernel Parameters and Function Objects
-
-```{eval-rst}
-.. module:: pystencils.codegen
-
-.. autosummary::
-  :toctree: generated
-  :nosignatures:
-  :template: autosummary/entire_class.rst
-
-  Parameter
-  Kernel
-  GpuKernel
-```
diff --git a/src/pystencils/jit/__init__.py b/src/pystencils/jit/__init__.py
index a47dc4aa6e1e0be03949c9fb48854ae385a41526..1ef8378d3000e95b12bb6a3a17062fb6488e1729 100644
--- a/src/pystencils/jit/__init__.py
+++ b/src/pystencils/jit/__init__.py
@@ -24,7 +24,7 @@ It is due to be replaced in the near future.
 
 from .jit import JitBase, NoJit, KernelWrapper
 from .legacy_cpu import LegacyCpuJit
-from .gpu_cupy import CupyJit
+from .gpu_cupy import CupyJit, CupyKernelWrapper, LaunchGrid
 
 no_jit = NoJit()
 """Disables just-in-time compilation for a kernel."""
@@ -36,4 +36,6 @@ __all__ = [
     "NoJit",
     "no_jit",
     "CupyJit",
+    "CupyKernelWrapper",
+    "LaunchGrid"
 ]