LAB 03 · KV、分页、前缀与请求调度

import math, json
layers=32; kv_heads=8; head_dim=128; bytes_per_number=2
per_token=2*layers*kv_heads*head_dim*bytes_per_number
print('每token:',per_token,'bytes =',per_token/1024,'KiB')
for length,requests in [(8192,1),(8192,8),(32768,1)]:
    print(length,'tokens ×',requests,'请求 =',per_token*length*requests/2**30,'GiB')
assert per_token==131072

block_size=4
physical=[7,1,5]
def page_table(length,size=4,physical=(7,1,5,9,3)):
    count=math.ceil(length/size)
    if count>len(physical): raise ValueError('物理空闲块不足')
    return [{'logical':i,'physical':physical[i],
             'token_start':i*size,'used':min(size,length-i*size)} for i in range(count)]
for length in (6,9):
    table=page_table(length,block_size,physical)
    print('长度',length,'表',table,'分配token位',len(table)*block_size,'尾块浪费',len(table)*block_size-length)
for token in range(9):
    print('token',token,'→ physical block',physical[token//4],'offset',token%4)

lengths=[6,9,17]
for size in (1,4,16):
    allocated=sum(math.ceil(n/size)*size for n in lengths)
    print('块长',size,'实际token',sum(lengths),'分配token位',allocated,'内部碎片',allocated-sum(lengths),'块表条目',sum(math.ceil(n/size) for n in lengths))

prefix=4096; suffix=2048; requests=3
no_share=requests*(prefix+suffix)*per_token/2**30
shared=(prefix+requests*suffix)*per_token/2**30
print('不共享',no_share,'GiB；精确共享',shared,'GiB；节省',no_share-shared,'GiB')
assert no_share==2.25 and shared==1.25
cache={('model-v1',0,(101,202,303)):'physical-prefix-7'}
def cache_hit(model,position,tokens): return (model,position,tuple(tokens)) in cache
for label,args in [('完全相同',('model-v1',0,[101,202,303])),('改一个token',('model-v1',0,[101,202,304])),('改模型',('model-v2',0,[101,202,303])),('改起始位置',('model-v1',1,[101,202,303]))]:
    print(label,'命中=',cache_hit(*args))

REQUESTS=[{'id':'A','arrival':0,'output_tokens':2},{'id':'B','arrival':0,'output_tokens':4},{'id':'C','arrival':1,'output_tokens':2}]
def simulate(requests,capacity=2,continuous=True):
    remaining={r['id']:r['output_tokens'] for r in requests}
    arrival={r['id']:r['arrival'] for r in requests}
    active=[]; done={}; started={}; trace=[]; step=0
    while len(done)<len(requests):
        waiting=[r['id'] for r in requests if arrival[r['id']]<=step and r['id'] not in done and r['id'] not in active]
        if continuous or not active:
            for name in waiting[:capacity-len(active)]: active.append(name); started[name]=step
        row={'step':step,'slots':active.copy()+['空']*(capacity-len(active))}
        for name in list(active):
            remaining[name]-=1
            if remaining[name]==0: done[name]=step+1; active.remove(name)
        trace.append(row); step+=1
        if step>1000: raise RuntimeError('模拟未终止')
    return {'trace':trace,'start':started,'finish':done,'makespan_steps':step,'occupied_slot_steps':sum(r['output_tokens'] for r in requests),'total_slot_steps':step*capacity}
for mode in (False,True):
    print('连续接纳' if mode else '静态批次',json.dumps(simulate(REQUESTS,continuous=mode),ensure_ascii=False,indent=2))
assert simulate(REQUESTS)['start']['C']==2
assert simulate(REQUESTS,continuous=False)['start']['C']==4

boundary=[{'id':'A','arrival':0,'output_tokens':2},{'id':'B','arrival':0,'output_tokens':2},{'id':'C','arrival':1,'output_tokens':2}]
for mode in (False,True): print('continuous=',mode,'finish=',simulate(boundary,continuous=mode)['finish'])
assert simulate(boundary,continuous=False)['finish']==simulate(boundary,continuous=True)['finish']