훈련이 끝난 뒤 part 11 of 13

Saved 1.6x the memory, lost 1.7x the time

guide / / 8 sections

Part ten stopped at the point where continuous batching makes the cache longer. Hold the cache as one dense tensor and its length is set by the longest row, and continuous batching deliberately parks a freshly admitted short row beside a long-running one. The closing paragraph noted that cutting the cache into fixed-size blocks would make that price disappear, and left it there.

So I cut it. The memory fell exactly as advertised. The time went the other way.

What is left after cutting

With block size S, a row holding L tokens gets ceil(L/S) blocks. The blocks live in a shared pool and each row has a block table saying which ones are its own. Nothing requires a row’s blocks to sit next to each other in the pool.

Taking the step with the largest allocation from part ten’s continuous trace:

Dense - every row padded to the longest Blocks of 16 - only the blocks a row needs 4,064 slots 2,560 slots Total slots by block size 0k 500k 1000k dense 1,009,601 actual tokens 567,139 1 2 4 8 16 32 64 128 block size
The step with the largest allocation in part ten's continuous trace. Its 32 rows hold caches of different lengths, and the dense tensor takes all of them to the longest, 127, for 4,064 slots. Cut into blocks of 16, only each row's last block is short-filled, for 2,560. Below, total slots over the whole run against block size - it rises with the block, and at 128 it is worse than dense.

Those 32 rows actually hold 2,337 tokens. The dense tensor takes all of them to the longest row at 127, so 32 x 127 = 4,064 slots. Blocks of 16 come to 2,560, the extra 223 being each row’s short-filled last block.

Over the whole run:

block      slots     occ     peak live    blocks to track
    1    567,139  100.0%         2,353              2,353
    2    571,956   99.2%         2,368              1,184
    4    581,580   97.5%         2,400                600
    8    600,824   94.4%         2,448                306
   16    639,184   88.7%         2,560                160
   32    714,240   79.4%         2,848                 89
   64    858,944   66.0%         3,456                 54
  128  1,232,640   46.0%         4,096                 32
dense  1,009,601   56.2%         4,064                 32

At block 16 the peak live allocation goes from 4,064 to 2,560, a factor of 1.59. A slot is 3 blocks x (k,v) x 4 heads x 32 x float32 = 3,072 bytes, so 12,192 KB becomes 7,680 KB. These figures are exact and independent of any implementation.

A block as large as the context is worse than dense

Set the last two rows of the table side by side and block 128 uses more slots than dense. Dense fills only to the longest row currently alive, while block 128 claims 128 slots whatever the length. A block the size of the maximum context turns paging off and keeps the waste.

The other end is not free either. Block 1 gives exactly 100% occupancy but leaves 2,353 blocks to manage, against 160 at block 16 - a table fifteen times longer.

The time goes the other way

Timing one step at 32 rows and cache 104:

              slots/row    time (us)   vs dense
dense               104         1010       1.00
block 4             104         1923       1.90
block 8             104         1825       1.81
block 16            112         1850       1.83
block 32            128         1893       1.87
block 64            128         1725       1.71

There is no trend in block size. The quartiles overlap each other, and block 4, which uses the fewest slots, is the slowest of them. Block size is a memory dial, not a time dial.

And every block size is 1.7 to 1.9 times slower than dense.

The price is in the gathering

Before attention runs, each row’s blocks have to be collected by following the block table. Timing that alone:

building the same shape (32 x 4 x 112 x 32)
  blocks scattered at random       65.8 us
  contiguous within each row       64.1 us
  whole pool in order              66.1 us
  dense copy of the same size      14.5 us

4.5 times, however the blocks are laid out in the pool. I had assumed my random block table was measuring a worst case; laying the blocks out in order changes nothing. The price is not scattered addresses, it is materialising a new tensor. That is why real systems fold the block lookup into the attention kernel, where the copy does not exist at all. This one does not.

Even after converting the memory into batch

Saved memory gets paid back as batch size. The 4,064 slots dense uses at batch 32 will hold batch 54 with blocks of 16. Does that make it back?

Dividing two configurations against each other in the same round, alternating which runs first, 21 rounds:

                                 median   quartiles         won
dense 32 / paged 32               1.663   1.509 ~ 1.766   21/21
dense 32 / paged 54               1.551   1.482 ~ 1.655   21/21
dense 54 / dense 32               1.224   1.142 ~ 1.361   21/21

At the same batch it loses by 1.663. Raising the batch to 54 recovers it only to 1.551, and the gap is narrow because the gathering cost grows with rows too.

The third row is the point of this part. Going from batch 32 to 54 is worth 1.224 by itself. Part ten put 49% of a step in fixed cost and 21% on rows, so 1.7 times the rows buys about that much and no more. It is a trade of 1.66 paid for 1.22 bought.

Squeezing the budget

What about when memory really is tight? Fixing a budget and giving each method the largest batch that fits inside it:

 budget  dense  block    steps (dense/block)   dense/paged      quartiles
    400      3      3          3216 / 3216           1.308   1.280 ~ 1.332
    600      4      5          2418 / 1948           1.108   1.071 ~ 1.119
    900      7     10          1402 /  991           0.990   0.966 ~ 1.012
  1,200      9     13          1097 /  775           1.113   1.067 ~ 1.138
  1,500     11     17           908 /  602           1.102   1.070 ~ 1.130
  2,560     20     32           518 /  336           1.212   1.183 ~ 1.287
  4,064     32     54           336 /  214           1.457   1.427 ~ 1.543
  6,000     47     75           237 /  164           1.703   1.690 ~ 1.775

Above 1 means dense is faster. The looser the budget the further dense pulls ahead; the tighter it gets the closer they come. In one row of eight, at 900, paging is ahead - 0.990, quartiles 0.966~1.012, which is closer to a tie than a win.

Why it loses again at 400 is plain: that budget holds batch 3 either way, so there is no batch for paging to buy. It pays the gathering and receives nothing.

What is left

What this part measured is paging that builds a dense tensor. Fold the block lookup into attention and the 4.5x copy disappears and every table here changes

  • I did not fold it, so I do not know by how much. These numbers cannot judge paging itself; they carry the condition “if the allocator changes and the kernel does not”.

The scale is different too. This model’s cache peaks at 12 MB, so memory has never once stopped it from running. Where caches run to gigabytes the question is not whether a little more batch fits but whether anything runs at all, and then the comparison is not against 1.22 but against zero.

The memory table does not depend on any of that. Slot counts reproduce exactly once the list of lengths is fixed, and 4,064 against 2,560 is the same on any machine.

So

  • Cutting into blocks takes peak live slots from 4,064 to 2,560, a factor of 1.59. Occupancy 56.2% to 88.7%
  • A block as large as the maximum context is worse than dense: 1,232,640 against 1,009,601
  • Block 1 is 100% occupied and leaves 2,353 blocks to manage
  • Time is 1.7 to 1.9 times worse than dense at every block size
  • The price is the copy, not the scatter: blocks laid out in order still cost 4.5 times a dense copy
  • Converting the saved memory into batch 32 to 54 is worth 1.224. A trade of 1.66 paid for 1.22 bought
  • Paging leads in one budget of eight, at 900, and only by 0.990

Comments