"""Opt-in batched native inference. No model, prompt, precision or sampling changes.""" import gc import json import time from datetime import UTC, datetime from importlib.metadata import version def require_fast_kernels(): """Reject the accelerated profile before weights download if CUDA kernels are missing.""" import torch from transformers.models.qwen3_5 import modeling_qwen3_5 if not torch.cuda.is_available() or not modeling_qwen3_5.is_fast_path_available: raise RuntimeError("Accelerated Flora needs CUDA, flash-linear-attention and causal-conv1d") if torch.cuda.get_device_capability(0) != (9, 0): raise RuntimeError("This accelerated Flora image is compiled for H200/Hopper SM90") return { name: version(name) for name in ("torch", "transformers", "flash-linear-attention", "causal-conv1d") } def collate_inputs(requests, pad_token_id, *, padding_multiple=1, minimum_width=0): """Left-pad separate sequences; concatenate their video tensors in the same order.""" import torch import torch.nn.functional as functional sequence_keys = {"input_ids", "attention_mask", "mm_token_type_ids"} media_keys = {"pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw"} keys = set(requests[0]) - {"video_metadata"} if not {"input_ids", "attention_mask"} <= keys or keys - sequence_keys - media_keys: raise ValueError("Unsupported native Flora input schema") if any(set(item) - {"video_metadata"} != keys for item in requests): raise ValueError("Only matching native input schemas may be batched") width = max(minimum_width, max(item["input_ids"].shape[1] for item in requests)) width = (width + padding_multiple - 1) // padding_multiple * padding_multiple result = {} for key in keys: tensors = [] for item in requests: tensor = item[key] if key in sequence_keys: if tensor.ndim != 2 or tensor.shape != item["input_ids"].shape or tensor.shape[0] != 1: raise ValueError("Each native request must contain one independent sequence") tensor = functional.pad( tensor, (width - tensor.shape[1], 0), value=pad_token_id if key == "input_ids" else 0 ) tensors.append(tensor) result[key] = torch.cat(tensors, dim=0) return result def decode_rows(output, width, tokenizer): """Validate each row's own EOS, including rows padded after an early finish.""" eos = tokenizer.eos_token_id ends = {eos} if isinstance(eos, int) else set(eos) captions = [] for row in output[:, width:].cpu().tolist(): stop = next((index for index, value in enumerate(row) if value in ends), None) if stop is None: raise ValueError("Flora returned an incomplete batched output") caption = tokenizer.decode(row[:stop], skip_special_tokens=True).strip() if not caption: raise ValueError("Flora returned an empty batched output") captions.append(caption) return captions def generate_many( model, processor, requests, output_limit, *, batch_size=4, max_padded_tokens=16384, max_padding_ratio=1.25, padding_multiple=1, pad_batch_power2=False, minimum_width=0, fixed_batch_size=0, check_deadline=lambda: None, ): """Batch similar lengths with bounded padding; split OOM batches without dropping requests.""" import torch if not requests: return [] if not isinstance(batch_size, int) or isinstance(batch_size, bool) or not 1 <= batch_size <= 64: raise ValueError("Flora generation batch size must be between 1 and 64") if type(max_padded_tokens) is not int or not 16384 <= max_padded_tokens <= 131072: raise ValueError("Flora padded-token budget is outside the reviewed H200 bounds") if type(max_padding_ratio) not in (int, float) or not 1 <= max_padding_ratio <= 8: raise ValueError("Flora padding ratio is outside the reviewed bounds") if ( type(padding_multiple) is not int or padding_multiple not in (1, 128) or type(pad_batch_power2) is not bool or type(minimum_width) is not int or minimum_width not in (0, 128, 1792) or type(fixed_batch_size) is not int or fixed_batch_size not in (0, 16, 64) or fixed_batch_size > batch_size ): raise ValueError("Unreviewed Flora shape padding") tokenizer = processor.tokenizer if tokenizer.pad_token_id is None: raise ValueError("Native tokenizer needs an explicit padding token") order = sorted(range(len(requests)), key=lambda index: requests[index]["input_ids"].shape[1]) results = [None] * len(requests) def run(indices, allow_fixed=True): check_deadline() batch = [requests[i] for i in indices] padded_count = 1 << (len(batch) - 1).bit_length() if pad_batch_power2 else len(batch) if allow_fixed: padded_count = max(padded_count, fixed_batch_size) raw_width = max(item["input_ids"].shape[1] for item in batch) multiple = padding_multiple if raw_width + padding_multiple - 1 + output_limit <= 32768 else 1 padded_width = (max(raw_width, minimum_width) + multiple - 1) // multiple * multiple if len(batch) * padded_width > max_padded_tokens: multiple, padded_width = 1, raw_width if padded_count * padded_width > max_padded_tokens: padded_count = len(batch) batch += [batch[-1]] * (padded_count - len(batch)) inputs = collate_inputs(batch, tokenizer.pad_token_id, minimum_width=padded_width) width = inputs["input_ids"].shape[1] if width + output_limit > 32768: raise ValueError("Flora input exceeds the native context budget") inputs = {key: value.to(next(model.parameters()).device) for key, value in inputs.items()} started = time.monotonic() failure = None try: with torch.inference_mode(): output = model.generate( **inputs, do_sample=False, num_beams=1, max_new_tokens=output_limit, eos_token_id=tokenizer.eos_token_id, pad_token_id=tokenizer.pad_token_id, use_cache=True, max_time=300, ) values = decode_rows(output[: len(indices)], width, tokenizer) if len(values) != len(indices): raise ValueError("Flora output count differs from the request count") for index, value in zip(indices, values, strict=True): results[index] = value print( "FLORA_BATCH " + json.dumps( { "event": "generation_end", "utc": datetime.now(UTC).isoformat(), "batch_size": len(indices), "padded_batch_size": padded_count, "input_width": width, "output_limit": output_limit, "generated_tokens": [ int((row == tokenizer.eos_token_id).nonzero()[0]) + 1 for row in output[:, width:] ] if isinstance(tokenizer.eos_token_id, int) else None, "elapsed_s": round(time.monotonic() - started, 3), } ), flush=True, ) except torch.OutOfMemoryError as exc: if len(indices) == 1: raise # Do not retain the failing frame and its GPU tensors during retry. exc.__traceback__ = None failure = "oom_split" if failure: del inputs gc.collect() torch.cuda.empty_cache() print("FLORA_BATCH " + json.dumps({"event": failure, "batch_size": len(indices)}), flush=True) middle = len(indices) // 2 # An OOM retry must actually shrink the allocation; dummy rows may # never restore the fixed outer batch size during recursive recovery. run(indices[:middle], allow_fixed=False) run(indices[middle:], allow_fixed=False) group = [] for index in order: width = max(requests[index]["input_ids"].shape[1], minimum_width) # Avoid a long request forcing excessive padding on short video clips. if group and ( len(group) >= batch_size or width > max_padding_ratio * max(minimum_width, requests[group[0]]["input_ids"].shape[1]) or width * (len(group) + 1) > max_padded_tokens ): run(group) group = [] group.append(index) if group: run(group) return results