| """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): |
| |
| 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) |
| |
| observed = observed.reshape(7).unsqueeze(0) |
| if self.execution_mode == 'compiled-cuda': |
| |
| |
| 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) |
|
|
