Продолжение red eyes is all you need

В общем, глобально было два пути оптимизации этого исчадия: через увеличение тактовой частоты железки и через уменьшение tokens per cycle. Первое в самом начале решалось простым раскидыванием регистров в нужных местах, но со временем я уперся в тот факт, что один большой BRAM модуль с весами слишком медленно доставляет данные из-за физического положения (веса/активации нужны многим модулям, и сам путь до нужного модуля занимает слишком много времени). Поэтому пришлось дробить на кучу блоков + конфигурировать все модули в свои pblocks.

Идея в том, что place&route по умолчанию раскидывает логику по кристаллу как ему удобнее, а удобнее ему обычно не там, где надо, и clock tree synthesis потом героически пытается развести тактовый сигнал через всю эту кашу. Если matvec u_qkv оказался в одном углу, а его weight_store в другом, то на трассы между ними уходят те самые наносекунды, из-за которых слайсится тайминг. Поэтому каждому matvec'у выделяется своя территория с примерно нужным количеством SLICE/DSP/BRAM, и его weight_store селится туда же.

Нижняя половина кристалла (X0..X145, Y0..Y149) поделена на четыре вертикальные полосы под четыре w32-matvec'а внутри transformer_layer: u_qkv в X0..X41, u_proj в X42..X69, u_ff_up в X70..X109, u_ff_down в X110..X145. Каждый pblock включает и сам matvec, и парный к нему weight_store, чтобы трассы от выхода weight_store (256 бит на u_qkv/u_proj/u_ff_up/u_ff_down, 128 на u_head_proj) не тянулись через полкристалла. Размеры полос пропорциональны размеру весов: pb_ff_up жирнее pb_proj по BRAM, потому что ff_up хранит 512x128 против 128x128 у proj, и так далее.

Head projection ушел на правый край (X146..X163) во всю высоту кристалла, потому что он живет в transformer_top, и его удобно держать на отшибе - он шарит tok_emb с embedding lookup, и весь этот weight-tied кусок логически отдельный. В верхней половине (Y150+) поселились остальные. Слева ln_f (X0..X100), правее в полосе X100..X145 живут u_sm_a и u_sm_b в одном pblock'е (чтобы оба softmax'а были рядом со своими ping-pong буферами), а сэмплер сидит в той же X-полосе сверху и снизу от softmax'а - в Y150..Y169 и Y211..Y249, обтекая его.

Благодаря этому удалось достичь Fmax ~98.8 MHz и зафиксировать рабочую частоту на 95 MHz