Skip to content

Use HGQ's trained softmax lookup tables in the generated HLS - #1531

Open
JanFSchulte wants to merge 2 commits into
fastmachinelearning:mainfrom
JanFSchulte:hgq2-softmax-luts
Open

Use HGQ's trained softmax lookup tables in the generated HLS#1531
JanFSchulte wants to merge 2 commits into
fastmachinelearning:mainfrom
JanFSchulte:hgq2-softmax-luts

Conversation

@JanFSchulte

@JanFSchulte JanFSchulte commented Aug 24, 2026

Copy link
Copy Markdown
Contributor

Description

hgq.layers.QSoftmax trains two lookup tables, exp_table and inv_table. QSoftmaxHandler
propagated only their sizes and types; the generated C++ rebuilt the contents at runtime
from std::exp and 1/x (nnet::init_exp_table / nnet::init_invert_table). Those
reconstructions are not guaranteed to reproduce the values HGQ trained, and the table was
addressed with a bit-slice whose width came from a different quantity than the address word.

The tables cannot be built at conversion time. For implementation='latency' their domain is the
softmax input precision, which only the bit_exact pass decides. QSoftmaxHandler therefore
stashes two domain-parameterised builders on the layer, and a new Vivado pass,
materialize_softmax_tables, calls them once the address type is final and turns the result into
weight arrays. It runs in the vivado:optimize flow next to fix_softmax_table_size, after
bit_exact and before transform_types.

The pass also forces exp_table_size == 2**width of the address type. softmax_idx_from_real_val
slices the top ceillog2(table_size) bits of the address word, so when the two disagree the table
is read at each bin's lower edge, or the slice runs off a narrower word. This is an independent
defect from the table contents and is a plausible larger contributor to the error reported in the
issue — see "Numerical behaviour" below.

The kernels gain softmax_lut / softmax_multidim_lut entry points taking the two tables as
arguments. softmax_latency / softmax_stable were split into an _impl body plus the existing
runtime-initialising entry point, which is unchanged, so plain Keras and QKeras softmax generate
exactly the code they did before. Separate names rather than overloads of softmax, so that
softmax_multidim's #pragma HLS allocation instances = softmax<CONFIG_T> still names a single
function.

The same issue reports invalid types such as ap_ufixed<2,32>. That spelling is bit_exact's
deliberate sentinel for B < 1, not itself a bug, but it is reached through a real one: HGQ's
heterogeneous quantizers can train a channel to a negative total bitwidth. extract_fixed_quantizer_config
now clamps B to 0, which is the constant-zero encoding generate_mask_fn already understands,
rather than to 1, which would bring a pruned channel back as a live 1-bit one. generate_mask_fn's
if b == 0 is widened to b <= 0 so the other four frontends that build mask_kbi are covered too.

Fixes #1523

Type of change

  • Bug fix (non-breaking)
  • New feature (non-breaking)
  • Breaking change (existing configs, APIs or generated code behave differently)
  • New frontend / backend
  • New layer / operator support
  • New configuration option — a new io_type, Strategy, or config attribute
  • Research paper implementation
  • Documentation
  • Build, CI or tooling
  • Refactor / cleanup (no functional change)

Affected areas

Backends:

  • Vivado, Vitis (VivadoAccelerator inherits vivado:optimize and so also gets the pass; not exercised)
  • Backend-independent (core IR, optimizer, hls4ml.model)

Nothing under hls4ml/model/ changes. Quartus, oneAPI, Catapult and Libero are untouched: the pass
is registered under the Vivado backend only, and their softmax kernels keep building tables at
runtime exactly as before.

Frontends:

  • Keras v3 (HGQ2 handlers only)
  • Not related to a frontend

Components:

  • IR (layers + model graph) / optimizer passes
  • C++ templates / HLS sources under hls4ml/templates/
  • Profiling, reporting or the CLI
  • Packaging, build or CI (one entry in generate_ci_yaml.py)

Configurations affected and exercised

The axis this change is keyed on is the softmax implementation, not Strategy.

Backend io_type implementation Verified
Vivado io_parallel stable, latency pytest + csim, 1D and multidim inputs
Vivado io_stream stable, latency pytest + csim, 1D and multidim inputs
Vitis io_parallel stable, latency pytest + csim, 1D and multidim inputs
Vitis io_stream stable, latency pytest + csim, 1D and multidim inputs
Vivado, Vitis both argmax, legacy pass returns early; covered by the existing test_softmax.py
VivadoAccelerator not run

Deliberately left out: no synthesis run of any configuration (see below), and no HGQ2 softmax on
Quartus / oneAPI / Catapult / Libero, which the HGQ2 frontend is not tested on today
(test_hgq_layers.py, test_hgq_players.py and test_hgq2_mha.py are all Vivado/Vitis only).

New configuration axis introduced by this PR: none.

Impact on generated HLS

  • This PR does not change the generated HLS for existing models.
  • This PR changes the generated HLS. Numbers below.

The change is confined to HGQ2 QSoftmax. For every other softmax the emitted call, the config
struct and the kernel body are byte-identical to main; test_hgq2_softmax.py carries an explicit
negative control asserting a plain keras.layers.Softmax still gets the two-argument call.

For HGQ2 QSoftmax the call gains two arguments and the project gains two weight headers
(exp_table<n>.h, inv_table<n>.h), while the runtime init_*_table loops disappear from that
layer.

Numerical behaviour: intentionally changed. After this PR, hls_model.predict on the test
models is bit-exact with Keras — np.testing.assert_array_equal, not assert_allclose — for all
16 covered combinations.

One caveat worth recording, since it bears on how the issue is attributed. Before implementing this
I compared HGQ's trained tables against what init_exp_table / init_invert_table reconstruct for
default QSoftmax configurations, and they matched bit for bit (the C++ assigns through
exp_table_t, which already carries HGQ's RND_CONV/SAT). On those models the reported error
therefore came from the address-width mismatch, not from the table values. Both are fixed here, and
carrying the trained tables makes the result independent of exp_scale and of the generic
reconstruction; but I have not isolated the two effects on the reporter's model.

Model / test Backend & version Part Latency (cycles) II LUT FF DSP BRAM
before not run
after not run

No synthesis was run. C simulation proves the arithmetic; it cannot see two things that need a
real build: whether #pragma HLS allocation instances = softmax_lut<CONFIG_T> resolves as intended
in softmax_multidim_lut, and the QoR effect of function_instantiate on the table arguments.
Both are flagged here rather than assumed.

Tests

New file test/pytest/test_hgq2_softmax.py, 19 tests:

  • test_hgq2_softmax_uses_trained_tables, parametrised over backend × io_type × stable ×
    input shape (16 cases). Asserts the weights exist and are sized 2**addr_width; that they carry
    the NamedType bit_exact derived rather than one re-inferred from the data (this catches
    AttributeDict.__setitem__ silently overwriting <name>_t when a WeightVariable is stored);
    that the contents equal an independent recomputation straight from hgq and differ from a
    naive exp(x * exp_scale) reconstruction, so the test cannot pass vacuously; that the generated
    call names both arrays and the headers are written and included; and finally
    assert_array_equal(r_hls, r_keras).
  • test_plain_keras_softmax_keeps_runtime_tables — negative control, two-argument call preserved.
  • test_dead_channels_stay_zero — a heterogeneous quantizer channel forced to a negative total
    width stays b == 0 and renders as out[3] = 0;.

test_hgq2_softmax is added to KERAS3_LIST in generate_ci_yaml.py, so CI runs it under
.pytest-keras3-only alongside test_hgq2_mha. Without that it would be batched into a default
job, whose environment has Keras 2 and no hgq.

Commands run, all from test/pytest on this branch:

pytest test_hgq2_softmax.py -q -p no:randomly            # 19 passed
pytest test_softmax.py -q -p no:randomly -k "Vivado or Vitis"   # 116 passed
pytest test_sparsepixels.py test_bit_exact_zeropadding.py -q -p no:randomly
                                                          # 17 passed, 2 failed
pre-commit run --files <every file in this diff>          # all hooks passed
pre-commit run --hook-stage manual --all-files check-manifest
                                                          # only uninitialised submodules reported

The two failures are test_bit_exact_zeropadding.py::test_bit_exact_zeropadding{1,2}d[io_parallel-oneAPI].
They fail identically on an unmodified origin/main worktree here, so they are pre-existing and
unrelated — this PR does not touch oneAPI.

test_hgq_layers.py and test_hgq_players.py, which cover the QUnaryLUTHandler refactor, could
not be run: HGQ v1 is not installed in my environment (ModuleNotFoundError: No module named 'HGQ'). Instead I checked the refactor directly, running the extracted helpers and the verbatim
pre-refactor algorithm side by side on QUnaryFunctionLUT layers with relu, sigmoid, tanh and exp
activations; the tables are byte-identical. That refactor also fixes a latent bug it inherited: the
old code called layer.oq(table[None, ...]), which fails on any layer built for rank > 2 — which
is exactly what QSoftmax's sublayers are for a multidimensional input.

Test configuration: Linux, Python 3.10, Keras 3 on the TensorFlow backend, hgq 0.1.8. No HLS
toolchain — compile() builds against the bundled ap_types/ with g++.

AI assistance disclosure

  • None — no AI tool was used.
  • Assisted — completion, refactoring, docstrings, tests; design and code are mine.
  • Substantial — significant AI-generated portions, reviewed and edited by me.
  • Agentic — produced largely end-to-end by an AI agent from my prompts.

Tool(s) and model(s): Claude Code (Claude Opus 5)

Where it was used: all files in this diff — the optimizer pass, the converter changes, the kernel
split, and the tests — from prompts describing the issue.

If anything other than None is ticked, confirm all of the following:

  • I am the author of this contribution and take full responsibility for it. ...
  • I verified the generated code against real hls4ml and HLS semantics ...
  • All numbers, logs and test results quoted in this PR come from runs I actually performed ...
  • I have the right to submit this work under the project's licence ...
  • No AI tool is credited as an author in any commit in this branch.

Checklist

Required:

  • I have read the contributing guidelines.
  • I installed and ran pre-commit on the files I edited.
  • I added tests under test/pytest covering this change.
  • I self-reviewed the full diff ...
  • The AI assistance disclosure above is complete and accurate.

If applicable:

  • Documentation under docs/ updated. — not applicable; this is a bug fix with no new
    user-facing option, and no docs/ page describes the softmax table construction.
  • No new build or synthesis warnings. — cannot confirm, no synthesis was run.
  • Public API / config schema changes are documented. — none.

Release note

HGQ2 QSoftmax layers now emit the lookup tables trained by HGQ instead of rebuilding them at
runtime in the generated C++, making Vivado and Vitis output bit-exact with Keras.

🤖 Generated with Claude Code

hgq.layers.QSoftmax trains two lookup tables, exp_table and inv_table.
QSoftmaxHandler propagated only their sizes and types, leaving the generated
C++ to rebuild the contents at runtime from std::exp and 1/x, which is not
guaranteed to reproduce the values HGQ trained.

The tables cannot be built at conversion time: for the latency implementation
their domain is the softmax input precision, which only bit_exact decides. The
converter therefore stashes two domain-parameterised builders on the layer, and
a new Vivado pass, materialize_softmax_tables, calls them once the address type
is final and turns the result into weight arrays. It also forces the table size
to 2**width of that address type, so that the top-bits slice
softmax_idx_from_real_val performs covers the whole word.

The kernels gain softmax_lut / softmax_multidim_lut entry points taking the two
tables as arguments; the existing runtime-initialised entry points are
unchanged, so plain Keras and QKeras softmax generate exactly the code they did
before.

Also fixes the negative total bitwidths reported in the same issue: a channel
trained below zero bits is clamped to zero, which is the constant-zero encoding
generate_mask_fn already understands, rather than to one, which would bring the
channel back as a live 1-bit one.

Fixes fastmachinelearning#1523
@JanFSchulte

Copy link
Copy Markdown
Contributor Author

I made this mostly as an example to test @vloncar's agent instructions from #1527, including the new PR template. At first glance, I think Claude did a decent job following the instructions. Still a bit verbose in the comments, but harmless compared to other things I have seen. It did flag a few things it would have fixed on the fly, and acknowledged that it was instructed to stay within the scope of the requested fix. So that part of the instructions seem to work well.

As for the fix itself, it seems a bit clunky to me, but does work, so we can discuss if we want to actually merge it or come up with something better.

@JanFSchulte JanFSchulte added bugfix please test Trigger testing by creating local PR branch labels Aug 24, 2026
@JanFSchulte
JanFSchulte marked this pull request as ready for review August 28, 2026 15:54
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bugfix please test Trigger testing by creating local PR branch

Projects

None yet

Development

Successfully merging this pull request may close these issues.

hls4ml HGQ2 issues: softmax LUT mismatch + negative bitwidths

1 participant