docs: update softmax backend guidance and dim contract

Author: debashishc

CHANGELOG.md

@@ -26,7 +26,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Documentation: README, DEVELOPMENT.md (#11), CONTRIBUTING.md with contribution policy (#7)
 - Pre-commit hooks installation instructions in CONTRIBUTING.md (#7)
 - Support for float16, bfloat16, and float32 dtypes
-- `FORGE_SOFTMAX_IMPL` softmax mode selection (`auto`, `ref`, `kernel`)
+- `softmax_online` backend registry APIs (`register/get/set/list`) with `ref` and `kernel` backends
 
 ### Changed
 - Modal benchmarks now use strict GPU matching (`!` suffix) for consistent hardware
@@ -38,7 +38,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Enhanced error messages with device information for better debugging
 - Set tolerance to 0 (exact comparison) for transpose correctness tests
 - Updated documentation to reflect current three-layer architecture and modern PyTorch patterns
-- Added `--impl` to `bench/benchmark_online_softmax.py` and updated benchmark handling for strict kernel mode
+- Added `--backend` to `bench/benchmark_online_softmax.py` and aligned benchmark handling with backend registry
+- Enforced strict `dim=-1` contract for `softmax_online` (row-wise only)
 
 ### Infrastructure
 - Package scaffolding with `pyproject.toml` and proper Python 3.13+ support

DEVELOPMENT.md

@@ -41,7 +41,7 @@ This enables usage via both `forge_cute_py.ops.op_name()` and
 | `copy_transpose` | CuTe DSL | Fully implemented with tile-based shared memory |
 | `reduce` | CuTe DSL | Placeholder (sum only) |
 | `reduce_sum` | Reference | Placeholder (awaiting benchmarking) |
-| `softmax_online` | Reference | Stub with autograd support (kernel TODO, mode-gated via `FORGE_SOFTMAX_IMPL`) |
+| `softmax_online` | Reference + CuTe backend | Backend-registry flow (`ref` default, `kernel` optional); row-wise `dim=-1` only |
 
 ### Development flow
 
@@ -128,16 +128,20 @@ non-zero tolerance in tests.
 uv run python bench/run.py --suite smoke
 uv run python bench/benchmark_copy_transpose.py --tile-size 16
 uv run python bench/benchmark_reduce.py
-uv run python bench/benchmark_online_softmax.py --impl auto
-uv run python bench/benchmark_online_softmax.py --impl kernel
+uv run python bench/benchmark_online_softmax.py --backend ref
+uv run python bench/benchmark_online_softmax.py --backend kernel
 modal run bench/modal_bench.py --suite smoke --out results.json
 modal run bench/modal_bench.py --suite smoke --op reduce_sum --out results.json
 ```
 
-`softmax_online` backend mode is controlled by `FORGE_SOFTMAX_IMPL`:
-- `auto` (default): try kernel if present, otherwise fallback to reference.
-- `ref`: force reference path.
-- `kernel`: require kernel path and fail fast if missing/incomplete.
+`softmax_online` backend selection is controlled by:
+- Python API: `set_softmax_online_backend("ref")` / `set_softmax_online_backend("kernel")`
+- Benchmark CLI: `bench/benchmark_online_softmax.py --backend {ref,kernel}`
+
+Current `softmax_online` constraints:
+- Input is 2D CUDA tensor
+- Dtype in `float16`, `bfloat16`, `float32`
+- `dim=-1` only
 
 > **Warning:** Modal benchmarks incur GPU costs. Review `bench/modal_bench.py`
 > and verify timeout/GPU settings before running. Start with `--suite smoke`

README.md

@@ -84,8 +84,8 @@ Run a standalone benchmark:
 
 ```bash
 uv run python bench/benchmark_copy_transpose.py --tile-size 16
-uv run python bench/benchmark_online_softmax.py --impl auto
-uv run python bench/benchmark_online_softmax.py --impl kernel
+uv run python bench/benchmark_online_softmax.py --backend ref
+uv run python bench/benchmark_online_softmax.py --backend kernel
 ```
 
 Profile a kernel with helper script:
@@ -110,7 +110,7 @@ ncu --set full -o profiles/copy_transpose uv run python bench/benchmark_copy_tra
 | copy_transpose | Implemented | tile_size=16/32 | CuTe DSL kernel with tiled shared memory |
 | reduce | Implemented (sum only) | - | Placeholder (awaiting benchmarking) |
 | reduce_sum | Stub (ref) | - | Placeholder (awaiting benchmarking) |
-| softmax_online | Stub (ref) | single-pass | Reference-backed by default; `FORGE_SOFTMAX_IMPL` can force `ref` or strict `kernel` mode |
+| softmax_online | Implemented (registry) | ref, kernel | Default backend is `ref`; `kernel` is available for contributor benchmarking and remains non-production |
 
 ---
 
@@ -153,15 +153,19 @@ uv run python bench/run.py --suite smoke              # Run benchmark suite
 uv run python bench/run.py --suite smoke --out out.json  # Save results
 uv run python bench/benchmark_copy_transpose.py       # Standalone benchmark
 uv run python bench/benchmark_reduce.py               # Standalone benchmark
-uv run python bench/benchmark_online_softmax.py --impl auto    # softmax fwd+bwd (fallback allowed)
-uv run python bench/benchmark_online_softmax.py --impl kernel  # softmax fwd+bwd (hard fail if kernel missing)
+uv run python bench/benchmark_online_softmax.py --backend ref     # softmax fwd+bwd on reference backend
+uv run python bench/benchmark_online_softmax.py --backend kernel  # softmax fwd+bwd on CuTe kernel backend
 modal run bench/modal_bench.py --suite smoke --out results.json  # Remote run on B200 via Modal
 ```
 
-`softmax_online` mode is controlled by `FORGE_SOFTMAX_IMPL`:
-- `auto` (default): try kernel first, then fallback to reference.
-- `ref`: force reference implementation.
-- `kernel`: require kernel implementation and fail fast if unavailable.
+`softmax_online` backend selection is controlled via Python API or benchmark CLI:
+- Python API: `set_softmax_online_backend("ref")` or `set_softmax_online_backend("kernel")`
+- Benchmark CLI: `bench/benchmark_online_softmax.py --backend {ref,kernel}`
+
+Current softmax contract:
+- supported tensors: 2D CUDA
+- supported dtypes: `float16`, `bfloat16`, `float32`
+- supported dim: `dim=-1` (row-wise only)
 
 ### Modal setup (remote benchmarks)
 ```bash