⚡ Thunderbolt: Softmax — Single-FMA range reduction + 8x unrolling

[bugparty]

💡 What: Added softmax_v6 to the ML kernels library, implementing an aggressive 8x unrolled AVX2 softmax utilizing a streamlined exp256_ps_v3 approximation.
🎯 Why: Previous AVX2 versions (e.g., softmax_v5) split ln(2) for exact precision during range reduction, creating higher instruction counts and register pressure that constrained unroll factors. Consolidating to a single FMA operation frees enough architectural registers to push loop unrolling to 8x (64 elements/iteration).
🏗️ How:

1. exp256_ps_v3 was updated to calculate the reduction remainder via a single _mm256_fnmadd_ps.
2. softmax_v6 uses an 8x unroll strategy across all three passes (max, exp/sum, norm), perfectly saturating execution ports and hiding latency.
3. Fallback scalar processing handles remaining elements seamlessly.
   📊 Impact:

• softmax_v5 (N=16384, Fixed Memory): 5.06 GFLOP/s
• softmax_v6 (N=16384, Fixed Memory): 5.82 GFLOP/s
• Speedup: ~15% higher throughput.
  🖥️ Tested on: x86-64 CPU (AVX2-capable execution environment via Linux runner, GCC 13.3.0).
  🔬 How to reproduce:

cd build
make ml_kernel_bench
DISABLE_CPU_BINDING=1 ./ml_kernels/ml_kernel_bench --filter softmax_v6

---

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

Summary by CodeRabbit

• New Features

  • Added a new high-performance softmax implementation for AVX2/FMA-capable systems.
  • Added a new exponential routine used by the softmax path for faster vectorized math.

• Bug Fixes

  • Improved numerical handling for softmax calculations, including edge cases and zero-sum protection.
  • Added a benchmarkable softmax variant with better large-input processing.

• Tests

  • Added coverage comparing the new softmax output against the existing reference implementation and validating probability normalization.

[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

Adds an AVX2/FMA exp256_ps_v3 exponential routine and a softmax_v6 kernel using 8-way unrolled max/sum/normalize phases, registers a SoftmaxV6Benchmark, adds a test_softmax_v6 unit test validating correctness against softmax_naive, and documents the optimization rationale in .jules/thunderbolt.md.

Changes

softmax_v6 kernel

Layer / File(s)	Summary
exp256_ps_v3 and softmax_v6 implementation ml_kernels/include/ml_kernels/softmax.h	Adds a single-FMA range-reduction exponential (exp256_ps_v3) and an 8-way unrolled softmax_v6 kernel with max reduction, exp/sum accumulation, and reciprocal-sum normalization.
Benchmark registration and validation test ml_kernels/src/kernel_bench.cpp, ml_kernels/src/test_naive_ops.cpp, .jules/thunderbolt.md	Registers SoftmaxV6Benchmark dispatching to softmax_v6, adds test_softmax_v6 comparing outputs against softmax_naive, and documents the single-FMA/8x-unroll optimization note.

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

Possibly related PRs

• bugparty/cpu_math_kernels_pri#31: Adds prior exp256_ps_v*/softmax_v* (v2/v5) kernel variants to the same softmax.h file, part of the same optimization series.

🚥 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: an optimized softmax implementation with single-FMA range reduction and 8x unrolling.
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-7538711097839376193

---

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 (1)

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

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

Drop the per-function AVX2/FMA target attributes here
ml_kernels already builds with -march=native, and the other AVX2/FMA helpers in softmax.h don’t use per-function target("avx2,fma"). Remove it from exp256_ps_v3 and softmax_v6, or apply it consistently across the file.

🤖 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 504, The per-function
AVX2/FMA target attributes on exp256_ps_v3 and softmax_v6 are inconsistent with
the rest of softmax.h and redundant with the -march=native build setup. Remove
the target("avx2,fma") attribute from exp256_ps_v3 and softmax_v6, or if
specialization is required, apply the same attribute pattern consistently to the
other AVX2/FMA helpers in this file so the compilation strategy is uniform.

🤖 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 504: The per-function AVX2/FMA target attributes on exp256_ps_v3 and
softmax_v6 are inconsistent with the rest of softmax.h and redundant with the
-march=native build setup. Remove the target("avx2,fma") attribute from
exp256_ps_v3 and softmax_v6, or if specialization is required, apply the same
attribute pattern consistently to the other AVX2/FMA helpers in this file so the
compilation strategy is uniform.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f8252d59-409d-4fa6-bc11-e7dcf815879b

📥 Commits

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

📒 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