File size: 4,128 Bytes
2e0601e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 | """Reuse request-local gather indices for fixed-shape token-local operations."""
from collections import defaultdict
from .branch_attention import REQUEST
from contextvars import ContextVar
import types
import torch
DEPTH=ContextVar('spev_vectorized_token_ops_depth',default=0)
class PackingCache:
def __init__(self):
self.tree=None;self.device=None;self.plan=None
def __call__(self,tree,device):
if tree is self.tree and device==self.device:return self.plan
buckets=defaultdict(list);start=tree.prefix_length
for index,length in enumerate(tree.branch_lengths):
buckets[((length+31)//32)*32].append((index,start,length));start+=length
assert start==tree.length
blocks=[];restore=[None]*len(tree.branch_lengths);offset=0
for width,branches in sorted(buckets.items()):
for begin in range(0,len(branches),8):
block=branches[begin:begin+8];indices=[0]*(8*width);valid=[False]*(8*width)
for slot,(index,start,length) in enumerate(block):
a=slot*width;indices[a:a+length]=range(start,start+length);valid[a:a+length]=[True]*length
restore[index]=torch.arange(length,device=device)+offset+a
blocks.append((width,torch.tensor(indices,device=device,dtype=torch.long),torch.tensor(valid,device=device,dtype=torch.bool)))
offset+=8*width
order=torch.cat(restore) if restore else torch.empty(0,device=device,dtype=torch.long)
self.tree=tree;self.device=device;self.plan=(blocks,order)
return self.plan
def grouped_forward(original,tree_getter,plan_getter=None):
cache=PackingCache() if plan_getter is None else plan_getter
def forward(module,value,*args,**kwargs):
if DEPTH.get():return original(value,*args,**kwargs)
tree=tree_getter()
if tree is None:raise ValueError('Grouped token operations require a tree layout')
if value.ndim not in (2,3) or value.shape[-2]!=tree.length:raise ValueError('Grouped token operations require the packed token sequence')
if value.ndim==3 and value.shape[0]!=1:raise ValueError('Grouped token operations support one packed request')
x=value[0] if value.ndim==3 else value
blocks,order=cache(tree,x.device);token=DEPTH.set(DEPTH.get()+1)
try:
prefix=original(x[:tree.prefix_length].contiguous(),*args,**kwargs);outputs=[]
for width,indices,valid in blocks:
packed=x.index_select(0,indices).masked_fill(~valid[:,None],0).reshape(8,width,x.shape[-1])
output=original(packed,*args,**kwargs)
assert isinstance(output,torch.Tensor) and output.shape[:2]==packed.shape[:2]
outputs.append(output.reshape(8*width,output.shape[-1]))
branch=(outputs[0] if len(outputs)==1 else torch.cat(outputs,dim=0)).index_select(0,order) if outputs else prefix[:0]
output=torch.cat([prefix,branch],dim=0)
return output[None] if value.ndim==3 else output
finally:DEPTH.reset(token)
return forward
def install(model,*,linears,moe):
if moe:raise ValueError('MoE dispatch is not a token-local linear operation')
originals=[]
def cache(tree, device):
request=REQUEST.get()
attention_cache=request[1]
if not hasattr(attention_cache, "packing"):
attention_cache.packing=PackingCache()
return attention_cache.packing(tree,device)
getter=lambda:REQUEST.get()[0] if REQUEST.get() else None
if linears:
for layer in model.backbone.model.layers:
for module in list(layer.modules()):
is_lora=all(hasattr(module,name) for name in ['a','b','base','scale'])
if isinstance(module,torch.nn.Linear) or is_lora:
original=module.forward;originals.append((module,original))
module.forward=types.MethodType(grouped_forward(original,getter,cache),module)
def restore():
for module,forward in reversed(originals):module.forward=forward
return restore
|