| |
| """Native K and U frontends in one isolated process/GPU, changed weight storage only.""" |
| import os,sys,json,time,threading,traceback,contextlib |
| from pathlib import Path |
| assert Path('/.dockerenv').exists() and os.environ.get('DUO_ISOLATED')=='1' |
| import torch,numpy as np |
| from autotune_guard import install as install_autotune_guard,pin_answer_configs |
| install_autotune_guard() |
| import fast_fp8_gemm |
| fast_fp8_gemm.install() |
| import transformers |
| from concurrent.futures import ThreadPoolExecutor |
| K_STREAM=None |
| G_PATH=os.environ['REAX_PART_G'] |
| Q_PATH=os.environ['REAX_PART_Q'] |
| from transformers import AutoConfig,AutoTokenizer,AutoModelForCausalLM,AutoModelForImageTextToText,FineGrainedFP8Config,BitsAndBytesConfig |
| sys.path.insert(0,str(Path(__file__).resolve().parent/'third_party/knowline')) |
| import knowline_server as ks |
| from knowline_layout import Layout |
| mode,scheme="duo","fp8";torch.set_num_threads(4) |
| def quant(): |
| if scheme=='bf16':return None |
| if scheme.startswith('fp8'): |
| return FineGrainedFP8Config(weight_block_size=(128,128) if scheme=='fp8' else (64,64),modules_to_not_convert=['lm_head','visual','model.visual','vision_tower','model.vision_tower','multi_modal_projector','model.multi_modal_projector','audio_tower','model.audio_tower','embed_vision','model.embed_vision','embed_audio','model.embed_audio','in_proj_a','in_proj_b']) |
| return BitsAndBytesConfig(load_in_8bit=True,llm_int8_threshold=6.0,llm_int8_skip_modules=['lm_head','visual','model.visual','vision_tower','model.vision_tower','multi_modal_projector','model.multi_modal_projector','audio_tower','model.audio_tower','embed_vision','model.embed_vision','embed_audio','model.embed_audio','in_proj_a','in_proj_b']) |
| def info(m): |
| from collections import Counter |
| return {'classes':dict(Counter(type(x).__name__ for x in m.modules())),'parameter_bytes':sum(p.numel()*p.element_size() for p in m.parameters()),'dtypes':dict(Counter(str(p.dtype) for p in m.parameters()))} |
| def kload(): |
| cfg=AutoConfig.from_pretrained(G_PATH);cls=getattr(transformers,'AutoModelForMultimodalLM',AutoModelForImageTextToText) if hasattr(cfg,'vision_config') or hasattr(cfg,'audio_config') else AutoModelForCausalLM |
| kw={'dtype':torch.bfloat16,'device_map':{'':'cuda:0'},'attn_implementation':'flex_attention','local_files_only':True} |
| if quant() is not None:kw['quantization_config']=quant() |
| m=cls.from_pretrained(G_PATH,**kw).eval();tok=AutoTokenizer.from_pretrained(G_PATH) |
| class Backend: |
| def warm(self,*a):pass |
| def score(self,text,media,ids): |
| if media:raise ValueError('text only') |
| with torch.inference_mode(): |
| inputs=torch.tensor([tok.encode(text,add_special_tokens=False)],device='cuda:0') |
| logits=m(input_ids=inputs,logits_to_keep=1,use_cache=False,kernel_options={'BLOCK_M':32,'BLOCK_N':32}).logits[0,-1].float().log_softmax(-1) |
| return [float(logits[i]) for i in ids],int(inputs.shape[1]) |
| temps=json.load(open(Path(G_PATH)/'temperatures.json' if (Path(G_PATH)/'temperatures.json').exists() else Path(__file__).resolve().parent/'third_party/knowline/temperatures.json'));engine=ks.KnowLine(tok,Backend(),1.0,1,temps,Layout(read_twice=True,twice_max_tokens=3000)) |
| return engine,m |
| try: |
| k,km=kload();u=None |
| tok=AutoTokenizer.from_pretrained(G_PATH);tiny=torch.tensor([tok.encode('Hello',add_special_tokens=False)],device='cuda:0') |
| with torch.inference_mode(): |
| one=km(input_ids=tiny,logits_to_keep=1,use_cache=True,kernel_options={'BLOCK_M':32,'BLOCK_N':32}).logits.float() |
| two=km(input_ids=tiny,logits_to_keep=1,use_cache=False,kernel_options={'BLOCK_M':32,'BLOCK_N':32}).logits.float() |
| cache_parity=float((one-two).abs().max());assert torch.allclose(one,two,atol=.04,rtol=.04), 'cache-free prefill parity' |
| del one,two,tiny |
|
|
| if mode=='duo': |
| sys.path.insert(0,Q_PATH) |
| import d3_runtime |
| original=d3_runtime.backbone_class |
| def backbone(kind): |
| cls=original(kind) |
| class QuantizedLoader: |
| @staticmethod |
| def from_pretrained(*a,**kw): |
| kw['quantization_config']=quant();kw['local_files_only']=True |
| return cls.from_pretrained(*a,**kw) |
| return QuantizedLoader |
| d3_runtime.backbone_class=backbone |
| |
| from modeling_d3 import D3Model |
| u=D3Model.from_pretrained(Q_PATH,device='cuda:0',local_files_only=True,batch_size=1) |
| original_run=u.runtime._run |
| def bounded_run(prepared): |
| wanted=u.runtime.batch_size |
| longest=max([len(x) for x in prepared.sequences.values()]+[len(x) for x in prepared.reversed_sequences.values()]+[1]) |
| u.runtime.batch_size=min(wanted,max(1,4096//longest)) |
| try:return original_run(prepared) |
| finally:u.runtime.batch_size=wanted |
| u.runtime._run=bounded_run |
| pin_answer_configs() |
| ba_patches=[] |
| state_once_backend=d3_runtime.StateOnce(u.backbone.language_model) |
| opt_active='plain' |
| from transformers.models.qwen3_5 import modeling_qwen3_5 as qwen_modeling |
| from causal_conv1d import causal_conv1d_fn as official_conv |
| def optimized_conv(x,weight,bias=None,activation=None,**kw):return official_conv(x,weight,bias,activation=activation) |
| def reference_conv(x,weight,bias=None,activation=None,**kw): |
| out=torch.nn.functional.conv1d(x.to(weight.dtype),weight[:,None,:],bias,padding=weight.shape[-1]-1,groups=weight.shape[0])[...,:x.shape[-1]] |
| return (torch.nn.functional.silu(out) if activation in ['silu','swish'] else out).to(x.dtype) |
| qwen_modeling.causal_conv1d_fn=reference_conv |
| cx=torch.randn(1,512,1024,device='cuda',dtype=torch.bfloat16);cw=torch.randn(512,4,device='cuda',dtype=torch.bfloat16) |
| ca=reference_conv(cx,cw,activation='silu');cb=optimized_conv(cx,cw,activation='silu') |
| conv_probe={'max_abs':float((ca-cb).abs().max()),'mean_abs':float((ca-cb).abs().mean())} |
| assert torch.allclose(ca,cb,atol=.08,rtol=.02), 'convolution parity prerequisite' |
| del cx,cw,ca,cb |
| if scheme!='bf16': |
| for model in [km]+([u.backbone] if u else []): |
| assert any('float8' in str(p.dtype) or p.dtype==torch.int8 for p in model.parameters()), 'quantization silently dequantized' |
| torch.cuda.synchronize() |
| fast_fp8_gemm.precompile() |
| print(json.dumps({'ready':True,'autotuner_guard':True,'mode':mode,'scheme':scheme,'cache_parity_max_abs':cache_parity,'gpu':torch.cuda.get_device_name(0),'k':info(km),'u':info(u.backbone) if u else None,'allocated':torch.cuda.memory_allocated(),'reserved':torch.cuda.memory_reserved(),'peak_load_bytes':torch.cuda.max_memory_allocated()}),flush=True) |
| except BaseException: |
| traceback.print_exc();raise |
|
|
| K_DONE=threading.Event() |
| def invoke_k(r): |
| try: |
| a,usage=k.run(r['state'],r['questions']) |
| if opt_active in ['longcache','fastconv_longcache','fast_all_longcache']:torch.cuda.current_stream().synchronize() |
| return {'answers':a,'usage':usage} |
| finally:K_DONE.set() |
| def invoke_u(r): |
| u.runtime.fast_route={'state_once':False,'oom_fallback':False,'serialized_for_memory':False} |
| if opt_active in ['state_once','fastconv_state_once','fast_all','longcache','fastconv_longcache','fast_all_longcache'] and len(r['questions'])>1: |
| prepared=u.runtime.prepare(r['state'],r['questions']) |
| prefix=max(list(prepared.states.values())+[0]);long_mode=opt_active in ['longcache','fastconv_longcache','fast_all_longcache'];use=0<prefix<=(24576 if long_mode else 16384) |
| u.runtime.fast_route.update(prefix_tokens=prefix,state_once=use) |
| if use and long_mode and prefix>16384: |
| |
| if not K_DONE.wait(300):raise RuntimeError('K completion wait exceeded memory-scheduling budget') |
| u.runtime.fast_route['serialized_for_memory']=True |
| u.runtime.state_once=state_once_backend if use else None |
| try: |
| result=None |
| if use: |
| try:result=u.runtime.run(prepared) |
| except torch.OutOfMemoryError:pass |
| if result is None: |
| if use: |
| u.runtime.state_once=None;u.runtime.fast_route.update(state_once=False,oom_fallback=True);torch.cuda.empty_cache() |
| result=u.runtime.run(prepared) |
| probabilities,tokens=result |
| return u.runtime.respond(prepared,probabilities,tokens) |
| finally:u.runtime.state_once=None |
| return u.system_one(state=r['state'],questions=r['questions']) |
|
|
| from probabilities import mix |
| norm_patches=[] |
| for mod_name,mod in km.named_modules(): |
| if type(mod).__name__.endswith('RMSNorm'): |
| norm_patches.append((mod,mod.forward,None)) |
| def apply_compile_norm(on): |
| global norm_patches |
| updated=[] |
| for mod,plain,compiled in norm_patches: |
| if on and compiled is None:compiled=torch.compile(plain,dynamic=True) |
| mod.forward=compiled if on else plain |
| updated.append((mod,plain,compiled)) |
| norm_patches=updated |
| stream_pool=ThreadPoolExecutor(2) |
| stream_k=torch.cuda.Stream();stream_u=torch.cuda.Stream() |
| stream_k.wait_stream(torch.cuda.current_stream());stream_u.wait_stream(torch.cuda.current_stream()) |
| def probe_graph(): |
| apply_compile_norm(False) |
| static=torch.tensor([AutoTokenizer.from_pretrained(G_PATH).encode('Hello, how are you?',add_special_tokens=False)],device='cuda:0') |
| def forward():return km(input_ids=static,logits_to_keep=1,use_cache=False,kernel_options={'BLOCK_M':32,'BLOCK_N':32}).logits.float() |
| with torch.inference_mode(): |
| for _ in range(3):ref=forward() |
| torch.cuda.synchronize();ref=ref.clone() |
| t=time.perf_counter() |
| for _ in range(20):out=forward() |
| torch.cuda.synchronize();eager=1000*(time.perf_counter()-t)/20 |
| graph=torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph):captured=forward() |
| graph.replay();torch.cuda.synchronize();diff=float((ref-captured).abs().max()) |
| t=time.perf_counter() |
| for _ in range(20):graph.replay() |
| torch.cuda.synchronize();graphed=1000*(time.perf_counter()-t)/20 |
| return {'status':'ok','scope':'diagnostic repeated tiny synthetic K forward only, not representative latency','tokens':static.shape[1],'eager_ms':eager,'graph_ms':graphed,'max_abs_logit_diff':diff} |
| for line in sys.stdin: |
| try: |
| r=json.loads(line);kind=r.pop('_kind','concurrent') |
| if kind=='graph_probe': |
| try:probe=probe_graph() |
| except BaseException as e:probe={'status':'error','scope':'K tiny diagnostic graph capture','error':repr(e)} |
| print(json.dumps(probe),flush=True);continue |
| ubatch=int(r.pop('_ubatch',8));opt_active=r.pop('_opt','fastconv_longcache');u.runtime.batch_size=ubatch |
| for mod,plain,fast in ba_patches:mod.forward=fast if opt_active in ['miner_ba','fast_all','fast_all_longcache'] else plain |
| apply_compile_norm(opt_active in ['compile_norm','fast_all','fast_all_longcache']) |
| qwen_modeling.causal_conv1d_fn=optimized_conv if opt_active in ['fastconv','fastconv_state_once','fast_all','fastconv_longcache','fast_all_longcache'] else reference_conv |
| K_DONE.clear() |
| if kind=='single_u':K_DONE.set() |
| torch.cuda.synchronize();torch.cuda.reset_peak_memory_stats();start=time.perf_counter() |
| a=b=None |
| if kind=='concurrent': |
| def call_k(): |
| with torch.cuda.stream(stream_k): |
| torch._dynamo.config.recompile_limit=128 |
| return invoke_k(r) |
| def call_u(): |
| with torch.cuda.stream(stream_u):return invoke_u(r) |
| fa=stream_pool.submit(call_k);fb=stream_pool.submit(call_u) |
| |
| ae=be=None |
| try:a=fa.result() |
| except BaseException as e:ae=e |
| try:b=fb.result() |
| except BaseException as e:be=e |
| torch.cuda.synchronize() |
| if ae or be:raise (ae or be) |
| mid=start |
| elif kind=='single_k': |
| a=invoke_k(r);torch.cuda.synchronize();mid=time.perf_counter() |
| elif kind=='single_u': |
| mid=start;b=invoke_u(r) |
| else: |
| a=invoke_k(r);torch.cuda.synchronize();mid=time.perf_counter();b=invoke_u(r) |
| torch.cuda.synchronize();end=time.perf_counter() |
| combined=mix(r,a,b) if a and b else None;duo_end=time.perf_counter() |
| result={'K':a,'U':b,'Duo':combined,'timings_ms':{'K':1000*(mid-start),'U':1000*(end-mid),'Duo':1000*(duo_end-start)},'peak_bytes':torch.cuda.max_memory_allocated(),'peak_reserved_bytes':torch.cuda.max_memory_reserved(),'optimization':{'name':opt_active,'miner_ba_patched_modules':len(ba_patches),'compile_norm_modules':len(norm_patches),'conv_probe':conv_probe,'u_route':getattr(u.runtime,'fast_route',None)}} |
| print(json.dumps({'status':'ok',**result}),flush=True) |
| except BaseException as e: |
| traceback.print_exc();torch.cuda.synchronize();torch.cuda.empty_cache() |
| print(json.dumps({'status':'error','error':str(e),'exception':type(e).__name__}),flush=True) |
|
|
