| """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 |
| |
| 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 |
| |
| |
| 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) |
| |
| 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 |
|
|
