File size: 3,832 Bytes
268e727 | 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 80 81 82 83 84 85 86 87 88 89 90 | """Optional CUDA compilation; exact trained weights and decision contract.
Compilation is an explicit runtime choice. Cold compilation is substantial;
hosts should warm the five supported sequence buckets before accepting audio.
"""
from .router import Router
BUCKETS = (32, 64, 128, 256, 384)
def _cudagraph_failure(error):
"""Identify graph-runtime failures without inspecting request data."""
seen = set()
while error is not None and id(error) not in seen:
seen.add(id(error))
trace = error.__traceback__
while trace is not None:
module = trace.tb_frame.f_globals.get('__name__', '')
if module.startswith('torch._inductor.cudagraph'):
return True
trace = trace.tb_next
error = error.__cause__ or error.__context__
return False
class _BucketedModel:
"""Keep graph failure recovery under the caller's existing device admission."""
def __init__(self, native, compiled, pad_id):
import torch
self.native = native
self.compiled = compiled
self.pad_id = pad_id
self.thread_budget = torch.get_num_threads()
self.execution_mode = 'compiled-cuda'
def _recover(self, input_ids, attention_mask, observed, reason):
# Validate native execution before committing the persistent fallback.
output = self.native(input_ids, attention_mask, observed)
self.execution_mode = 'eager-cuda-fallback'
print(f'[jev] compiled CUDA unavailable ({reason}); '
'continuing with same-weight eager CUDA inference', flush=True)
return output
def __call__(self, input_ids, attention_mask, observed):
import torch
from torch.nn import functional as F
if input_ids.shape[0] != 2 or observed.shape != (1, 7):
raise ValueError('Compiled router is a single-decision runtime')
length = input_ids.shape[1]
bucket = next((n for n in BUCKETS if n >= length), None)
if bucket is None:
raise ValueError('Sequence exceeds decision contract')
if bucket > length:
input_ids = F.pad(input_ids, (0, bucket-length), value=self.pad_id)
attention_mask = F.pad(attention_mask, (0, bucket-length), value=0)
# Match the layout used by prewarming, including a size-one dimension.
observed = observed.reshape(7).unsqueeze(0)
if self.execution_mode == 'compiled-cuda':
# A provider can change this setting on an input worker. Avoid an
# expensive recompilation before the known graph TLS failure.
if torch.get_num_threads() != self.thread_budget:
return self._recover(input_ids, attention_mask, observed, 'thread-budget-changed')
try:
torch.compiler.cudagraph_mark_step_begin()
return self.compiled(input_ids, attention_mask, observed)
except Exception as error:
if isinstance(error, torch.OutOfMemoryError) or not _cudagraph_failure(error):
raise
return self._recover(input_ids, attention_mask, observed, type(error).__name__)
return self.native(input_ids, attention_mask, observed)
class CompiledCudaRouter(Router):
def __init__(self, path, *, device='cuda:0', admission=None):
if not device.startswith('cuda'):
raise ValueError('Compiled runtime is CUDA-only')
super().__init__(path, backend='eager', device=device, admission=admission)
def _load(self):
import torch
super()._load()
native = self.model
compiled = torch.compile(native, mode='reduce-overhead', dynamic=False)
pad_id = native.encoder.config.pad_token_id
self.model = _BucketedModel(native, compiled, pad_id)
|