diff --git a/docs/documentation/guides.md b/docs/documentation/guides.md deleted file mode 100644 index af23bec09..000000000 --- a/docs/documentation/guides.md +++ /dev/null @@ -1,14 +0,0 @@ ---- -title: Guides -subtitle: Practical guides for profiling, optimizing, scaling and third-party tool integrations in JAX. -slug: guides ---- - -- [**Writing High-Performance CuTe DSL kernels in JAX**](https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html) -- [**Tips for High-Performance LLMs with JAX on GPUs**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/GPU_performance.md) -- [**AutoPGLE on GPU workflows**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/PGLE.md) -- [**Profiling JAX programs on GPU**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/profiling.md) - - [**`nsys-jax`**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/nsys-jax.md) -- [**Resilient Training with Ray and JAX**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/resiliency/ray_resilient_jax.md) -- [**User Guide on Native XLA-FP8**](https://github.com/NVIDIA/JAX-Toolbox/blob/main/docs/NATIVE_FP8.md) - diff --git a/docs/gpu-kernels.md b/docs/gpu-kernels.md new file mode 100644 index 000000000..31937c231 --- /dev/null +++ b/docs/gpu-kernels.md @@ -0,0 +1,11 @@ +--- +title: Customize JAX with GPU kernels +subtitle: Writing custom kernels for JAX on GPUs. +slug: gpu-kernels +--- + +Developers need customization of the JAX stack to customize their models and to optimize beyond what's natively available in JAX/XLA stack. + +NVIDIA offers a suite of tools, libraries and kernel DSLs, helping customization. + +- [**Writing High-Performance CuTe DSL kernels in JAX**](https://docs.jax.dev/en/latest/notebooks/cute_dsl_jax.html) diff --git a/docs/index.yml b/docs/index.yml index 696626a8e..ccfac3e37 100644 --- a/docs/index.yml +++ b/docs/index.yml @@ -29,11 +29,6 @@ navigation: - page: Environment Variables path: getting-started/environment-variables.md - # ==================== Documentation ==================== - - section: Documentation - contents: - - page: Guides - path: documentation/guides.md # ==================== Frameworks ==================== - section: Frameworks contents: @@ -74,11 +69,11 @@ navigation: - page: JAX-vLLM Rollout Offloading Bridge path: inference/jax-inference-offloading.md - # ==================== Tools ==================== - - section: Tools + # ==================== Customization ==================== + - section: Customization contents: - - page: Triage Tool - path: triage-tool.md + - page: GPU Kernels + path: gpu-kernels.md # ==================== Reference ==================== - section: Reference