FluxML / FluxML/NNlib.jl

Upstream candidates from LuxLib

Open
#712 0 comments 1 reaction 0 assignees View on GitHub
Dominant language
Julia
Stars
258
Forks
138
Avg merge
2d 1h
Merged PRs (30d)
3

Description

Tracking issue for functionality in [`Lux.jl/lib/LuxLib`](https://github.com/LuxDL/Lux.jl/tree/main/lib/LuxLib) that could be upstreamed into NNlib so the wider ecosystem (Flux, etc.) can share it.

NNlib already depends on `KernelAbstractions`, `ChainRulesCore`, `GPUArraysCore`, `Atomix`, and has weakdeps on `CUDA`, `cuDNN`, `EnzymeCore`, `ForwardDiff`, `Metal`, `AMDGPU`, so the KA kernels and AD rules port cleanly. The optimized-CPU/vendor pieces can mostly live in package extensions.

## High value

- [x] ~**Normalization layers** — `batchnorm`, `groupnorm`, `instancenorm`, `layernorm`. Generic `normalization`/`affine_normalize` core + running-statistics machinery + KA CPU/GPU kernels, plus the cuDNN batchnorm fast path.~ We have our own generic implementations.

## Medium value — generalizes/enhances existing NNlib ops
- [ ] **Fused dense** — `fused_dense_bias_activation` (`σ.(W*x .+ b)` with buffer reuse), the cache-size-aware `matmul`/`matmuladd` polyalgorithm (BLAS / Octavian / LoopVectorization, with MKL/AppleAccelerate/BLISBLAS extensions), and the cuBLASLt fused-dense path.
- [ ] **`fast_activation` / `fast_activation!!`** — in-place-when-possible activation broadcast with an optimized CPU SIMD loop, `∇activation` gradient kernel, traits to skip caching the pre-activation when the derivative only needs the output, and SLEEFPirates Float32 fast paths.
- [ ] **Generalized `batched_matmul`** — arbitrary contracting/batching dims (NNlib's `batched_mul` is 3-D only), LoopVectorization CPU path, Reactant `dot_general` lowering.
- [ ] **Enhanced attention** — `scaled_dot_product_attention` with grouped-/multi-query attention (KV-group sharing), flexible `head_dim`/`token_dim`, `is_causal` kwarg, additive bias, Reactant `dot_general` lowering. NNlib's `dot_product_attention` lacks GQA/MQA and flexible dim placement (cf. #572, #699).
- [ ] **`bias_activation` / `bias_activation!!`** — N-D bias-add + activation with multi-AD fallbacks; generalizes `bias_act!` (candidate to consolidate).
- [ ] **`fused_conv_bias_activation`** — cuDNN fused path for `identity`/`relu`, Julia-side fusion otherwise, mixed-precision handling; overlaps/extends `conv_bias_act`.

## Infrastructure / lower value

- [ ] **`alpha_dropout`**, plus `dropout` with explicit mask/`update_mask` and `training` auto-detection from the AD context.
- [ ] **`internal_operation_mode` + traits dispatch** (`GenericBroadcastOp`/`GPUBroadcastOp`/`LoopedArrayOp` selector and the supporting traits) — prerequisite for the optimized kernels; decide whether to adopt or slim down.
- [ ] **CPU cache helpers** (`fits_in_l{1,2,3}cache`, `use_octavian`, hardware detection) needed by the matmul polyalg.
- [ ] **Cross-AD glue** — hand-written Enzyme rules on the mutating kernels + Tracker/ReverseDiff/ForwardDiff extensions.

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.