Use HGQ's trained softmax lookup tables in the generated HLS - #1531
Open
JanFSchulte wants to merge 2 commits into
Open
Use HGQ's trained softmax lookup tables in the generated HLS#1531JanFSchulte wants to merge 2 commits into
JanFSchulte wants to merge 2 commits into
Conversation
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
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
marked this pull request as ready for review
August 28, 2026 15:54
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
hgq.layers.QSoftmaxtrains two lookup tables,exp_tableandinv_table.QSoftmaxHandlerpropagated only their sizes and types; the generated C++ rebuilt the contents at runtime
from
std::expand1/x(nnet::init_exp_table/nnet::init_invert_table). Thosereconstructions 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 thesoftmax input precision, which only the
bit_exactpass decides.QSoftmaxHandlerthereforestashes 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 intoweight arrays. It runs in the
vivado:optimizeflow next tofix_softmax_table_size, afterbit_exactand beforetransform_types.The pass also forces
exp_table_size == 2**widthof the address type.softmax_idx_from_real_valslices the top
ceillog2(table_size)bits of the address word, so when the two disagree the tableis 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_lutentry points taking the two tables asarguments.
softmax_latency/softmax_stablewere split into an_implbody plus the existingruntime-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 thatsoftmax_multidim's#pragma HLS allocation instances = softmax<CONFIG_T>still names a singlefunction.
The same issue reports invalid types such as
ap_ufixed<2,32>. That spelling isbit_exact'sdeliberate sentinel for
B < 1, not itself a bug, but it is reached through a real one: HGQ'sheterogeneous quantizers can train a channel to a negative total bitwidth.
extract_fixed_quantizer_confignow clamps
Bto 0, which is the constant-zero encodinggenerate_mask_fnalready understands,rather than to 1, which would bring a pruned channel back as a live 1-bit one.
generate_mask_fn'sif b == 0is widened tob <= 0so the other four frontends that buildmask_kbiare covered too.Fixes #1523
Type of change
io_type, Strategy, or config attributeAffected areas
Backends:
vivado:optimizeand so also gets the pass; not exercised)hls4ml.model)Nothing under
hls4ml/model/changes. Quartus, oneAPI, Catapult and Libero are untouched: the passis registered under the Vivado backend only, and their softmax kernels keep building tables at
runtime exactly as before.
Frontends:
Components:
hls4ml/templates/generate_ci_yaml.py)Configurations affected and exercised
The axis this change is keyed on is the softmax
implementation, not Strategy.io_typeimplementationtest_softmax.pyDeliberately 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.pyandtest_hgq2_mha.pyare all Vivado/Vitis only).New configuration axis introduced by this PR: none.
Impact on generated HLS
The change is confined to HGQ2
QSoftmax. For every other softmax the emitted call, the configstruct and the kernel body are byte-identical to
main;test_hgq2_softmax.pycarries an explicitnegative control asserting a plain
keras.layers.Softmaxstill gets the two-argument call.For HGQ2
QSoftmaxthe call gains two arguments and the project gains two weight headers(
exp_table<n>.h,inv_table<n>.h), while the runtimeinit_*_tableloops disappear from thatlayer.
Numerical behaviour: intentionally changed. After this PR,
hls_model.predicton the testmodels is bit-exact with Keras —
np.testing.assert_array_equal, notassert_allclose— for all16 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_tablereconstruct fordefault
QSoftmaxconfigurations, and they matched bit for bit (the C++ assigns throughexp_table_t, which already carries HGQ'sRND_CONV/SAT). On those models the reported errortherefore 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_scaleand of the genericreconstruction; but I have not isolated the two effects on the reporter's model.
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 intendedin
softmax_multidim_lut, and the QoR effect offunction_instantiateon 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 carrythe
NamedTypebit_exactderived rather than one re-inferred from the data (this catchesAttributeDict.__setitem__silently overwriting<name>_twhen aWeightVariableis stored);that the contents equal an independent recomputation straight from
hgqand differ from anaive
exp(x * exp_scale)reconstruction, so the test cannot pass vacuously; that the generatedcall 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 totalwidth stays
b == 0and renders asout[3] = 0;.test_hgq2_softmaxis added toKERAS3_LISTingenerate_ci_yaml.py, so CI runs it under.pytest-keras3-onlyalongsidetest_hgq2_mha. Without that it would be batched into a defaultjob, whose environment has Keras 2 and no
hgq.Commands run, all from
test/pyteston this branch: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/mainworktree here, so they are pre-existing andunrelated — this PR does not touch oneAPI.
test_hgq_layers.pyandtest_hgq_players.py, which cover theQUnaryLUTHandlerrefactor, couldnot 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 verbatimpre-refactor algorithm side by side on
QUnaryFunctionLUTlayers with relu, sigmoid, tanh and expactivations; 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 — whichis exactly what
QSoftmax's sublayers are for a multidimensional input.Test configuration: Linux, Python 3.10, Keras 3 on the TensorFlow backend,
hgq0.1.8. No HLStoolchain —
compile()builds against the bundledap_types/with g++.AI assistance disclosure
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:
Checklist
Required:
pre-commiton the files I edited.test/pytestcovering this change.If applicable:
docs/updated. — not applicable; this is a bug fix with no newuser-facing option, and no
docs/page describes the softmax table construction.Release note
🤖 Generated with Claude Code