| """Causal reward-guided chunk search with persistent global/codec KV caches. |
| |
| Winner-of-N: keep one prefix and fork it N ways at each chunk. |
| Top-2: retain up to two global winners, split N children (extra to the best). |
| Only common acoustic |
| token prefixes are irreversible streaming output; unrelated beams are not spliced. |
| """ |
| import math |
| import time |
|
|
| import numpy as np |
| import torch |
|
|
| from performance_core import combine_branches, StreamingCodec |
|
|
|
|
| def expansion_indices(winners, n): |
| if not winners or n < len(winners): |
| raise ValueError('Candidate budget must cover retained beam count') |
| quotient, remainder = divmod(n, len(winners)) |
| return [index for position,index in enumerate(winners) |
| for _ in range(quotient + (position < remainder))] |
|
|
|
|
| def common_prefix(codes, rows): |
| if len(rows) == 1: |
| return int(codes[rows[0], :, 0].ne(-1).sum()) |
| same = codes[rows].eq(codes[rows[0]].unsqueeze(0)).all(-1).all(0) |
| same &= codes[rows[0], :, 0].ne(-1) |
| differing = (~same).nonzero() |
| return int(differing[0]) if len(differing) else codes.shape[1] |
|
|
|
|
| def reorder_codec(codec, indexes, n): |
| """Pinned codec state has per-row tensors, all with batch axis zero. |
| |
| Copy via index_select, including duplicates; never share mutable child KV. |
| Fail closed if a new codec revision adds an unsupported state shape. |
| """ |
| for module in codec.modules(): |
| state = getattr(module, '_streaming_state', None) |
| if state is None: |
| continue |
| for name, value in vars(state).items(): |
| if torch.is_tensor(value): |
| if value.ndim == 0 or value.shape[0] != n: |
| raise ValueError(f'Unsupported codec state tensor {name}: {tuple(value.shape)}') |
| setattr(state, name, value.index_select(0, indexes)) |
|
|
|
|
| class ChunkSearch: |
| def __init__(self, engine, scorer): |
| self.engine, self.scorer = engine, scorer |
|
|
| @torch.inference_mode() |
| def generate(self, spec, n=8, beam_width=1, chunk_frames=10, max_seconds=6.6, |
| guidance=1.5, temperature=1., seed=777, reference_path=None, on_commit=None): |
| if not 1 <= n <= 16 or beam_width not in [1,2]: |
| raise ValueError('Search supports batch sizes 1 through 16 and beam width 1/2') |
| requested_beam_width = beam_width |
| beam_width = min(n, beam_width) |
| if chunk_frames not in [10,20]: |
| raise ValueError('Search chunk must be 800 or 1600 ms') |
| engine, mdl = self.engine, self.engine.model |
| cfg = mdl.config |
| budget = min(750, math.ceil(max_seconds/.08)) |
| p = {'max_new_tokens':budget, 'audio_temperature':temperature, 'audio_top_p':.95, |
| 'audio_top_k':50, 'audio_repetition_penalty':1., 'text_temperature':0. if temperature==0 else 1., |
| 'text_top_p':1., 'text_top_k':50} |
| timeline, committed, first_commit = [], 0, None |
| started = time.perf_counter() |
| with engine.lock, torch.cuda.device(engine.device), torch.autocast('cuda',dtype=torch.bfloat16): |
| reference = engine.reference(reference_path) if reference_path else None |
| branches = [engine.pack(spec['prompt'],spec,reference,n)] |
| if guidance != 1 and spec['prompt'] != spec['unconditional']: |
| branches.append(engine.pack(spec['unconditional'],spec,reference,n)) |
| ids, initial = combine_branches(branches,cfg.pad_token_id,cfg.audio_pad_token_id) |
| nb, length = len(branches), ids.shape[1] |
| masks = torch.zeros(n*nb,length+budget,dtype=torch.bool,device=engine.device) |
| masks[:,:length] = initial |
| past, model_input, steps = None, ids, 0 |
| codes = torch.empty(n,0,12,dtype=torch.long,device=engine.device) |
| finished = torch.zeros(n,dtype=torch.bool,device=engine.device) |
| pcm = [np.empty(0,dtype=np.float32) for _ in range(n)] |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
| capture_start = engine.fast.capture_seconds |
| decoder = engine.proc.audio_tokenizer |
| with StreamingCodec(decoder,n,cfg.audio_pad_token_id) as streaming: |
| while steps < budget: |
| phase_start = time.perf_counter() |
| event_start,event_end = [torch.cuda.Event(enable_timing=True) for _ in range(2)] |
| event_start.record() |
| new = [] |
| for _ in range(min(chunk_frames,budget-steps)): |
| out = mdl.transformer(inputs_embeds=mdl._build_inputs_embeds(model_input),input_ids=None, |
| attention_mask=masks[:,:length+steps],past_key_values=past,use_cache=True, |
| position_ids=None,output_attentions=False,output_hidden_states=False, |
| return_dict=True,cu_seqlens=None,num_sequences=None) |
| past = out.past_key_values |
| local = mdl._global_hidden_to_local(out.last_hidden_state[:,-1]).to(mdl.local_transformer.ln_f.weight.dtype) |
| stop,sampled = engine.fast.frame(local,n,nb,guidance,p,graph=True) |
| continues = stop.eq(cfg.audio_assistant_slot_token_id) & ~finished |
| finished |= stop.eq(cfg.audio_end_token_id) |
| if not bool(continues.any()): |
| break |
| frame = sampled.masked_fill(~continues.unsqueeze(-1),cfg.audio_pad_token_id) |
| new.append(frame) |
| row = mdl._build_generation_row(n,ids.device,frame) |
| row[:,0,0] = torch.where(continues,row[:,0,0],cfg.pad_token_id) |
| model_input = row.repeat(nb,1,1) |
| masks[:,length+steps] = continues.repeat(nb) |
| steps += 1 |
| if new: |
| chunk = torch.stack(new,dim=1) |
| codes = torch.cat((codes,chunk),dim=1) |
| wave,lengths = streaming.decode(chunk) |
| for i in range(n): |
| pcm[i] = np.concatenate((pcm[i],wave[i,:int(lengths[i])].float().cpu().numpy())) |
| event_end.record() |
| event_end.synchronize() |
| generation_seconds = time.perf_counter()-phase_start |
| endpoint = steps>=budget or bool(finished.all()) or not new |
| |
| |
| with torch.autocast('cuda', enabled=False): |
| scores = self.scorer.score(pcm,spec['text'],[endpoint or bool(x) for x in finished.cpu()]) |
| order = sorted(range(n),key=lambda i:(-scores[i]['reward'],scores[i]['prefix_wer'],i)) |
| winners = order[:beam_width] if not endpoint else order[:1] |
| safe = common_prefix(codes,winners) |
| if safe < committed: |
| raise RuntimeError('Search attempted to revise committed acoustic prefix') |
| committed_seconds = max(0,safe-committed)*.08 |
| entry = {'step':len(timeline)+1,'frame':steps,'retained_rows':winners, |
| 'committed_audio_seconds':committed_seconds,'committed_until_seconds':safe*.08, |
| 'generation_wall_seconds':generation_seconds,'tts_codec_device_seconds':event_start.elapsed_time(event_end)/1000, |
| 'scores':scores, **self.scorer.last_timing} |
| if safe > committed: |
| if first_commit is None: |
| first_commit = time.perf_counter()-started |
| segment = pcm[winners[0]][committed*3840:safe*3840].copy() |
| if on_commit: |
| on_commit(segment,entry) |
| committed = safe |
| timeline.append(entry) |
| if endpoint: |
| winner = winners[0] |
| break |
| indexes = torch.tensor(expansion_indices(winners,n),device=engine.device) |
| branch_indexes = torch.cat([indexes+branch*n for branch in range(nb)]) |
| past = tuple(tuple(value.index_select(0,branch_indexes) for value in layer) for layer in past) |
| masks = masks.index_select(0,branch_indexes) |
| model_input = model_input.index_select(0,branch_indexes) |
| codes = codes.index_select(0,indexes) |
| finished = finished.index_select(0,indexes) |
| pcm = [pcm[i].copy() for i in indexes.cpu().tolist()] |
| reorder_codec(decoder,indexes,n) |
| else: |
| winner = 0 |
| output = pcm[winner] |
| engine.last_codes = [codes[winner][codes[winner,:,0].ne(cfg.audio_pad_token_id)].cpu()] |
| capture_seconds = engine.fast.capture_seconds-capture_start |
| total = time.perf_counter()-started |
| seconds = len(output)/48000 |
| if not seconds: |
| raise RuntimeError('Search generated no audio; try another seed/reference') |
| performance = {'n':n,'beam_width':beam_width,'requested_beam_width':requested_beam_width, |
| 'chunk_ms':chunk_frames*80,'max_seconds':max_seconds, |
| 'actual_seconds':seconds,'search_wall_seconds':total,'search_rtf':total/seconds, |
| 'first_commit_seconds':first_commit,'tts_codec_seconds':sum(x['tts_codec_device_seconds'] for x in timeline), |
| 'scoring_seconds':sum(x['total_seconds'] for x in timeline),'graph_capture_seconds':capture_seconds, |
| 'streaming_rule':'Winner commits immediately; top-2 emits only common acoustic-token prefixes', |
| 'includes':'reference encoding, packing, global/local TTS, codec, GPU transfers, batched ASR/timestamps, shared VoiceCLAP, heads, CPU selection and KV cloning', |
| 'excludes':'model loading, one-time style text-query embedding, Luna, MP3 encoding, final optional alignment/evaluation, network/playback', |
| 'device':str(engine.device),'scorer_device':str(self.scorer.device), |
| 'timeline':timeline,'final_search_score':scores[winner], |
| 'hit_frame_limit':steps>=budget and not bool(finished[winner])} |
| return output,performance |
|
|
