⚡ Thunderbolt: softmax_v6 — 8x unrolled AVX2 single-FMA exp256

[bugparty]

💡 What: Added softmax_v6 and exp256_ps_v3 to ml_kernels/include/ml_kernels/softmax.h. softmax_v6 unrolls the softmax loops 8x (64 elements per iteration). exp256_ps_v3 utilizes a single FMA for the r = x - n * ln(2) operation rather than splitting ln(2).

🎯 Why: To shift the bottleneck away from execution port contention/instruction latency and directly towards L1/L2 cache bandwidth limits by fully utilizing all 16 YMM registers.

🏗️ How: softmax_v5 unrolled 4x. softmax_v6 aggressively unrolls 8x and merges the precision-split subtraction in range reduction into a single _mm256_fnmadd_ps. The single FMA sacrifices ~1e-7 precision (still well within the 1e-4 tolerance of softmax ML kernels) to eliminate an instruction dependency.

📊 Impact:

• In Fixed Memory Mode (N=262144): Throughput increased from ~4.37 GFLOP/s (softmax_v5) to ~4.91 GFLOP/s (softmax_v6), a ~12% speedup.
• In Pool Mode (N=262144): Throughput increased from ~2.69 GFLOP/s (softmax_v5) to ~3.03 GFLOP/s (softmax_v6), a ~12% speedup.

🖥️ Tested on: Target: AVX2 (Haswell+). Environment: GCC 13.3.0.

🔬 How to reproduce:

DISABLE_CPU_BINDING=1 ./build/ml_kernels/ml_kernel_bench --filter "softmax_v[56]" --sizes 262144 --iters 200

---

PR created automatically by Jules for task 7534808963543397349 started by @bugparty

Summary by CodeRabbit

• New Features

  • Added a faster softmax implementation with improved AVX2 performance.
  • Expanded benchmark coverage to include the new softmax variant.

• Documentation

  • Added notes summarizing recent softmax optimization findings and benchmark results.

• Tests

  • Added validation to confirm the new softmax output matches the reference implementation and still sums to 1.

[google-labs-jules]

👋 Jules, reporting for duty! I'm here to lend a hand with this pull request.

When you start a review, I'll add a 👀 emoji to each comment to let you know I've read it. I'll focus on feedback directed at me and will do my best to stay out of conversations between you and other bots or reviewers to keep the noise down.

I'll push a commit with your requested changes shortly after. Please note there might be a delay between these steps, but rest assured I'm on the job!

For more direct control, you can switch me to Reactive Mode. When this mode is on, I will only act on comments where you specifically mention me with @jules. You can find this option in the Pull Request section of your global Jules UI settings. You can always switch back!

New to Jules? Learn more at jules.google/docs.

---

For security, I will only act on instructions from the user who triggered this task.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR adds an AVX2 exp approximation helper exp256_ps_v3 and a new softmax implementation softmax_v6 using 8x loop unrolling and single-FMA polynomial evaluation. It registers a corresponding benchmark, adds a correctness test against softmax_naive, and documents the optimization approach.

Changes

AVX2 softmax_v6 kernel addition

Layer / File(s)	Summary
exp256_ps_v3 and softmax_v6 implementation ml_kernels/include/ml_kernels/softmax.h	Adds exp256_ps_v3 (range clamping, base-2 reduction, single-FMA Horner polynomial, exponent reconstruction) and softmax_v6 (8x-unrolled max/exp/sum/normalize with scalar tail handling and early return on zero sum).
Benchmark registration and correctness test ml_kernels/src/kernel_bench.cpp, ml_kernels/src/test_naive_ops.cpp	Adds SoftmaxV6Benchmark calling softmax_v6 over pooled buffers, and test_softmax_v6 validating output against softmax_naive, wired into main().
Optimization notes documentation .jules/thunderbolt.md	Documents the single-FMA exp256 and unrolling approach with benchmark comparisons to softmax_v5.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Possibly related PRs

• bugparty/cpu_math_kernels_pri#31: Builds on the same softmax_v5/exp256_ps_v2 AVX2 range-reduction and Horner polynomial scheme extended here into exp256_ps_v3/softmax_v6.
• bugparty/cpu_math_kernels_pri#28: Adds analogous exp256 approximation helper and softmax variant (exp256_ps_estrin/softmax_v4) in the same file.
• bugparty/cpu_math_kernels_pri#7: Adds related test coverage comparing softmax implementations against softmax_naive.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title clearly matches the main change: a new 8x-unrolled AVX2 softmax implementation with a single-FMA exp approximation.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch thunderbolt-softmax-v6-7534808963543397349

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🧹 Nitpick comments (2)

ml_kernels/src/test_naive_ops.cpp (1)

217-218: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Consider adding test_softmax_v6() at the end of main() for consistency.

It's called first, before all existing tests. Not a bug, but placing it after test_softmax_v5() would be more consistent with the version-ordered progression of the other test calls.

♻️ Suggested ordering

 int main() {
-    test_softmax_v6();
     test_relu_naive();
     test_max_naive();
     test_softmax_v3();
     test_softmax_v4();
     test_softmax_v5();
+    test_softmax_v6();
     std::cout << "All tests passed successfully!" << std::endl;
 }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ml_kernels/src/test_naive_ops.cpp` around lines 217 - 218, Move the
test_softmax_v6() call in main() to the end of the existing softmax test
sequence so the versioned test calls stay in order. Keep the change localized in
main() in test_naive_ops.cpp, placing test_softmax_v6() after test_softmax_v5()
and before any later cleanup/return logic if present.

ml_kernels/include/ml_kernels/softmax.h (1)

505-505: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Brace placement violates the repo style rule.

Function-body opening braces sit on the signature line for both exp256_ps_v3 (Line 505) and softmax_v6 (Line 539). The rest of the file shares this pattern, but the guideline calls for braces on their own line.

As per coding guidelines: "Keep braces on their own lines for function bodies".

Also applies to: 539-539

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@ml_kernels/include/ml_kernels/softmax.h` at line 505, Move the function-body
opening braces for exp256_ps_v3 and softmax_v6 onto their own lines to match the
repo’s brace style. Update the function declarations so the opening brace is
separated from the signature for each of these functions, and keep the rest of
the body unchanged.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@ml_kernels/include/ml_kernels/softmax.h`:
- Line 505: Move the function-body opening braces for exp256_ps_v3 and
softmax_v6 onto their own lines to match the repo’s brace style. Update the
function declarations so the opening brace is separated from the signature for
each of these functions, and keep the rest of the body unchanged.

In `@ml_kernels/src/test_naive_ops.cpp`:
- Around line 217-218: Move the test_softmax_v6() call in main() to the end of
the existing softmax test sequence so the versioned test calls stay in order.
Keep the change localized in main() in test_naive_ops.cpp, placing
test_softmax_v6() after test_softmax_v5() and before any later cleanup/return
logic if present.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 4d6a629c-81a5-495e-85ad-eff90567b4a4

📥 Commits

Reviewing files that changed from the base of the PR and between acca01e and ce5dc0c.

📒 Files selected for processing (4)

• .jules/thunderbolt.md
• ml_kernels/include/ml_kernels/softmax.h
• ml_kernels/src/kernel_bench.cpp
• ml_kernels/src/test_naive_ops.cpp