| |
| import os |
| from pathlib import Path |
| import sys |
|
|
| sys.dont_write_bytecode = True |
| ROOT = Path(__file__).resolve().parent |
| sys.path.insert(0, str(ROOT)) |
| sys.path.insert(0, str(ROOT / 'runtime')) |
| os.environ.pop('PYTHONPATH', None) |
|
|
| import argparse |
| import json |
| import runpy |
| import time |
|
|
| from inference import ROOT, Engine, dump |
|
|
|
|
| def benchmark(args): |
| import torch |
| engine = Engine(args.model, args.batch, args.profile, backend='cuda') |
| groups = args.batch // args.k |
| positions = torch.arange(args.prefill).unsqueeze(0) |
| group_ids = torch.arange(groups).unsqueeze(1) |
| unique = (positions * 104729 + 8191 + group_ids * 15485863).remainder(248320 - 4096) + 2048 |
| ids = unique.repeat_interleave(args.k, dim=0).cuda() |
| mask = torch.ones_like(ids) |
| common = dict(slot_batch_size=args.batch, max_new_tokens=args.max_new_tokens, |
| temperature=1.0, top_k=20, top_p=0.95, presence_penalty=1.5, |
| eos_token_ids=(), pad_token_id=engine.tokenizer.pad_token_id, stop_check_interval_tokens=128, |
| attempt_ids=torch.arange(args.batch, device='cuda'), |
| prefill_group_ids=torch.arange(args.batch, device='cuda').div(args.k, rounding_mode='floor'), |
| refill_batch_size=min(args.batch,16),refill_prefetch_size=min(args.batch,16), |
| prefill_chunk_size=min(4 if args.prefill>=8192 else 16,max(1,32768//args.prefill)), |
| prefill_token_chunk_size=16384, sampling_vocab_size=len(engine.tokenizer), |
| counter_sampling_backend='triton_top20', adaptive_tail_compaction=False, |
| mechanical_repetition_early_stop=False) |
| with torch.inference_mode(): |
| started=time.perf_counter() |
| warm=engine.benchmark_generate(engine.model,ids,mask,**common,counter_seed=20260827, |
| max_new_tokens_per_attempt=torch.full((args.batch,),args.warmup,device='cuda',dtype=torch.long)) |
| cold=dict(seconds=time.perf_counter()-started,ttft_seconds=warm.metrics['initial_prefill_seconds']) |
| del warm |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats() |
| result=engine.benchmark_generate(engine.model,ids,mask,**common, |
| counter_seed=20260828+args.prefill+(args.batch if args.batch>1 else 0)) |
| metrics=result.metrics |
| prefill=metrics['initial_prefill_seconds']+metrics['refill_prefill_seconds'] |
| decode=metrics['generation_wall_seconds']-prefill |
| receipt=dict(status='COMPLETE',batch=args.batch,prefill_tokens=args.prefill, |
| generated_tokens=int(result.generated_lengths.sum()),decode_tokens=args.max_new_tokens, |
| ttft_seconds=metrics['initial_prefill_seconds'],decode_seconds=decode, |
| decode_tokens_per_second=metrics['decode_useful_token_events']/decode, |
| prefill_unique_tokens_per_second=groups*args.prefill/prefill, |
| end_to_end_output_tokens_per_second=int(result.generated_lengths.sum())/metrics['generation_wall_seconds'], |
| peak_allocated_gib=torch.cuda.max_memory_allocated()/2**30, |
| peak_reserved_gib=torch.cuda.max_memory_reserved()/2**30, |
| warmup=cold,model_load_seconds=engine.load_seconds,operators=engine.operators, |
| torch=torch.__version__,gpu=torch.cuda.get_device_name(),metrics=metrics) |
| dump(args.output,receipt) |
| print(json.dumps(receipt,ensure_ascii=False),flush=True) |
|
|
|
|
| def serve(args): |
| from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer |
| from queue import Queue |
| import threading |
| import uuid |
| from transformers import TextStreamer |
| engine = Engine(args.model, args.batch, args.profile, backend=args.backend, device=args.device, dtype=args.dtype) |
| warmup_started = time.perf_counter() |
| if args.batch == 1: |
| engine.generate([dict(prompt='Compute 17 × 23.', max_tokens=16)], cap=16, mechanical=False) |
| lock = threading.Lock() |
|
|
| class Handler(BaseHTTPRequestHandler): |
| def send_json(self, value, status=200): |
| body = json.dumps(value, ensure_ascii=False).encode() |
| self.send_response(status) |
| self.send_header('Content-Type', 'application/json') |
| self.send_header('Content-Length', str(len(body))) |
| self.end_headers() |
| self.wfile.write(body) |
|
|
| def do_GET(self): |
| if self.path == '/v1/models': |
| self.send_json(dict(object='list', data=[dict(id='EdgeRock-4B', object='model', owned_by='rockmatrix-ai')])) |
| else: |
| self.send_json(dict(status='READY', model='EdgeRock-4B', backend=engine.backend, batch_capacity=args.batch)) |
|
|
| def do_POST(self): |
| try: |
| request = json.loads(self.rfile.read(int(self.headers['Content-Length']))) |
| standard = self.path == '/v1/chat/completions' |
| if self.path not in {'/', '/v1/chat/completions'}: |
| self.send_json({'error': {'message': 'Unknown endpoint'}}, 404) |
| return |
| unsupported = [key for key in ['tools', 'tool_choice', 'response_format', 'stop'] if request.get(key)] |
| if unsupported or request.get('n', 1) != 1: |
| self.send_json({'error': {'message': 'This endpoint supports plain text chat with n=1.'}}, 400) |
| return |
| prompts = [{'messages': request['messages']}] if standard else request.get('prompts') or [{'messages': request['messages']}] |
| settings = {key: request[key] for key in ['temperature', 'top_k', 'top_p', 'presence_penalty', 'seed'] if key in request} |
| cap = int(request.get('max_completion_tokens', request.get('max_tokens', request.get('max_new_tokens', args.max_new_tokens)))) |
| identity, created = 'chatcmpl-' + uuid.uuid4().hex, int(time.time()) |
|
|
| def completion(rows): |
| row = rows[0] |
| text = engine.tokenizer.decode([t for t in row['response_ids'] if t not in (248044, 248046)], skip_special_tokens=False) |
| usage = dict(prompt_tokens=row['prompt_tokens'], completion_tokens=row['generated_tokens'], total_tokens=row['prompt_tokens']+row['generated_tokens']) |
| reason = 'length' if row['cap_hit'] else 'stop' |
| return text, usage, reason |
|
|
| if standard and request.get('stream'): |
| queue = Queue() |
| class TokenStream(TextStreamer): |
| def on_finalized_text(self, text, stream_end=False): |
| if text: |
| queue.put(('text', text)) |
| decoder = TokenStream(engine.tokenizer, skip_prompt=False, skip_special_tokens=False) |
| def on_tokens(attempt, tokens): |
| tokens = [t for t in tokens if t not in (248044, 248046)] |
| if tokens: |
| decoder.put(engine.torch.tensor(tokens)[None]) |
| def generate(): |
| try: |
| with lock: |
| rows, metrics = engine.generate(prompts, cap=cap, on_tokens=on_tokens, **settings) |
| decoder.end() |
| queue.put(('done', completion(rows))) |
| except Exception as error: |
| queue.put(('error', str(error))) |
| threading.Thread(target=generate, daemon=True).start() |
| self.send_response(200) |
| self.send_header('Content-Type', 'text/event-stream') |
| self.send_header('Cache-Control', 'no-cache') |
| self.end_headers() |
| def event(value): |
| self.wfile.write(('data: ' + json.dumps(value, ensure_ascii=False) + '\n\n').encode()) |
| self.wfile.flush() |
| event(dict(id=identity, object='chat.completion.chunk', created=created, model='EdgeRock-4B', choices=[dict(index=0, delta={'role':'assistant'}, finish_reason=None)])) |
| while True: |
| kind, value = queue.get() |
| if kind == 'error': |
| event({'error': {'message': value}}) |
| break |
| if kind == 'text': |
| event(dict(id=identity, object='chat.completion.chunk', created=created, model='EdgeRock-4B', choices=[dict(index=0, delta={'content':value}, finish_reason=None)])) |
| else: |
| text, usage, reason = value |
| event(dict(id=identity, object='chat.completion.chunk', created=created, model='EdgeRock-4B', choices=[dict(index=0, delta={}, finish_reason=reason)])) |
| if request.get('stream_options', {}).get('include_usage'): |
| event(dict(id=identity, object='chat.completion.chunk', created=created, model='EdgeRock-4B', choices=[], usage=usage)) |
| break |
| self.wfile.write(b'data: [DONE]\n\n') |
| self.wfile.flush() |
| return |
| with lock: |
| rows, metrics = engine.generate(prompts, cap=cap, **settings) |
| if standard: |
| text, usage, reason = completion(rows) |
| self.send_json(dict(id=identity, object='chat.completion', created=created, model='EdgeRock-4B', choices=[dict(index=0, message={'role':'assistant', 'content':text}, finish_reason=reason)], usage=usage)) |
| else: |
| self.send_json(dict(outputs=rows, metrics=metrics)) |
| except (KeyError, TypeError, ValueError) as error: |
| self.send_json({'error': {'message': str(error)}}, 400) |
| except (BrokenPipeError, ConnectionResetError): |
| return |
|
|
| print(json.dumps(dict(status='READY', host=args.host, port=args.port, backend=engine.backend, |
| batch=args.batch, warmup_seconds=time.perf_counter()-warmup_started)), flush=True) |
| ThreadingHTTPServer((args.host, args.port), Handler).serve_forever() |
|
|
|
|
| def main(): |
| os.environ.setdefault('TOKENIZERS_PARALLELISM', 'false') |
| os.environ.setdefault('NLTK_DATA', str(ROOT/'runtime/third_party/verifiers/nltk_data')) |
| if len(sys.argv)>1 and sys.argv[1]=='_verify': |
| kind=sys.argv[2] |
| sys.argv=[sys.argv[0],*sys.argv[3:]] |
| runpy.run_module('evaluation.'+{'math':'math_verifier','code':'code_verifier','instruction':'instruction_verifier'}[kind],run_name='__main__') |
| return |
| parser=argparse.ArgumentParser(description='EdgeRock-4B inference and reproducible evaluation') |
| sub=parser.add_subparsers(dest='command',required=True) |
| for name in ['demo','generate','serve','benchmark','evaluate','score','status','_worker']: |
| p=sub.add_parser(name) |
| p.add_argument('--model',type=Path,default=ROOT) |
| p.add_argument('--output',type=Path,default=Path('outputs')/name) |
| if name in {'demo','generate','serve','benchmark','evaluate','_worker'}: |
| p.add_argument('--batch',type=int,default=1 if name in {'demo','serve'} else 320) |
| p.add_argument('--backend',choices=['auto','torch','cuda'],default='auto') |
| p.add_argument('--device') |
| p.add_argument('--dtype',choices=['float32','float16','bfloat16']) |
| p.add_argument('--profile',choices=['fast','reference'],default='fast') |
| p.add_argument('--max-new-tokens',type=int,default=4096 if name in {'demo','benchmark'} else 131072) |
| if name in {'demo','generate','evaluate','_worker'}: |
| p.add_argument('--seed',type=int,default=20260923) |
| if name=='demo': |
| p.add_argument('--prompt',default='计算 17 × 23，并简短说明计算过程。') |
| if name=='generate': |
| p.add_argument('--input',type=Path,required=True) |
| if name=='serve': |
| p.add_argument('--host',default='127.0.0.1');p.add_argument('--port',type=int,default=8000) |
| if name=='benchmark': |
| p.add_argument('--prefill',type=int,default=4096) |
| p.add_argument('--warmup',type=int,default=256) |
| p.add_argument('--k',type=int,default=None) |
| if name in {'evaluate','score'}: |
| p.add_argument('--suite',choices=['text','memory','vision'],default='text') |
| p.add_argument('--sources',nargs='*') |
| p.add_argument('--cpu-workers',type=int,default=160) |
| if name=='evaluate': |
| p.add_argument('--gpus',default='0') |
| p.add_argument('--attempts',type=Path,help='Rerun selected source/question/sample identities from a Parquet file.') |
| p.add_argument('--limit-per-source',type=int) |
| p.add_argument('--k',type=int,help='Override automatic K; such runs are reported separately.') |
| p.add_argument('--generate-only',action='store_true') |
| if name=='_worker': |
| p.add_argument('--rank',type=int,required=True) |
| p.add_argument('--port',type=int,required=True) |
| p.add_argument('--suite',choices=['text','memory','vision'],default='text') |
| args=parser.parse_args() |
| if args.command=='benchmark': |
| args.k = args.k if args.k is not None else min(8, args.batch) |
| if min(args.batch, args.prefill, args.max_new_tokens, args.warmup, args.k) < 1: |
| parser.error('benchmark batch, lengths, warmup, and k must be positive') |
| if args.batch % args.k: |
| parser.error('benchmark k must divide batch size') |
| benchmark(args) |
| elif args.command=='serve': |
| serve(args) |
| elif args.command in {'demo','generate'}: |
| prompts=[args.prompt] if args.command=='demo' else [json.loads(line) for line in args.input.read_text().splitlines() if line.strip()] |
| engine=Engine(args.model,args.batch,args.profile,backend=args.backend,device=args.device,dtype=args.dtype) |
| rows,metrics=engine.generate(prompts,cap=args.max_new_tokens,seed=args.seed) |
| args.output.mkdir(parents=True,exist_ok=True) |
| (args.output/'responses.jsonl').write_text(''.join(json.dumps(r,ensure_ascii=False)+'\n' for r in rows)) |
| dump(args.output/'metrics.json',metrics) |
| for row in rows: |
| print(row['response'],flush=True) |
| print(json.dumps({'ttft_seconds':metrics['initial_prefill_seconds'],'decode_tokens_per_second':metrics['decode_tokens_per_second'],'output':str(args.output)},ensure_ascii=False)) |
| elif args.command=='status': |
| print((args.output/'STATUS.json').read_text()) |
| else: |
| from evaluate import evaluate,score,worker |
| {'evaluate':evaluate,'score':score,'_worker':worker}[args.command](args) |
|
|
| if __name__ == '__main__': |
| main() |
|
|
