diff --git a/extensions_built_in/captioner/Qwen32Captioner.py b/extensions_built_in/captioner/Qwen32Captioner.py new file mode 100644 index 0000000..e69de29 diff --git a/extensions_built_in/captioner/Qwen38Captioner.py b/extensions_built_in/captioner/Qwen38Captioner.py new file mode 100644 index 0000000..67cb36f --- /dev/null +++ b/extensions_built_in/captioner/Qwen38Captioner.py @@ -0,0 +1,164 @@ +"""Qwen3.8-27B image / video captioner on the qwen38 arch from +extensions_built_in/llm_models: ComfyUI W4A8 weights kept in their quantization, +custom DeltaNet + attention stack, CUDA-graph decode and optional MTP speculative +decoding. LoRAs trained with the qwen38 arch attach as sidechains (BaseCaptioner). +Text only in the output; audio files are not captioned (use Qwen2.5/3-Omni).""" + +from collections import OrderedDict + +import concurrent.futures +import os +import traceback +from collections import deque + +import torch + +from .BaseCaptioner import BaseCaptioner, CaptionConfig +import logging +import warnings + +warnings.filterwarnings("ignore") +logging.disable(logging.WARNING) + +IMAGE_EXTENSIONS = {"jpg", "jpeg", "png", "bmp", "webp"} +# reasoning budget on top of max_new_tokens when thinking is on +MAX_THINKING_TOKENS = 4096 + + +class Qwen38CaptionConfig(CaptionConfig): + def __init__(self, **kwargs): + super().__init__(**kwargs) + # MTP speculative decoding (draft depth tokens per step, verified in one pass) + self.mtp = bool(kwargs.get("mtp", True)) + self.mtp_depth = int(kwargs.get("mtp_depth", 3)) + # capture the decode step as a CUDA graph (off = eager, for debugging) + self.cuda_graph = bool(kwargs.get("cuda_graph", True)) + # vision tower on dequantized bf16 weights (int8 activations drift image embeddings ~6%) + self.vision_bf16 = bool(kwargs.get("vision_bf16", True)) + # frames per second sampled from videos + self.video_fps = float(kwargs.get("video_fps", 2)) + self.system_prompt = kwargs.get("system_prompt", None) or None + + +class Qwen38Captioner(BaseCaptioner): + caption_config_class = Qwen38CaptionConfig + + def __init__(self, process_id: int, job, config: OrderedDict, **kwargs): + super(Qwen38Captioner, self).__init__(process_id, job, config, **kwargs) + self.llm = None + + def load_model(self): + from toolkit.config_modules import ModelConfig + from extensions_built_in.llm_models.qwen38 import Qwen38LLM + + cc = self.caption_config + if cc.low_vram: + print("[AITK] low_vram has no effect for Qwen3.8 (W4A8 weights are ~17 GB and stay resident).") + model_config = ModelConfig( + arch="qwen38", + name_or_path=cc.model_name_or_path, + quantize=cc.quantize, + qtype=cc.qtype, + model_kwargs={ + "instruction": cc.caption_prompt, + "system_prompt": cc.system_prompt, + "enable_thinking": cc.thinking, + "max_pixels": cc.max_res * cc.max_res, + "max_new_tokens": cc.max_new_tokens + (MAX_THINKING_TOKENS if cc.thinking else 0), + "mtp_depth": cc.mtp_depth if cc.mtp else 0, + "cuda_graph": cc.cuda_graph, + "vision_bf16": cc.vision_bf16, + "video_fps": cc.video_fps, + }, + ) + llm = Qwen38LLM(self.device_torch, model_config, dtype=cc.dtype) + llm.add_status_update_hook(lambda s: self.update_status("running", s)) + llm.load_model() + self.llm = llm + # LoRA sidechains and the stop/compile machinery key off self.model + self.model = llm.model + self.processor = llm.processor + mode = f"MTP depth {cc.mtp_depth}" if cc.mtp else "plain decode" + print(f"[AITK] Qwen3.8 ready: {mode}, cuda graphs {'on' if cc.cuda_graph else 'off'}") + + def maybe_compile_models(self): + # the decode step is already a captured CUDA graph; torch.compile would only + # re-trace the custom stack + return + + @staticmethod + def _is_image_file(file_path: str) -> bool: + return os.path.splitext(file_path)[1].lower().lstrip(".") in IMAGE_EXTENSIONS + + def _prep_media(self, file_path: str) -> dict: + """CPU side of one file (image decode / frame extraction, processor); runs in a + worker thread ahead of the GPU.""" + item = self.llm.prepare_caption_inputs(file_path, self.caption_config.caption_prompt) + item["file"] = file_path + return item + + def _generate(self, item: dict) -> str: + out_ids = self.llm.generate_from_inputs(item) + return self.llm.decode_text(out_ids) + + def run_caption_loop(self): + """Worker threads preprocess files ahead of the GPU; generation runs one file + at a time on the main thread (the decode step is weight-bandwidth bound and + CUDA-graphed, so there is no CPU slack to hide).""" + import tqdm as tqdm_mod + + pbar = tqdm_mod.tqdm(total=len(self.file_paths), desc="Captioning files", unit="file", smoothing=0.9) + + def finish(file_path, caption): + if caption is not None: + self.save_caption_for_file(file_path, caption) + self.step_num += 1 + self.update_step() + pbar.update(1) + + executor = concurrent.futures.ThreadPoolExecutor(max_workers=max(1, int(self.caption_config.num_workers))) + try: + futures = deque() + file_iter = iter(self.file_paths) + lookahead = max(2, int(self.caption_config.num_workers)) + 1 + for _ in range(lookahead): + path = next(file_iter, None) + if path is None: + break + futures.append((path, executor.submit(self._prep_media, path))) + while futures: + if self.is_ui_captioner: + self.maybe_stop() + if self.is_stopping: + break + path, fut = futures.popleft() + nxt = next(file_iter, None) + if nxt is not None: + futures.append((nxt, executor.submit(self._prep_media, nxt))) + try: + item = fut.result() + except Exception as e: + print(f"Error preprocessing {path}: {e}") + finish(path, None) + continue + try: + caption = self._generate(item) + st = self.llm.last_generation_stats + if st.get("seconds"): + pbar.set_postfix_str(f"{st['tokens'] / st['seconds']:.1f} tok/s, prompt {st['prompt_tokens']}") + except Exception as e: + print(f"Error captioning {path}: {e}") + traceback.print_exc() + caption = None + finish(path, caption) + finally: + executor.shutdown(wait=False, cancel_futures=True) + pbar.close() + + def get_caption_for_file(self, file_path: str) -> str: + try: + return self._generate(self._prep_media(file_path)) + except Exception as e: + print(f"Error processing {file_path}: {e}") + traceback.print_exc() + return None diff --git a/extensions_built_in/captioner/__init__.py b/extensions_built_in/captioner/__init__.py index e42683f..c5194df 100644 --- a/extensions_built_in/captioner/__init__.py +++ b/extensions_built_in/captioner/__init__.py @@ -61,6 +61,18 @@ class Qwen25OmniCaptionerExtension(Extension): return Qwen25OmniCaptioner +class Qwen38CaptionerExtension(Extension): + uid = "Qwen38Captioner" + name = "Qwen 3.8 Captioner" + + @classmethod + def get_process(cls): + # import your process class here so it is only loaded when needed and return it + from .Qwen38Captioner import Qwen38Captioner + + return Qwen38Captioner + + class Ideogram4CaptionerExtension(Extension): uid = "Ideogram4Captioner" name = "Ideogram4 Captioner" @@ -79,5 +91,6 @@ AI_TOOLKIT_EXTENSIONS = [ Qwen3VLCaptionerExtension, Qwen3OmniCaptionerExtension, Qwen25OmniCaptionerExtension, + Qwen38CaptionerExtension, Ideogram4CaptionerExtension, ] diff --git a/extensions_built_in/llm_models/__init__.py b/extensions_built_in/llm_models/__init__.py index 393325f..14b0b7f 100644 --- a/extensions_built_in/llm_models/__init__.py +++ b/extensions_built_in/llm_models/__init__.py @@ -1,5 +1,7 @@ from .qwen25_omni import Qwen25OmniLLM +from .qwen38 import Qwen38LLM AI_TOOLKIT_MODELS = [ Qwen25OmniLLM, + Qwen38LLM, ] diff --git a/extensions_built_in/llm_models/qwen38.py b/extensions_built_in/llm_models/qwen38.py new file mode 100644 index 0000000..7c9cc6e --- /dev/null +++ b/extensions_built_in/llm_models/qwen38.py @@ -0,0 +1,421 @@ +"""Qwen3.8-27B as a trainable text-generation model (arch ``qwen38``). + +Image or video in (plus a .txt caption), text out. The frozen vision tower encodes +the media into embeddings [N, H] plus the processor's expanded placeholder span and +the MRoPE grid metadata (cached as the "latent" when latent caching is on). Training +runs the text stack on ``prefix + media span + instruction + assistant header + +caption`` with next-token cross-entropy on the caption (LoRA on the text stack; +towers, embeddings and lm_head stay frozen). ``is_llm`` routes the trainer to +``train_llm_accumulation`` -> ``get_llm_loss``: no noise, scheduler, VAE or prompt +encoding. + +Weights: the ComfyUI single file (W4A8 codebook language model, int8 convrot +embeddings / lm_head / MTP / vision) is loaded with its quantization kept as-is +through the toolkit's convrotw4a8 / convrot8 backends, or an HF layout (bf16 shards) +is loaded and quantized to the requested qtype. See src/qwen38_model.py for the +model itself (static caches, CUDA-graph decode, MTP speculative decoding). + +Samples: ``ctrl_img`` is the media file, the sample prompt is the instruction; the +generated text is written as a .txt. + +model_kwargs: + instruction user-turn text used for every training item (default below) + system_prompt optional system message + enable_thinking keep the model's thinking mode on (default false: direct answers) + max_pixels/min_pixels image and video frame budget for the vision tower + video_fps fps the dataset extracted video frames at (dataset fps) + max_new_tokens sample generation budget (default 512) + mtp_depth MTP speculative draft depth for samples (default 3, 0 = off) + cuda_graph CUDA-graph the decode step (default true) + vision_bf16 run the vision tower with bf16 activations on its dequantized int8 + weights (default true; int8 activations drift the image embeddings ~6%) + loss_chunk positions per lm_head chunk in the CE (default 512) +""" + +import os +import time +from typing import List, Optional + +import numpy as np +import torch +import torch.nn.functional as F +import torch.utils.checkpoint +from transformers import AutoConfig, AutoProcessor + +from toolkit.basic import flush +from toolkit.config_modules import GenerateImageConfig, ModelConfig +from toolkit.dto import DTO +from toolkit.models.base_model import BaseModel +from toolkit.models.v2.resolver import resolve_component_file +from toolkit.print import print_acc + +from .src.qwen38_model import Qwen38Generator, build_sequence_embeds + +BASE_REPO = "Qwen/Qwen3.8-27B" +DEFAULT_CHECKPOINT = "Comfy-Org/Qwen3.8-27B/text_encoders/qwen3.8_27b_w4a8.safetensors" +DEFAULT_INSTRUCTION = "Describe this in detail." + +KIND_IMAGE, KIND_VIDEO = 1, 2 +VIDEO_EXTS = (".mp4", ".avi", ".mov", ".webm", ".mkv", ".wmv", ".m4v", ".flv") +AUDIO_EXTS = (".mp3", ".wav", ".flac", ".aac", ".ogg", ".m4a") + + +def _media_kind(path: str) -> int: + ext = os.path.splitext(path)[1].lower() + if ext in AUDIO_EXTS: + raise ValueError(f"qwen38 takes images and videos, not audio: {path}") + return KIND_VIDEO if ext in VIDEO_EXTS else KIND_IMAGE + + +class Qwen38LLM(BaseModel): + arch = "qwen38" + is_llm = True + # the data loader accepts image and video files in one dataset + is_multimodal_llm = True + sample_rate = 16000 + + def __init__(self, device, model_config: ModelConfig, dtype="bf16", custom_pipeline=None, noise_scheduler=None, **kwargs): + super().__init__(device, model_config, dtype, custom_pipeline, noise_scheduler, **kwargs) + self.is_transformer = True + # LoRA goes on the text stack only (vision tower, embeddings and lm_head stay frozen) + self.target_lora_modules = ["Qwen38TextModel"] + kw = self.model_config.model_kwargs + self.instruction = str(kw.get("instruction", DEFAULT_INSTRUCTION)) + self.system_prompt: Optional[str] = kw.get("system_prompt", None) + self.enable_thinking = bool(kw.get("enable_thinking", False)) + self.max_pixels = int(kw.get("max_pixels", 1024 * 1024)) + self.min_pixels = int(kw.get("min_pixels", 64 * 32 * 32)) + self.video_fps = float(kw.get("video_fps", 24)) + self.max_new_tokens = int(kw.get("max_new_tokens", 512)) + self.mtp_depth = int(kw.get("mtp_depth", 3)) + self.cuda_graph = bool(kw.get("cuda_graph", True)) + self.vision_bf16 = bool(kw.get("vision_bf16", True)) + self.loss_chunk = int(kw.get("loss_chunk", 512)) + self.debug = bool(kw.get("debug", False)) + self.processor = None + self.additional_loss_logs = {} + + @staticmethod + def get_train_scheduler(): + # no diffusion: the trainer never samples timesteps for this model + return None + + # ------------------------------------------------------------------ + # loading + # ------------------------------------------------------------------ + def _config_source(self) -> str: + extras = self.model_config.extras_name_or_path + if extras and extras != self.model_config.name_or_path and not extras.endswith(".safetensors"): + return extras + return BASE_REPO + + def load_model(self): + dtype = self.torch_dtype + device = self.device_torch + name_or_path = self.model_config.name_or_path or DEFAULT_CHECKPOINT + cfg_src = self._config_source() + want_quant = self.model_config.quantize + qtype = self.model_config.qtype + if self.model_config.layer_offloading and self.model_config.layer_offloading_transformer_percent > 0: + raise NotImplementedError("Layer offloading is not implemented for qwen38") + + from .src.qwen38_load import dequantize_vision_tower, load_qwen38 + + if name_or_path.endswith(".safetensors"): + ckpt = resolve_component_file(name_or_path, folder="text_encoders", component="qwen3.8") + self.print_and_status_update(f"Loading Qwen3.8 from {os.path.basename(ckpt)}") + source = ckpt + else: + self.print_and_status_update(f"Loading Qwen3.8 from {name_or_path}") + source = name_or_path + if os.path.isdir(name_or_path) or "/" in name_or_path: + cfg_src = name_or_path if self.model_config.extras_name_or_path in (None, "", name_or_path) else cfg_src + config = AutoConfig.from_pretrained(cfg_src) + model, prequantized = load_qwen38(source, config, dtype, device, status_fn=self.print_and_status_update) + if prequantized: + model.aitk_is_quantized = True + if want_quant and qtype not in ("convrotw4a8", "convrot8"): + print_acc(f"Checkpoint is pre-quantized (W4A8 / int8 convrot); keeping it instead of {qtype}") + want_quant = False + if self.vision_bf16 and prequantized: + n = dequantize_vision_tower(model, dtype) + print_acc(f"Vision tower: {n} int8 linears dequantized to {dtype} (model_kwargs.vision_bf16)") + model.eval() + model.requires_grad_(False) + if want_quant: + from toolkit.util.quantize import get_qtype, quantize + + self.print_and_status_update(f"Quantizing Qwen3.8 ({qtype})") + quantize(model, weights=get_qtype(qtype)) + model.aitk_is_quantized = True + flush() + self.model = model + self.model.generator = Qwen38Generator(model, use_cuda_graph=self.cuda_graph) + self.processor = AutoProcessor.from_pretrained(cfg_src) + self.tokenizer = self.processor.tokenizer + self.text_encoder = None + self.vae = None + self.pipeline = self + self._build_template() + flush() + + def _template_kwargs(self): + return {"enable_thinking": self.enable_thinking} + + def _build_template(self): + """Split the chat template once around the media span and the instruction so + training sequences can be assembled from cached media spans without the processor.""" + mark = "<>" + msgs = [] + if self.system_prompt is not None: + msgs.append({"role": "system", "content": self.system_prompt}) + msgs.append({"role": "user", "content": [{"type": "image", "image": "x"}, {"type": "text", "text": mark}]}) + text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, **self._template_kwargs()) + span = "<|vision_start|><|image_pad|><|vision_end|>" + if span not in text or mark not in text: + raise RuntimeError("Unexpected Qwen3.8 chat template; cannot locate the media span") + before, rest = text.split(span, 1) + mid, after = rest.split(mark, 1) + enc = lambda t: self.tokenizer(t, add_special_tokens=False).input_ids + self._prefix_ids = enc(before) + self._mid_ids = enc(mid) + self._after_ids = enc(after) + self._probe = { + KIND_IMAGE: text.replace(mark, "x"), + KIND_VIDEO: text.replace(span, "<|vision_start|><|video_pad|><|vision_end|>").replace(mark, "x"), + } + self._tail_ids = enc(mid + "x" + after) + cfg = self.model.config + self._placeholder = {KIND_IMAGE: cfg.image_token_id, KIND_VIDEO: cfg.video_token_id} + self._im_end = self.tokenizer.convert_tokens_to_ids("<|im_end|>") + self._eos_ids = [self._im_end, self.tokenizer.convert_tokens_to_ids("<|endoftext|>")] + + def get_transformer_block_names(self) -> Optional[List[str]]: + return ["model.language_model.layers"] + + def get_model_has_grad(self): + return False + + def get_te_has_grad(self): + return False + + def get_bucket_divisibility(self): + return 32 + + def save_model(self, output_path, meta, save_dtype): + raise NotImplementedError("qwen38: only LoRA training is supported (no full-model save)") + + # LoRA keys are model-relative (model.language_model.layers.N....), the HF layout + def convert_lora_weights_before_save(self, state_dict): + return {(k[len("transformer.") :] if k.startswith("transformer.") else k): v for k, v in state_dict.items()} + + def convert_lora_weights_before_load(self, state_dict): + return {(k if k.startswith("transformer.") else "transformer." + k): v for k, v in state_dict.items()} + + # ------------------------------------------------------------------ + # media encoding (latent cache) + # ------------------------------------------------------------------ + def _span(self, input_ids: torch.Tensor, types: torch.Tensor): + """The processor-expanded media span (tokens between the template prefix and + the instruction tail), with its token-type ids.""" + ids = input_ids[0] + n = ids.shape[0] + p, t = len(self._prefix_ids), len(self._tail_ids) + if ids[:p].tolist() != self._prefix_ids or ids[n - t :].tolist() != self._tail_ids: + raise RuntimeError("Processor output does not start/end with the template prefix/tail") + return ids[p : n - t].to(torch.int32), types[0, p : n - t].to(torch.int32) + + def _pack(self, embeds, media_ids, media_types, kind, grid): + dev = embeds.device + return DTO( + embeds[None].to(self.torch_dtype), + media_ids=media_ids[None].to(dev), + media_types=media_types[None].to(dev), + kind=torch.tensor([kind], dtype=torch.int32, device=dev), + grid_thw=torch.tensor([list(grid)], dtype=torch.int32, device=dev), + ) + + @torch.no_grad() + def encode_images(self, image_list: torch.Tensor, device=None, dtype=None): + """[1, 3, H, W] image in [-1, 1] | [1, F, 3, H, W] frames.""" + if image_list.shape[0] != 1: + raise ValueError("qwen38 encodes one item at a time (batch_size 1)") + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + x = image_list[0] + if x.ndim == 2: + raise ValueError("qwen38 does not take audio") + if x.ndim == 3: + return self._encode_image_item(x) + if x.ndim == 4: + return self._encode_video_item(x) + raise ValueError(f"Unexpected media tensor shape {tuple(image_list.shape)}") + + def encode_audio(self, audio_data_list): + # the data loader offers a video's own audio track here; video is vision-only + return torch.zeros(len(audio_data_list), 1) + + @staticmethod + def _to_uint8(img: torch.Tensor) -> np.ndarray: + return ((img.float().clamp(-1, 1) + 1.0) * 127.5).round().to(torch.uint8).permute(1, 2, 0).cpu().numpy() + + def _encode_image_item(self, img: torch.Tensor): + from PIL import Image + + pil = Image.fromarray(self._to_uint8(img)) + inputs = self.processor(text=self._probe[KIND_IMAGE], images=[pil], return_tensors="pt", min_pixels=self.min_pixels, max_pixels=self.max_pixels) + grid = inputs["image_grid_thw"] + feats = self.model.get_image_features(inputs["pixel_values"].to(self.device_torch), grid.to(self.device_torch)) + ids, types = self._span(inputs["input_ids"], inputs["mm_token_type_ids"]) + return self._pack(feats, ids, types, KIND_IMAGE, grid[0].tolist()) + + def _encode_video_item(self, frames: torch.Tensor): + video = np.stack([self._to_uint8(f) for f in frames]) # [F, H, W, 3] + inputs = self.processor( + text=self._probe[KIND_VIDEO], videos=[video], return_tensors="pt", fps=self.video_fps, + min_pixels=self.min_pixels, max_pixels=self.max_pixels, + ) + grid = inputs["video_grid_thw"] + feats = self.model.get_image_features(inputs["pixel_values_videos"].to(self.device_torch), grid.to(self.device_torch)) + ids, types = self._span(inputs["input_ids"], inputs["mm_token_type_ids"]) + return self._pack(feats, ids, types, KIND_VIDEO, grid[0].tolist()) + + # ------------------------------------------------------------------ + # training + # ------------------------------------------------------------------ + def _sequence(self, media_ids: torch.Tensor, media_types: torch.Tensor, caption: str): + """ids [L], token types [L] and the index where the caption (loss span) starts.""" + instr = self.tokenizer(self.instruction, add_special_tokens=False).input_ids + target = self.tokenizer(caption, add_special_tokens=False).input_ids + [self._im_end] + head = self._prefix_ids + media_ids.tolist() + self._mid_ids + instr + self._after_ids + ids = torch.tensor(head + target, dtype=torch.long, device=self.device_torch) + types = torch.zeros_like(ids) + p = len(self._prefix_ids) + types[p : p + media_ids.shape[0]] = media_types.to(types.dtype) + return ids, types, len(head) + + def _lm_loss(self, ids, types, start: int, media: torch.Tensor, kind: int, grid): + model = self.model + embeds = build_sequence_embeds(model, ids, media, [self._placeholder[kind]]) + grid_t = torch.tensor([list(grid)], device=ids.device) + position_ids, _ = model.get_rope_index( + ids, types, image_grid_thw=grid_t if kind == KIND_IMAGE else None, video_grid_thw=grid_t if kind == KIND_VIDEO else None + ) + hidden = model.forward_hidden(embeds, position_ids)[0, start - 1 : -1] + targets = ids[start:] + lm_head = model.lm_head + + def chunk_ce(h, t): + return F.cross_entropy(lm_head(h).float(), t, reduction="sum") + + total = hidden.new_zeros((), dtype=torch.float32) + for s0 in range(0, targets.shape[0], self.loss_chunk): + h, t = hidden[s0 : s0 + self.loss_chunk], targets[s0 : s0 + self.loss_chunk] + if torch.is_grad_enabled(): + total = total + torch.utils.checkpoint.checkpoint(chunk_ce, h, t, use_reentrant=False) + else: + total = total + chunk_ce(h, t) + return total / targets.shape[0] + + def get_llm_loss(self, batch) -> torch.Tensor: + """Next-token cross-entropy on the caption, from the cached media embeddings. + Called by the trainer's train_llm_accumulation; logs ride in additional_loss_logs.""" + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + lat = batch.latents + if not isinstance(lat, DTO) or lat.get("media_ids") is None: + if batch.tensor is None: + raise ValueError("qwen38 got a batch with neither media tensors nor cached embeddings") + lat = self.encode_images(batch.tensor) + if lat.shape[0] != 1: + raise ValueError("qwen38 trains with batch_size 1") + captions = batch.get_caption_list() + media = lat.tensor[0].to(self.device_torch, self.torch_dtype) + kind = int(lat.get("kind")[0]) + ids, types, start = self._sequence(lat.get("media_ids")[0].long(), lat.get("media_types")[0].long(), captions[0]) + ce = self._lm_loss(ids, types, start, media, kind, lat.get("grid_thw")[0].tolist()) + if self.debug: + print_acc( + f"qwen38: seq {ids.shape[0]} target {ids.shape[0] - start} ce {ce.item():.4f} " + f"peak {torch.cuda.max_memory_allocated() / 1e9:.1f} GB" + ) + self.additional_loss_logs = {"loss/ce": ce.detach().float().item()} + return ce + + # ------------------------------------------------------------------ + # sampling: ctrl_img is the media file, prompt is the instruction, output is text + # ------------------------------------------------------------------ + def get_generation_pipeline(self): + return self + + def _load_sample_media(self, path: str): + kind = _media_kind(path) + if kind == KIND_IMAGE: + from PIL import Image + + return kind, {"type": "image", "image": "x"}, {"images": [Image.open(path).convert("RGB")]} + from transformers.video_utils import load_video + + frames = load_video(path, fps=self.video_fps) + if isinstance(frames, tuple): + frames = frames[0] + return kind, {"type": "video", "video": "x"}, {"videos": [frames], "fps": self.video_fps} + + def prepare_caption_inputs(self, media_path: Optional[str], instruction: Optional[str] = None) -> dict: + """CPU half of caption(): media decode + processor. Safe in a worker thread.""" + instruction = instruction.strip() if instruction and instruction.strip() else self.instruction + content = [] + media_kwargs = {} + kind = None + if media_path: + kind, entry, media_kwargs = self._load_sample_media(media_path) + content.append(entry) + content.append({"type": "text", "text": instruction}) + msgs = ([{"role": "system", "content": self.system_prompt}] if self.system_prompt is not None else []) + [{"role": "user", "content": content}] + text = self.processor.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False, **self._template_kwargs()) + inputs = self.processor(text=text, return_tensors="pt", min_pixels=self.min_pixels, max_pixels=self.max_pixels, **media_kwargs) + return {"kind": kind, "inputs": inputs} + + @torch.no_grad() + def generate_from_inputs(self, prepared: dict, max_new_tokens: Optional[int] = None, mtp_depth: Optional[int] = None) -> List[int]: + """GPU half of caption(): vision tower, prefill and decode. Returns token ids.""" + if self.model.device == torch.device("cpu"): + self.model.to(self.device_torch) + kind, inputs = prepared["kind"], prepared["inputs"] + ids = inputs["input_ids"][0].to(self.device_torch) + types = inputs["mm_token_type_ids"][0].to(self.device_torch) + model = self.model + media = None + image_grid = video_grid = None + if kind == KIND_IMAGE: + image_grid = inputs["image_grid_thw"].to(self.device_torch) + media = model.get_image_features(inputs["pixel_values"].to(self.device_torch), image_grid) + elif kind == KIND_VIDEO: + video_grid = inputs["video_grid_thw"].to(self.device_torch) + media = model.get_image_features(inputs["pixel_values_videos"].to(self.device_torch), video_grid) + embeds = build_sequence_embeds(model, ids, media, [self._placeholder[KIND_IMAGE], self._placeholder[KIND_VIDEO]]) + position_ids, delta = model.get_rope_index(ids, types, image_grid_thw=image_grid, video_grid_thw=video_grid) + t0 = time.perf_counter() + out_ids = model.generator.generate( + embeds, ids[None], position_ids, delta, + max_new_tokens=self.max_new_tokens if max_new_tokens is None else max_new_tokens, + eos_ids=self._eos_ids, mtp_depth=self.mtp_depth if mtp_depth is None else mtp_depth, + ) + self.last_generation_stats = dict(model.generator.last_stats, seconds=time.perf_counter() - t0, prompt_tokens=int(ids.shape[0])) + return out_ids + + def decode_text(self, out_ids: List[int]) -> str: + text = self.tokenizer.decode(out_ids, skip_special_tokens=True) + # thinking mode: keep only the answer after the think block + if "" in text: + text = text.split("")[-1] + return text.strip() + + @torch.no_grad() + def caption(self, media_path: Optional[str], instruction: Optional[str] = None, max_new_tokens: Optional[int] = None, mtp_depth: Optional[int] = None) -> str: + """Generate text for one media file (or none). Also the sample path.""" + return self.decode_text(self.generate_from_inputs(self.prepare_caption_inputs(media_path, instruction), max_new_tokens, mtp_depth)) + + def generate_single_image(self, pipeline, gen_config: GenerateImageConfig, conditional_embeds, unconditional_embeds, generator, extra): + # a str return is saved as .txt by GenerateImageConfig.save_image + return self.caption(gen_config.ctrl_img, gen_config.prompt) diff --git a/extensions_built_in/llm_models/src/qwen38_kernels.py b/extensions_built_in/llm_models/src/qwen38_kernels.py new file mode 100644 index 0000000..1db88ed --- /dev/null +++ b/extensions_built_in/llm_models/src/qwen38_kernels.py @@ -0,0 +1,274 @@ +"""Small fused Triton kernels for the Qwen3.8 decode path. Inference only: every op +falls back to the eager formulation under autograd or off the GPU.""" + + + +import torch + +_kernels = None + + +def _get(): + global _kernels + if _kernels is not None: + return _kernels + import triton + import triton.language as tl + + @triton.jit + def rmsnorm_kernel(x_ptr, w_ptr, o_ptr, D, eps, ADD_ONE: tl.constexpr, BLOCK: tl.constexpr): + # one row per program: fp32 math, out = x * rsqrt(mean(x^2) + eps) * ((1 +) w) + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + mask = offs < D + x = tl.load(x_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + var = tl.sum(x * x, axis=0) / D + w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + if ADD_ONE: + w = w + 1.0 + y = x * (1.0 / tl.sqrt(var + eps)) * w + tl.store(o_ptr + row * D + offs, y.to(o_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def decode_attn_kernel( + q_ptr, k_ptr, v_ptr, qpos_ptr, o_ptr, + L, scale, min_key, + stride_qb, stride_qh, stride_qm, + stride_kb, stride_kh, stride_kl, + stride_ob, stride_oh, stride_om, + G: tl.constexpr, M: tl.constexpr, MPAD: tl.constexpr, ROWS: tl.constexpr, D: tl.constexpr, BLOCK_L: tl.constexpr, + ): + # one program per (batch, kv head): the G query heads of the group x M query + # tokens (rows padded to ROWS >= 16 for the MMA) attend the cache up to the + # largest query position, read from a device tensor so a CUDA graph never bakes a + # length in. Online softmax in fp32; scores in fp32 straight from the bf16 dot. + b = tl.program_id(0) + hk = tl.program_id(1) + rows = tl.arange(0, ROWS) + rmask = rows < G * M + g_idx = tl.where(rmask, rows // M, 0) + m_idx = tl.where(rmask, rows % M, 0) + offs_d = tl.arange(0, D) + qpos = tl.load(qpos_ptr + m_idx) + offs_mp = tl.arange(0, MPAD) + max_pos = tl.max(tl.load(qpos_ptr + offs_mp, mask=offs_mp < M, other=0), axis=0) + q = tl.load( + q_ptr + b * stride_qb + (hk * G + g_idx)[:, None] * stride_qh + m_idx[:, None] * stride_qm + offs_d[None, :], + mask=rmask[:, None], other=0.0, + ) # [ROWS, D] bf16 + m_i = tl.full((ROWS,), float("-inf"), tl.float32) + l_i = tl.zeros((ROWS,), tl.float32) + acc = tl.zeros((ROWS, D), tl.float32) + n_blocks = (max_pos + BLOCK_L) // BLOCK_L + for blk in range(0, n_blocks): + offs_l = blk * BLOCK_L + tl.arange(0, BLOCK_L) + kmask = offs_l < L + k = tl.load(k_ptr + b * stride_kb + hk * stride_kh + offs_l[:, None] * stride_kl + offs_d[None, :], + mask=kmask[:, None], other=0.0) # [BLOCK_L, D] + s = tl.dot(q, tl.trans(k), out_dtype=tl.float32) * scale # [ROWS, BLOCK_L] + allowed = (offs_l[None, :] <= qpos[:, None]) & (offs_l[None, :] >= min_key) & kmask[None, :] + s = tl.where(allowed, s, float("-inf")) + m_new = tl.maximum(m_i, tl.max(s, axis=1)) + alpha = tl.exp(m_i - m_new) + p = tl.exp(s - m_new[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + v = tl.load(v_ptr + b * stride_kb + hk * stride_kh + offs_l[:, None] * stride_kl + offs_d[None, :], + mask=kmask[:, None], other=0.0) + acc = acc * alpha[:, None] + tl.dot(p.to(tl.bfloat16), v, out_dtype=tl.float32) + m_i = m_new + out = acc / l_i[:, None] + tl.store(o_ptr + b * stride_ob + (hk * G + g_idx)[:, None] * stride_oh + m_idx[:, None] * stride_om + offs_d[None, :], + out.to(o_ptr.dtype.element_ty), mask=rmask[:, None]) + + @triton.jit + def add_rmsnorm_kernel(h_ptr, x_ptr, w_ptr, hout_ptr, y_ptr, D, eps, ADD_ONE: tl.constexpr, BLOCK: tl.constexpr): + # h_new = h + x (bf16, written back) and y = rmsnorm(h_new) in one pass + row = tl.program_id(0) + offs = tl.arange(0, BLOCK) + mask = offs < D + h = tl.load(h_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + x = tl.load(x_ptr + row * D + offs, mask=mask, other=0.0).to(tl.float32) + hn = (h + x).to(hout_ptr.dtype.element_ty) + tl.store(hout_ptr + row * D + offs, hn, mask=mask) + hf = hn.to(tl.float32) + var = tl.sum(hf * hf, axis=0) / D + w = tl.load(w_ptr + offs, mask=mask, other=0.0).to(tl.float32) + if ADD_ONE: + w = w + 1.0 + tl.store(y_ptr + row * D + offs, (hf * (1.0 / tl.sqrt(var + eps)) * w).to(y_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def swiglu_kernel(g_ptr, u_ptr, o_ptr, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + g = tl.load(g_ptr + offs, mask=mask, other=0.0).to(tl.float32) + u = tl.load(u_ptr + offs, mask=mask, other=0.0).to(tl.float32) + tl.store(o_ptr + offs, (g * tl.sigmoid(g) * u).to(o_ptr.dtype.element_ty), mask=mask) + + @triton.jit + def _conv_in(state_ptr, x_ptr, b, offs_c, mask, C, i: tl.constexpr, M: tl.constexpr, KS: tl.constexpr): + # input i of the window [state (KS-1 past inputs) | M new inputs] for channels offs_c + if i < KS - 1: + return tl.load(state_ptr + b * C * (KS - 1) + offs_c * (KS - 1) + i, mask=mask, other=0.0).to(tl.float32) + else: + return tl.load(x_ptr + (b * M + (i - (KS - 1))) * C + offs_c, mask=mask, other=0.0).to(tl.float32) + + @triton.jit + def conv_step_kernel(state_ptr, x_ptr, w_ptr, o_ptr, C, M: tl.constexpr, KS: tl.constexpr, WRITE_STATE: tl.constexpr, BLOCK_C: tl.constexpr): + # depthwise causal conv (KS taps) over [state | M new inputs] + silu, x and out token + # major [B, M, C] (as the projection produced them), state [B, C, KS-1]. + # Optionally slides the state window forward by M. + b = tl.program_id(0) + offs_c = tl.program_id(1) * BLOCK_C + tl.arange(0, BLOCK_C) + mask = offs_c < C + for m in tl.static_range(M): + acc = tl.zeros((BLOCK_C,), tl.float32) + for j in tl.static_range(KS): + w = tl.load(w_ptr + offs_c * KS + j, mask=mask, other=0.0).to(tl.float32) + acc += w * _conv_in(state_ptr, x_ptr, b, offs_c, mask, C, m + j, M, KS) + tl.store(o_ptr + (b * M + m) * C + offs_c, (acc * tl.sigmoid(acc)).to(o_ptr.dtype.element_ty), mask=mask) + if WRITE_STATE: + for j in tl.static_range(KS - 1): + v = _conv_in(state_ptr, x_ptr, b, offs_c, mask, C, M + j, M, KS) + tl.store(state_ptr + b * C * (KS - 1) + offs_c * (KS - 1) + j, v.to(state_ptr.dtype.element_ty), mask=mask) + + _kernels = (rmsnorm_kernel, decode_attn_kernel, add_rmsnorm_kernel, swiglu_kernel, conv_step_kernel) + return _kernels + + +@torch.library.custom_op("ostris::qwen38_rmsnorm", mutates_args=()) +def _rmsnorm_op(x: torch.Tensor, weight: torch.Tensor, eps: float, add_one: bool) -> torch.Tensor: + shape = x.shape + d = shape[-1] + x2 = x.reshape(-1, d) + if not x2.is_contiguous(): + x2 = x2.contiguous() + out = torch.empty_like(x2) + kernel = _get()[0] + block = 1 << (d - 1).bit_length() + kernel[(x2.shape[0],)](x2, weight, out, d, eps, ADD_ONE=add_one, BLOCK=block, num_warps=8 if block >= 2048 else 2) + return out.view(shape) + + +@_rmsnorm_op.register_fake +def _rmsnorm_fake(x, weight, eps, add_one): + return torch.empty_like(x) + + +def rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float, add_one: bool = True) -> torch.Tensor: + """Fused RMSNorm when no gradient is needed; eager fp32 math otherwise.""" + if x.is_cuda and not torch.is_grad_enabled() and x.shape[-1] <= 16384: + return _rmsnorm_op(x, weight, eps, add_one) + xf = x.float() + w = weight.float() + if add_one: + w = 1.0 + w + return (xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps) * w).type_as(x) + + +@torch.library.custom_op("ostris::qwen38_decode_attn", mutates_args=()) +def _decode_attn_op(q: torch.Tensor, k_cache: torch.Tensor, v_cache: torch.Tensor, qpos: torch.Tensor, scale: float, min_key: int) -> torch.Tensor: + """q [B, Hq, M, D] (M power of two <= 16), caches [B, Hk, L, D], qpos [M] int64. + Reads the cache only up to the largest query position.""" + b, hq, m, d = q.shape + hk, L = k_cache.shape[1], k_cache.shape[2] + g = hq // hk + q = q.contiguous() + out = torch.empty_like(q) + kernel = _get()[1] + rows = max(16, 1 << (g * m - 1).bit_length()) + kernel[(b, hk)]( + q, k_cache, v_cache, qpos, out, + L, scale, min_key, + q.stride(0), q.stride(1), q.stride(2), + k_cache.stride(0), k_cache.stride(1), k_cache.stride(2), + out.stride(0), out.stride(1), out.stride(2), + G=g, M=m, MPAD=1 << max(0, (m - 1).bit_length()), ROWS=rows, D=d, BLOCK_L=32, num_warps=4, num_stages=2, + ) + return out + + +@_decode_attn_op.register_fake +def _decode_attn_fake(q, k_cache, v_cache, qpos, scale, min_key): + return torch.empty_like(q) + + +# off by default: one program per kv head leaves most SMs idle, so past ~500 keys the +# cuBLAS path (full-cache masked bmm) is faster; a split-key version would fix that. +DECODE_ATTN_FUSED = False + + +def decode_attention_fused_ok(q: torch.Tensor, k_cache: torch.Tensor) -> bool: + b, hq, m, d = q.shape + hk = k_cache.shape[1] + return DECODE_ATTN_FUSED and q.is_cuda and m <= 16 and (hq // hk) * m <= 128 and d in (64, 128, 256) and not torch.is_grad_enabled() + + +@torch.library.custom_op("ostris::qwen38_add_rmsnorm", mutates_args=()) +def _add_rmsnorm_op(h: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, eps: float, add_one: bool) -> list[torch.Tensor]: + shape = h.shape + d = shape[-1] + h2 = h.reshape(-1, d).contiguous() + x2 = x.reshape(-1, d).contiguous() + hout = torch.empty_like(h2) + y = torch.empty_like(h2) + kernel = _get()[2] + block = 1 << (d - 1).bit_length() + kernel[(h2.shape[0],)](h2, x2, weight, hout, y, d, eps, ADD_ONE=add_one, BLOCK=block, num_warps=8 if block >= 2048 else 2) + return [hout.view(shape), y.view(shape)] + + +@_add_rmsnorm_op.register_fake +def _add_rmsnorm_fake(h, x, weight, eps, add_one): + return [torch.empty_like(h), torch.empty_like(h)] + + +def add_rmsnorm(h: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, eps: float, add_one: bool = True): + """(h + x, rmsnorm(h + x)): fused on the no-grad GPU path.""" + if h.is_cuda and not torch.is_grad_enabled() and h.shape[-1] <= 16384 and h.dtype == x.dtype: + hn, y = _add_rmsnorm_op(h, x, weight, eps, add_one) + return hn, y + hn = h + x + return hn, rmsnorm(hn, weight, eps, add_one) + + +@torch.library.custom_op("ostris::qwen38_swiglu", mutates_args=()) +def _swiglu_op(g: torch.Tensor, u: torch.Tensor) -> torch.Tensor: + g2 = g.contiguous() + u2 = u.contiguous() + out = torch.empty_like(g2) + n = g2.numel() + kernel = _get()[3] + kernel[(-(-n // 4096),)](g2, u2, out, n, BLOCK=4096, num_warps=4) + return out + + +@_swiglu_op.register_fake +def _swiglu_fake(g, u): + return torch.empty_like(g) + + +def swiglu(g: torch.Tensor, u: torch.Tensor) -> torch.Tensor: + if g.is_cuda and not torch.is_grad_enabled() and g.dtype == u.dtype: + return _swiglu_op(g, u) + return torch.nn.functional.silu(g) * u + + +@torch.library.custom_op("ostris::qwen38_conv_step", mutates_args={"state"}) +def _conv_step_op(state: torch.Tensor, x: torch.Tensor, weight: torch.Tensor, write_state: bool) -> torch.Tensor: + """Depthwise causal conv + silu of M new tokens x [B, M, C] against state [B, C, KS-1] + (conv weight [C, 1, KS]); returns [B, M, C]. write_state slides the window forward.""" + b, m, c = x.shape + ks = weight.shape[-1] + x2 = x.contiguous() + out = torch.empty_like(x2) + kernel = _get()[4] + block_c = 1024 + kernel[(b, -(-c // block_c))](state, x2, weight.reshape(c, ks).contiguous(), out, c, M=m, KS=ks, WRITE_STATE=write_state, BLOCK_C=block_c, num_warps=4) + return out + + +@_conv_step_op.register_fake +def _conv_step_fake(state, x, weight, write_state): + return torch.empty_like(x) diff --git a/extensions_built_in/llm_models/src/qwen38_load.py b/extensions_built_in/llm_models/src/qwen38_load.py new file mode 100644 index 0000000..b876060 --- /dev/null +++ b/extensions_built_in/llm_models/src/qwen38_load.py @@ -0,0 +1,129 @@ +"""Build a Qwen38ForConditionalGeneration from a ComfyUI single file (quantized layers +attached to the toolkit backends as-is) or an HF layout (bf16 shards).""" + +import json +import os +from typing import Callable, Optional, Tuple + +import torch +from safetensors import safe_open + +from toolkit.print import print_acc + +from .qwen38_model import Qwen38ForConditionalGeneration + + +def _list_source_files(source: str): + """(files, is_single_file) for a .safetensors path, a local HF dir or a hub repo id.""" + if os.path.isfile(source): + return [source], True + if os.path.isdir(source): + d = source + else: + import huggingface_hub + + d = huggingface_hub.snapshot_download(source, allow_patterns=["*.safetensors", "*.json"]) + index = os.path.join(d, "model.safetensors.index.json") + if os.path.isfile(index): + files = sorted(set(json.load(open(index))["weight_map"].values())) + return [os.path.join(d, f) for f in files], False + single = os.path.join(d, "model.safetensors") + if os.path.isfile(single): + return [single], False + raise FileNotFoundError(f"No safetensors weights found under {d}") + + +def _read_keys(files): + keys = set() + for f in files: + with safe_open(f, "pt") as sf: + keys.update(sf.keys()) + return keys + + +def _load_into_device(files, device) -> dict: + sd = {} + dev = str(device) + for f in files: + with safe_open(f, "pt", device=dev) as sf: + for k in sf.keys(): + sd[k] = sf.get_tensor(k) + return sd + + +@torch.no_grad() +def load_qwen38( + source: str, + config, + dtype: torch.dtype, + device, + status_fn: Optional[Callable[[str], None]] = None, +) -> Tuple[Qwen38ForConditionalGeneration, bool]: + """Returns (model on `device`, is_prequantized). Quantized ComfyUI layers keep + their stored format (convrotw4a8 / convrot8 OstrisLinear, int8 embeddings).""" + from accelerate import init_empty_weights + + from toolkit.util.comfy_quant_import import import_comfy_quantized_layers + + files, _ = _list_source_files(source) + keys = _read_keys(files) + with_mtp = any(k.startswith("mtp.") for k in keys) + config.vision_config._attn_implementation = "sdpa" + with init_empty_weights(include_buffers=False): + model = Qwen38ForConditionalGeneration(config, with_mtp=with_mtp) + model.eval() + + if status_fn is not None: + status_fn(f"Reading {len(files)} weight file(s) onto {device}") + sd = _load_into_device(files, device) + prequantized = any(k.endswith(".comfy_quant") for k in sd) + if prequantized: + sd, n = import_comfy_quantized_layers(model, sd, orig_dtype=dtype, rotated_embedding_rows=True) + print_acc(f"Attached {n} pre-quantized ComfyUI layers") + sd = {k: (v.to(dtype) if v.is_floating_point() else v) for k, v in sd.items()} + result = model.load_state_dict(sd, assign=True, strict=False) + allowed = set() + for name, module in model.named_modules(): + if hasattr(module, "ostris_quantizer"): + allowed.add(f"{name}.weight") + allowed.add(f"{name}.bias") + missing = [k for k in result.missing_keys if k not in allowed] + if missing or result.unexpected_keys: + raise RuntimeError(f"Qwen3.8 checkpoint mismatch. missing: {missing[:8]} unexpected: {result.unexpected_keys[:8]}") + leftover = [n for n, p in model.named_parameters() if p.device.type == "meta"] + if leftover: + raise RuntimeError(f"Params never loaded: {leftover[:8]}") + # buffers built in __init__ (rope tables) still live on the cpu + for module in model.modules(): + if hasattr(module, "ostris_quantizer"): + continue + for pname, prm in list(module.named_parameters(recurse=False)): + if prm.is_floating_point() and prm.dtype != dtype: + prm.data = prm.data.to(dtype) + model.to(device) + return model, prequantized + + +@torch.no_grad() +def dequantize_vision_tower(model: Qwen38ForConditionalGeneration, dtype: torch.dtype) -> int: + """Replace the vision tower's quantized linears with plain Linears holding the + dequantized weights. The stored int8 weights are kept exactly; only the per-token + int8 activation quantization goes away, which measured a 6% drift of the image + embeddings against the bf16 reference. Runs once per image, so the cost is memory + only (~0.4 GB for the 27B tower). Returns the number of layers converted.""" + from toolkit.util.ostris_quant import OstrisLinear + + visual = model.model.visual + n = 0 + for name, mod in list(visual.named_modules()): + if not isinstance(mod, OstrisLinear): + continue + lin = torch.nn.Linear(mod.in_features, mod.out_features, bias=mod.bias is not None, device="meta", dtype=dtype) + lin.weight = torch.nn.Parameter(mod.ostris_quantizer.dequantize(mod).to(dtype), requires_grad=False) + if mod.bias is not None: + lin.bias = torch.nn.Parameter(mod.bias.detach().to(dtype), requires_grad=False) + parent_name, _, attr = name.rpartition(".") + parent = visual.get_submodule(parent_name) if parent_name else visual + setattr(parent, attr, lin) + n += 1 + return n diff --git a/extensions_built_in/llm_models/src/qwen38_model.py b/extensions_built_in/llm_models/src/qwen38_model.py new file mode 100644 index 0000000..16da7dc --- /dev/null +++ b/extensions_built_in/llm_models/src/qwen38_model.py @@ -0,0 +1,808 @@ +"""Qwen3.8 (Qwen3.5-family hybrid VLM) built for throughput. + +Module tree and parameter names mirror transformers' Qwen3_5ForConditionalGeneration +(model.language_model.layers.N.{linear_attn|self_attn,mlp,...}, model.visual.*, +lm_head, mtp.*) so the ComfyUI / HF checkpoints load 1:1 and LoRA keys stay standard. +The vision tower is transformers' own Qwen3_5VisionModel (runs once per image); the +text stack, caches, decode step and speculative loop are this file's. + +Text stack: 3 of every 4 layers are Gated DeltaNet (linear attention, fla kernels: +chunked for prefill / training, fused recurrent for decode), the 4th is gated full +attention (GQA, partial interleaved MRoPE). Linears are whatever the loader attached +(W4A8 codebook / int8 convrot OstrisLinear, or plain nn.Linear). + +Decode: static KV + recurrent-state caches with a device-side position, so one +decode step is a fixed graph that CUDA-graph replays. With the MTP draft head the +step is draft D tokens -> verify D+1 in one pass -> accept -> commit states -> +refresh the MTP cache from the verified hidden states. +""" + +from dataclasses import dataclass +from typing import Dict, List, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.checkpoint +from fla.modules import FusedRMSNormGated +from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule +from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel + +from toolkit.util.convrot_w4a8_quant import apply_forward_hooks, module_hooks_only, w4a8_grouped_linear + +from .qwen38_kernels import _conv_step_op, _decode_attn_op, add_rmsnorm, decode_attention_fused_ok, rmsnorm, swiglu + +# ------------------------------------------------------------------ norms / rope + + +class RMSNorm(nn.Module): + """Qwen3.5 RMSNorm: fp32 math, (1 + weight) scaling.""" + + def __init__(self, dim: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.zeros(dim)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return rmsnorm(x, self.weight, self.eps, add_one=True) + + +def rotate_half(x): + x1, x2 = x[..., : x.shape[-1] // 2], x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_partial_rope(q, k, cos, sin): + """q [B,H,T,D], k [B,Hk,T,D]; cos/sin [B,T,R] rotate the first R dims.""" + cos = cos.unsqueeze(1) + sin = sin.unsqueeze(1) + r = cos.shape[-1] + q_rot, q_pass = q[..., :r], q[..., r:] + k_rot, k_pass = k[..., :r], k[..., r:] + q_rot = q_rot * cos + rotate_half(q_rot) * sin + k_rot = k_rot * cos + rotate_half(k_rot) * sin + return torch.cat([q_rot, q_pass], dim=-1), torch.cat([k_rot, k_pass], dim=-1) + + +class InterleavedMRope(nn.Module): + """Qwen3.5 interleaved MRoPE for the partial rotary dims: frequency i takes its + position from axis i % 3 (T, H, W). Text tokens have equal positions on all axes, + so decode uses a 1-D table.""" + + def __init__(self, rotary_dim: int, theta: float, mrope_section): + super().__init__() + half = rotary_dim // 2 + inv_freq = 1.0 / (theta ** (torch.arange(0, rotary_dim, 2, dtype=torch.int64).float() / rotary_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + # axis per frequency: [T,H,W] interleaved as in apply_interleaved_mrope + axis = torch.zeros(half, dtype=torch.long) + for dim, offset in enumerate((1, 2), start=1): + axis[offset : mrope_section[dim] * 3 : 3] = dim + self.register_buffer("axis", axis, persistent=False) + + def forward(self, position_ids: torch.Tensor, dtype: torch.dtype): + """position_ids [3, B, T] -> cos, sin [B, T, rotary_dim].""" + pos = position_ids.float() + freqs = pos[..., None] * self.inv_freq.float()[None, None, None, :] # [3,B,T,half] + freqs = torch.gather(freqs, 0, self.axis.view(1, 1, 1, -1).expand(1, *freqs.shape[1:]))[0] + emb = torch.cat([freqs, freqs], dim=-1) + return emb.cos().to(dtype), emb.sin().to(dtype) + + def table(self, max_len: int, dtype: torch.dtype): + """1-D cos/sin table [max_len, rotary_dim] for text-only positions.""" + pos = torch.arange(max_len, device=self.inv_freq.device).view(1, 1, -1).expand(3, 1, -1) + cos, sin = self.forward(pos, dtype) + return cos[0], sin[0] + + +# ------------------------------------------------------------------ caches + + +@dataclass +class DecodeCtx: + """Per-forward decode context: where the M new tokens live in the caches.""" + + pos_idx: torch.Tensor # [M] int64 device: cache slots of the new tokens + qpos: torch.Tensor # [M] int64 device: absolute position of each query (for the causal mask) + cos: torch.Tensor # [B, M, R] + sin: torch.Tensor + commit: bool = True # linear layers: write the new recurrent/conv state now + min_key: int = 0 # keys below this slot are never attended (MTP slot 0 is empty) + + +class LayerCache: + """Static per-layer state. Attention layers: k/v [B, Hk, L, D]. Linear layers: + conv [B, C, ks-1] (last conv inputs) and rec [B, Hv, Dk, Dv] fp32.""" + + def __init__(self): + self.k = self.v = None + self.conv = self.rec = None + # speculative verify stash (uncommitted step inputs) + self.spec = None + + +class Qwen38Cache: + def __init__(self, cfg, batch: int, max_len: int, device, dtype, layer_types=None, with_linear=True): + self.max_len = max_len + self.batch = batch + layer_types = cfg.layer_types if layer_types is None else layer_types + self.layers: List[LayerCache] = [] + for lt in layer_types: + lc = LayerCache() + if lt == "full_attention": + lc.k = torch.zeros(batch, cfg.num_key_value_heads, max_len, cfg.head_dim, device=device, dtype=dtype) + lc.v = torch.zeros_like(lc.k) + elif with_linear: + conv_dim = cfg.linear_num_key_heads * cfg.linear_key_head_dim * 2 + cfg.linear_num_value_heads * cfg.linear_value_head_dim + lc.conv = torch.zeros(batch, conv_dim, cfg.linear_conv_kernel_dim - 1, device=device, dtype=dtype) + lc.rec = torch.zeros( + batch, cfg.linear_num_value_heads, cfg.linear_key_head_dim, cfg.linear_value_head_dim, + device=device, dtype=torch.float32, + ) + self.layers.append(lc) + # committed length, device-side so the decode graph never syncs + self.pos = torch.zeros(1, dtype=torch.long, device=device) + + def reset(self): + for lc in self.layers: + for t in (lc.k, lc.v, lc.conv, lc.rec): + if t is not None: + t.zero_() + self.pos.zero_() + + +# ------------------------------------------------------------------ mixers + + +class GatedDeltaNet(nn.Module): + def __init__(self, cfg, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.hidden_size = cfg.hidden_size + self.num_v_heads = cfg.linear_num_value_heads + self.num_k_heads = cfg.linear_num_key_heads + self.head_k_dim = cfg.linear_key_head_dim + self.head_v_dim = cfg.linear_value_head_dim + self.key_dim = self.head_k_dim * self.num_k_heads + self.value_dim = self.head_v_dim * self.num_v_heads + self.conv_kernel_size = cfg.linear_conv_kernel_dim + self.conv_dim = self.key_dim * 2 + self.value_dim + self.conv1d = nn.Conv1d( + self.conv_dim, self.conv_dim, bias=False, kernel_size=self.conv_kernel_size, + groups=self.conv_dim, padding=self.conv_kernel_size - 1, + ) + self.dt_bias = nn.Parameter(torch.ones(self.num_v_heads)) + self.A_log = nn.Parameter(torch.zeros(self.num_v_heads)) + self.norm = FusedRMSNormGated(self.head_v_dim, eps=cfg.rms_norm_eps, activation="silu") + self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False) + self.in_proj_qkv = nn.Linear(self.hidden_size, self.conv_dim, bias=False) + self.in_proj_z = nn.Linear(self.hidden_size, self.value_dim, bias=False) + self.in_proj_b = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + self.in_proj_a = nn.Linear(self.hidden_size, self.num_v_heads, bias=False) + + def _heads(self, mixed: torch.Tensor): + b, t, _ = mixed.shape + q, k, v = torch.split(mixed, [self.key_dim, self.key_dim, self.value_dim], dim=-1) + q = q.reshape(b, t, self.num_k_heads, self.head_k_dim) + k = k.reshape(b, t, self.num_k_heads, self.head_k_dim) + v = v.reshape(b, t, self.num_v_heads, self.head_v_dim) + # fla's kernels take fewer key heads than value heads natively (same result as + # the repeat_interleave transformers does, measured bit-identical) + return q, k, v + + def _gates(self, a: torch.Tensor, b: torch.Tensor): + beta = b.sigmoid() + g = -self.A_log.float().exp() * F.softplus(a.float() + self.dt_bias.float()) + return beta, g + + def _raw_ab(self, x: torch.Tensor): + """Pre-activation gate inputs (a, b) [B, T, Hv]: one GEMM over the stacked weights + when the two tiny projections are not wrapped (e.g. by an active LoRA).""" + pa, pb = self.in_proj_a, self.in_proj_b + if not (module_hooks_only(pa) and module_hooks_only(pb)): + return pa(x), pb(x) + w = getattr(self, "_ab_weight", None) + if w is None or w.device != x.device or w.dtype != x.dtype: + w = torch.cat([pa.weight, pb.weight], dim=0).to(x.dtype) + self._ab_weight = w + ab = F.linear(x, w) + a, b = ab[..., : self.num_v_heads], ab[..., self.num_v_heads :] + if pa._forward_hooks: + a = apply_forward_hooks(pa, x, a) + if pb._forward_hooks: + b = apply_forward_hooks(pb, x, b) + return a, b + + def _finish(self, core: torch.Tensor, z: torch.Tensor): + b, t = core.shape[:2] + out = self.norm(core.reshape(-1, self.head_v_dim), z.reshape(-1, self.head_v_dim)) + return self.out_proj(out.reshape(b, t, -1)) + + def forward(self, x: torch.Tensor, cache: Optional[LayerCache] = None, ctx: Optional[DecodeCtx] = None): + if ctx is not None: + return self.forward_decode(x, cache, ctx) + return self.forward_prefill(x, cache) + + def forward_prefill(self, x: torch.Tensor, cache: Optional[LayerCache]): + b, t, _ = x.shape + mixed, z = w4a8_grouped_linear((self.in_proj_qkv, self.in_proj_z), x) + beta, g = self._gates(self.in_proj_a(x), self.in_proj_b(x)) + xt = mixed.transpose(1, 2) # [B, C, T] + if cache is not None: + ks = self.conv_kernel_size - 1 + cache.conv.copy_(F.pad(xt, (max(ks - t, 0), 0))[:, :, -ks:]) + conv = F.silu(F.conv1d(xt, self.conv1d.weight, None, padding=self.conv_kernel_size - 1, groups=self.conv_dim)[:, :, :t]) + q, k, v = self._heads(conv.transpose(1, 2)) + core, final = chunk_gated_delta_rule( + q, k, v, g=g, beta=beta, initial_state=None, output_final_state=cache is not None, + use_qk_l2norm_in_kernel=True, + ) + if cache is not None: + cache.rec.copy_(final) + return self._finish(core, z) + + def forward_decode(self, x: torch.Tensor, cache: LayerCache, ctx: DecodeCtx): + b, m, _ = x.shape + mixed, z = w4a8_grouped_linear((self.in_proj_qkv, self.in_proj_z), x) + a_raw, b_raw = self._raw_ab(x) + # one kernel: causal depthwise conv over [state | new tokens] + silu, token major + conv = _conv_step_op(cache.conv, mixed, self.conv1d.weight, ctx.commit) # [B, M, C] + q, k, v = self._heads(conv) + # the decay / beta activations run inside the recurrent kernel + core, final = fused_recurrent_gated_delta_rule( + q, k, v, g=a_raw.float(), beta=b_raw, initial_state=cache.rec, output_final_state=True, + use_qk_l2norm_in_kernel=True, use_gate_in_kernel=True, A_log=self.A_log.float(), + dt_bias=self.dt_bias.float(), use_beta_sigmoid_in_kernel=True, + ) + if ctx.commit: + cache.rec.copy_(final) + else: + cache.spec = (q, k, v, a_raw, b_raw, mixed) + return self._finish(core, z) + + def commit_speculative(self, cache: LayerCache, keep: torch.Tensor, accepted: torch.Tensor): + """After verification: advance the state through the first accepted+1 of the + stashed tokens. Rejected tokens get g=0, beta=0 (no decay, no update), so one + more recurrent pass from the committed state lands exactly on S_{accepted+1}.""" + q, k, v, a_raw, b_raw, mixed = cache.spec + full = torch.cat([cache.conv, mixed.transpose(1, 2)], dim=-1) # [B, C, ks-1+M] + beta, g = self._gates(a_raw, b_raw) + keep_h = keep.view(1, -1, 1).to(g.dtype) + _, final = fused_recurrent_gated_delta_rule( + q, k, v, g=g * keep_h, beta=beta * keep_h.to(beta.dtype), initial_state=cache.rec, + output_final_state=True, use_qk_l2norm_in_kernel=True, + ) + cache.rec.copy_(final) + ks = self.conv_kernel_size - 1 + idx = accepted.view(1) + 1 + torch.arange(ks, device=full.device) + cache.conv.copy_(full.index_select(-1, idx)) + cache.spec = None + + +def decode_attention(q, k_cache, v_cache, qpos, scale, min_key: int = 0): + """Single-step GQA attention of M queries against the whole static cache. + q [B,Hq,M,D]; caches [B,Hk,L,D]; qpos [M] absolute positions. Scores in fp32 + straight from the bf16 GEMM (sink logits overflow a bf16 score tensor).""" + if decode_attention_fused_ok(q, k_cache): + return _decode_attn_op(q, k_cache, v_cache, qpos, scale, min_key) + b, hq, m, d = q.shape + hk, L = k_cache.shape[1], k_cache.shape[2] + g = hq // hk + qg = q.reshape(b, hk, g, m, d).reshape(b * hk, g * m, d) + kT = k_cache.reshape(b * hk, L, d).transpose(-1, -2) + scores = torch.bmm(qg, kT, out_dtype=torch.float32).view(b * hk, g, m, L) * scale + keys = torch.arange(L, device=q.device) + allowed = keys[None, :] <= qpos[:, None] # [M, L] + if min_key > 0: + allowed = allowed & (keys[None, :] >= min_key) + scores = scores.masked_fill(~allowed[None, None], float("-inf")) + probs = torch.softmax(scores, dim=-1).to(q.dtype).view(b * hk, g * m, L) + out = torch.bmm(probs, v_cache.reshape(b * hk, L, d)) # [B*hk, g*m, d] + return out.view(b, hk, g, m, d).reshape(b, hq, m, d) + + +class GatedAttention(nn.Module): + def __init__(self, cfg, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + self.head_dim = cfg.head_dim + self.num_heads = cfg.num_attention_heads + self.num_kv_heads = cfg.num_key_value_heads + self.scaling = self.head_dim**-0.5 + self.q_proj = nn.Linear(cfg.hidden_size, self.num_heads * self.head_dim * 2, bias=cfg.attention_bias) + self.k_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=cfg.attention_bias) + self.v_proj = nn.Linear(cfg.hidden_size, self.num_kv_heads * self.head_dim, bias=cfg.attention_bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, cfg.hidden_size, bias=cfg.attention_bias) + self.q_norm = RMSNorm(self.head_dim, eps=cfg.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=cfg.rms_norm_eps) + + def _qkv(self, x, cos, sin): + b, t, _ = x.shape + qg, k, v = w4a8_grouped_linear((self.q_proj, self.k_proj, self.v_proj), x) + qg = qg.view(b, t, self.num_heads, 2 * self.head_dim) + q, gate = qg[..., : self.head_dim], qg[..., self.head_dim :] + gate = gate.reshape(b, t, -1) + q = self.q_norm(q).transpose(1, 2) + k = self.k_norm(k.view(b, t, self.num_kv_heads, self.head_dim)).transpose(1, 2) + v = v.view(b, t, self.num_kv_heads, self.head_dim).transpose(1, 2) + q, k = apply_partial_rope(q, k, cos, sin) + return q, k, v, gate + + def forward(self, x, cos, sin, cache: Optional[LayerCache] = None, ctx: Optional[DecodeCtx] = None, start: int = 0): + b, t, _ = x.shape + q, k, v, gate = self._qkv(x, cos, sin) + if ctx is None: + if cache is not None: + cache.k[:, :, start : start + t].copy_(k) + cache.v[:, :, start : start + t].copy_(v) + o = F.scaled_dot_product_attention(q, k, v, is_causal=True, enable_gqa=True, scale=self.scaling) + else: + cache.k.index_copy_(2, ctx.pos_idx, k) + cache.v.index_copy_(2, ctx.pos_idx, v) + o = decode_attention(q, cache.k, cache.v, ctx.qpos, self.scaling, ctx.min_key) + o = o.transpose(1, 2).reshape(b, t, -1) * torch.sigmoid(gate) + return self.o_proj(o) + + +class MLP(nn.Module): + def __init__(self, cfg): + super().__init__() + self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) + self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False) + self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False) + + def forward(self, x): + gate, up = w4a8_grouped_linear((self.gate_proj, self.up_proj), x) + return self.down_proj(swiglu(gate, up)) + + +class DecoderLayer(nn.Module): + def __init__(self, cfg, layer_idx: int, layer_type: Optional[str] = None): + super().__init__() + self.layer_type = layer_type or cfg.layer_types[layer_idx] + if self.layer_type == "linear_attention": + self.linear_attn = GatedDeltaNet(cfg, layer_idx) + else: + self.self_attn = GatedAttention(cfg, layer_idx) + self.mlp = MLP(cfg) + self.input_layernorm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + self.post_attention_layernorm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + + def forward(self, h, cos, sin, cache=None, ctx=None, start: int = 0, pending=None): + """pending: the previous layer's MLP output, not yet added to the residual stream; + its add is fused with this layer's input norm. Returns (h, mlp_out) the same way.""" + if pending is not None: + h, x = add_rmsnorm(h, pending, self.input_layernorm.weight, self.input_layernorm.eps) + else: + x = self.input_layernorm(h) + if self.layer_type == "linear_attention": + x = self.linear_attn(x, cache, ctx) + else: + x = self.self_attn(x, cos, sin, cache, ctx, start) + h, x = add_rmsnorm(h, x, self.post_attention_layernorm.weight, self.post_attention_layernorm.eps) + return h, self.mlp(x) + + +# ------------------------------------------------------------------ text model / mtp / top + + +class Qwen38TextModel(nn.Module): + def __init__(self, cfg): + super().__init__() + self.config = cfg + self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size, cfg.pad_token_id) + self.layers = nn.ModuleList([DecoderLayer(cfg, i) for i in range(cfg.num_hidden_layers)]) + self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + rp = cfg.rope_parameters + self.rotary = InterleavedMRope( + int(cfg.head_dim * rp.get("partial_rotary_factor", 1.0)), rp["rope_theta"], rp.get("mrope_section", [11, 11, 10]) + ) + self.gradient_checkpointing = False + + def forward(self, inputs_embeds, cos, sin, cache: Optional[Qwen38Cache] = None, ctx: Optional[DecodeCtx] = None, start: int = 0): + h = inputs_embeds + pending = None + for i, layer in enumerate(self.layers): + lc = cache.layers[i] if cache is not None else None + if self.gradient_checkpointing and torch.is_grad_enabled(): + h, pending = torch.utils.checkpoint.checkpoint(layer, h, cos, sin, lc, ctx, start, pending, use_reentrant=False) + else: + h, pending = layer(h, cos, sin, lc, ctx, start, pending) + _, y = add_rmsnorm(h, pending, self.norm.weight, self.norm.eps) + return y + + def commit_speculative(self, cache: Qwen38Cache, keep: torch.Tensor, accepted: torch.Tensor): + for i, layer in enumerate(self.layers): + if layer.layer_type == "linear_attention": + layer.linear_attn.commit_speculative(cache.layers[i], keep, accepted) + + +class Qwen38MTP(nn.Module): + """Multi-token-prediction draft head: fc(cat[norm(embed(token_p)), norm(hidden_{p-1})]) + -> one full-attention decoder layer with its own KV cache -> norm -> lm_head.""" + + def __init__(self, cfg): + super().__init__() + self.fc = nn.Linear(cfg.hidden_size * 2, cfg.hidden_size, bias=False) + self.pre_fc_norm_embedding = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + self.pre_fc_norm_hidden = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + self.layers = nn.ModuleList([DecoderLayer(cfg, 0, layer_type="full_attention")]) + self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps) + + def forward(self, embeds, hidden, cos, sin, cache: LayerCache, ctx: Optional[DecodeCtx] = None, start: int = 0): + x = self.fc(torch.cat([self.pre_fc_norm_embedding(embeds), self.pre_fc_norm_hidden(hidden)], dim=-1)) + h, mlp = self.layers[0](x, cos, sin, cache, ctx, start) + h, y = add_rmsnorm(h, mlp, self.norm.weight, self.norm.eps) + return y, h + + +class Qwen38Inner(nn.Module): + def __init__(self, config): + super().__init__() + self.visual = Qwen3_5VisionModel._from_config(config.vision_config) + self.language_model = Qwen38TextModel(config.text_config) + + +class Qwen38ForConditionalGeneration(nn.Module): + def __init__(self, config, with_mtp: bool = True): + super().__init__() + self.config = config + self.model = Qwen38Inner(config) + tc = config.text_config + self.lm_head = nn.Linear(tc.hidden_size, tc.vocab_size, bias=False) + self.mtp = Qwen38MTP(tc) if with_mtp else None + self.generator: Optional["Qwen38Generator"] = None + + @property + def text_config(self): + return self.config.text_config + + @property + def device(self): + return self.model.language_model.norm.weight.device + + def get_input_embeddings(self): + return self.model.language_model.embed_tokens + + # -------------------------------------------------- vision + + @torch.no_grad() + def get_image_features(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor: + """Merged vision embeddings [N_merged, hidden] for one or more images / videos.""" + visual = self.model.visual + out = visual(pixel_values.to(visual.patch_embed.proj.weight.dtype), grid_thw=grid_thw, return_dict=True) + return out.pooler_output + + # -------------------------------------------------- positions + + def get_rope_index(self, input_ids: torch.Tensor, mm_token_type_ids: torch.Tensor, image_grid_thw=None, video_grid_thw=None): + """Port of transformers' Qwen3_5 get_rope_index for one unpadded sequence. + input_ids / mm_token_type_ids [T]; returns position_ids [3, 1, T] and the rope + delta (first text position after the sequence minus its length).""" + merge = self.config.vision_config.spatial_merge_size + device = input_ids.device + if video_grid_thw is not None: + video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0) + video_grid_thw[:, 0] = 1 + grids = {1: iter(image_grid_thw) if image_grid_thw is not None else None, 2: iter(video_grid_thw) if video_grid_thw is not None else None} + types = mm_token_type_ids.tolist() + groups = [] + i = 0 + while i < len(types): + j = i + while j < len(types) and types[j] == types[i]: + j += 1 + groups.append((types[i], i, j)) + i = j + cur = 0 + parts = [] + for kind, s, e in groups: + if kind == 0: + n = e - s + parts.append(torch.arange(n, device=device).view(1, -1).expand(3, -1) + cur) + cur += n + else: + thw = next(grids[kind]) + t, h, w = int(thw[0]), int(thw[1]) // merge, int(thw[2]) // merge + n = t * h * w + pw = torch.arange(cur, cur + w, device=device).repeat(h * t) + ph = torch.arange(cur, cur + h, device=device).repeat_interleave(w * t) + pt = torch.full((n,), cur, device=device, dtype=torch.long) + parts.append(torch.stack([pt, ph, pw], dim=0)) + cur += max(h, w) + pos = torch.cat(parts, dim=1) + delta = int(pos.max().item()) + 1 - input_ids.shape[0] + return pos.view(3, 1, -1), delta + + # -------------------------------------------------- forward (prefill / training) + + def forward_hidden(self, inputs_embeds, position_ids, cache: Optional[Qwen38Cache] = None): + """Post-norm hidden states [B, T, H] for a full sequence (prefill or training).""" + lm = self.model.language_model + cos, sin = lm.rotary(position_ids, inputs_embeds.dtype) + return lm(inputs_embeds, cos, sin, cache=cache) + + def enable_gradient_checkpointing(self): + self.model.language_model.gradient_checkpointing = True + + +# ------------------------------------------------------------------ generation + + +class Qwen38Generator: + """Greedy generation with static caches and CUDA-graph decode steps; optional MTP + speculative decoding (draft depth D). One instance per model; caches and graphs + are reused across prompts and rebuilt when the length budget grows.""" + + def __init__(self, model: Qwen38ForConditionalGeneration, use_cuda_graph: bool = True): + self.model = model + self.cfg = model.text_config + self.use_cuda_graph = use_cuda_graph + self.cache: Optional[Qwen38Cache] = None + self.mtp_cache: Optional[Qwen38Cache] = None + self.max_len = 0 + self.graphs: Dict[tuple, Tuple[torch.cuda.CUDAGraph, dict]] = {} + self.dtype = model.model.language_model.norm.weight.dtype + self.device = model.device + self.rope_table = None + self.last_stats = {} + + # ---------------- buffers + + def _ensure(self, max_len: int, batch: int = 1): + if self.cache is not None and self.max_len >= max_len and self.cache.batch == batch: + return + self.graphs = {} + self.max_len = max_len + self.cache = Qwen38Cache(self.cfg, batch, max_len, self.device, self.dtype) + if self.model.mtp is not None: + self.mtp_cache = Qwen38Cache(self.cfg, batch, max_len, self.device, self.dtype, layer_types=["full_attention"], with_linear=False) + cos, sin = self.model.model.language_model.rotary.table(max_len + 64, self.dtype) + self.rope_table = (cos, sin) + self.rope_delta = torch.zeros(1, dtype=torch.long, device=self.device) + hs = self.cfg.hidden_size + self.nt_buf = torch.zeros(batch, 1, dtype=torch.long, device=self.device) + self.h_buf = torch.zeros(batch, 1, hs, dtype=self.dtype, device=self.device) + self.out_buf = torch.zeros(1 + 16, dtype=torch.long, device=self.device) + + def _cos_sin(self, positions: torch.Tensor): + cos, sin = self.rope_table + idx = (positions + self.rope_delta).clamp_(min=0) + return cos.index_select(0, idx)[None], sin.index_select(0, idx)[None] + + def _embed(self, ids: torch.Tensor): + return self.model.get_input_embeddings()(ids).to(self.dtype) + + def _lm_head_argmax(self, hidden: torch.Tensor) -> torch.Tensor: + return self.model.lm_head(hidden).argmax(dim=-1) + + # ---------------- step bodies (graph-capturable: no host syncs) + + def _plain_step(self): + cache = self.cache + pos = cache.pos + ctx = DecodeCtx(pos_idx=pos, qpos=pos, cos=None, sin=None, commit=True) + ctx.cos, ctx.sin = self._cos_sin(pos) + hidden = self.model.model.language_model(self._embed(self.nt_buf), ctx.cos, ctx.sin, cache=cache, ctx=ctx) + nt = self._lm_head_argmax(hidden[:, -1:]) + self.out_buf[0].fill_(0) + self.out_buf[1].copy_(nt.view(-1)[0]) + self.nt_buf.copy_(nt) + self.h_buf.copy_(hidden[:, -1:]) + pos.add_(1) + + def _mtp_step(self, depth: int): + model, cache, mtp_cache = self.model, self.cache, self.mtp_cache + lm = model.model.language_model + pos = cache.pos + ar = torch.arange(depth + 1, device=self.device) + tok_in, hid_in = self.nt_buf, self.h_buf + drafts = [] + for k in range(depth): + p = pos + k + ctx = DecodeCtx(pos_idx=p, qpos=p, cos=None, sin=None, min_key=1) + ctx.cos, ctx.sin = self._cos_sin(p) + normed, pre = model.mtp(self._embed(tok_in), hid_in, ctx.cos, ctx.sin, mtp_cache.layers[0], ctx) + tok_in = self._lm_head_argmax(normed) + hid_in = pre + drafts.append(tok_in) + drafts = torch.cat(drafts, dim=1) # [B, D] + tokens = torch.cat([self.nt_buf, drafts], dim=1) # [B, D+1] + positions = pos + ar + ctx = DecodeCtx(pos_idx=positions, qpos=positions, cos=None, sin=None, commit=False) + ctx.cos, ctx.sin = self._cos_sin(positions) + hidden = lm(self._embed(tokens), ctx.cos, ctx.sin, cache=cache, ctx=ctx) # [B, D+1, H] + pred = self._lm_head_argmax(hidden) # [B, D+1] + match = (pred[:, :depth] == drafts).long() + accepted = match.cumprod(dim=1).sum(dim=1)[0] # scalar in [0, D] + keep = (ar <= accepted).long() + lm.commit_speculative(cache, keep, accepted) + # refresh the draft cache from the verified hidden states: pair (token_p, hidden_{p-1}) + prev_hidden = torch.cat([self.h_buf, hidden[:, :depth]], dim=1) + rctx = DecodeCtx(pos_idx=positions, qpos=positions, cos=ctx.cos, sin=ctx.sin, min_key=1) + model.mtp(self._embed(tokens), prev_hidden, rctx.cos, rctx.sin, mtp_cache.layers[0], rctx) + # committed tokens: drafts[:a] then pred[a]; report [a, tok_0..tok_D] + drafts_ext = torch.cat([drafts, pred[:, -1:]], dim=1)[0] + committed = torch.where(ar < accepted, drafts_ext, pred[0]) + self.out_buf[0].copy_(accepted) + self.out_buf[1 : depth + 2].copy_(committed) + a_idx = accepted.view(1, 1) + self.nt_buf.copy_(pred.gather(1, a_idx)) + self.h_buf.copy_(hidden.gather(1, a_idx.view(1, 1, 1).expand(1, 1, hidden.shape[-1]))) + pos.add_(accepted + 1) + + def _adapter_signature(self): + """Captured graphs bake in the adapter wrappers' Python-side decisions (active / + merged / multiplier), so the graph cache is keyed on them.""" + hooks = 0 + for mod in self.model.model.language_model.modules(): + if not isinstance(mod, nn.Linear): + continue + hooks += len(mod._forward_hooks) + len(mod._forward_pre_hooks) + fwd = mod.__dict__.get("forward") + if fwd is None: + continue + lora = getattr(fwd, "__self__", None) + net_ref = getattr(lora, "network_ref", None) + net = net_ref() if callable(net_ref) else None + if net is None: + return ("wrapped", id(lora)) + return (id(net), bool(getattr(net, "is_active", True)), bool(getattr(net, "is_merged_in", False)), repr(getattr(net, "_multiplier", None))) + # inference LoRA sidechains are forward hooks; their count keys the graph too + return ("hooks", hooks) if hooks else None + + def _get_step(self, depth: int): + key = (("mtp", depth) if depth > 0 else ("plain",)) + (self._adapter_signature(),) + fn = (lambda: self._mtp_step(depth)) if depth > 0 else self._plain_step + if not self.use_cuda_graph: + return fn + if key not in self.graphs: + # warm up (triton autotune / allocator) on a side stream, then capture + s = torch.cuda.Stream() + s.wait_stream(torch.cuda.current_stream()) + saved = self._snapshot_state() + with torch.cuda.stream(s): + for _ in range(2): + fn() + self._restore_state(saved) + torch.cuda.current_stream().wait_stream(s) + # two instantiations of the same step: a graph cannot be relaunched until its + # previous replay has finished, so alternating two hides the (ms-scale) CPU + # launch cost of a ~2000-node graph behind the other one's execution + graphs = [] + for _ in range(2): + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g): + fn() + self._restore_state(saved) + graphs.append(g) + self.graphs[key] = graphs + graphs = self.graphs[key] + state = {"i": 0} + + def replay(): + graphs[state["i"]].replay() + state["i"] ^= 1 + + return replay + + def _snapshot_state(self): + c = self.cache + snap = {"pos": c.pos.clone(), "nt": self.nt_buf.clone(), "h": self.h_buf.clone(), "layers": []} + for lc in c.layers: + snap["layers"].append(tuple(t.clone() if t is not None else None for t in (lc.conv, lc.rec))) + snap["mtp_pos"] = self.mtp_cache.pos.clone() if self.mtp_cache is not None else None + return snap + + def _restore_state(self, snap): + c = self.cache + c.pos.copy_(snap["pos"]) + self.nt_buf.copy_(snap["nt"]) + self.h_buf.copy_(snap["h"]) + for lc, (conv, rec) in zip(c.layers, snap["layers"]): + if conv is not None: + lc.conv.copy_(conv) + lc.rec.copy_(rec) + if self.mtp_cache is not None and snap["mtp_pos"] is not None: + self.mtp_cache.pos.copy_(snap["mtp_pos"]) + + # ---------------- prefill + loop + + @torch.no_grad() + def generate( + self, + inputs_embeds: torch.Tensor, + input_ids: torch.Tensor, + position_ids: torch.Tensor, + rope_delta: int, + max_new_tokens: int, + eos_ids, + mtp_depth: int = 0, + progress=None, + ) -> List[int]: + """inputs_embeds [1, T, H] (media already scattered in), input_ids [1, T], + position_ids [3, 1, T]. Returns the generated token ids (eos included).""" + model = self.model + b, t, _ = inputs_embeds.shape + if b != 1: + raise ValueError("Qwen38Generator decodes one sequence at a time") + if mtp_depth > 0 and model.mtp is None: + mtp_depth = 0 + need = t + max_new_tokens + mtp_depth + 2 + if self.cache is None or self.max_len < need: + # round up so a slightly longer prompt does not rebuild caches and graphs + self._ensure(max(need, 2048) if self.cache is None else -(-need // 1024) * 1024) + cache, mtp_cache = self.cache, self.mtp_cache + cache.reset() + if mtp_cache is not None: + mtp_cache.reset() + self.rope_delta.fill_(rope_delta) + eos = set(int(e) for e in eos_ids) + + # prefill (eager): all layers over the prompt, caches written at slots 0..T-1 + lm = model.model.language_model + cos, sin = lm.rotary(position_ids, inputs_embeds.dtype) + hidden = lm(inputs_embeds, cos, sin, cache=cache) # [1, T, H] + nt = self._lm_head_argmax(hidden[:, -1:]) # [1, 1] + cache.pos.fill_(t) + if mtp_depth > 0: + # draft cache over the prompt: slot p = (token_p, hidden_{p-1}), p = 1..T-1 + model.mtp(self._embed(input_ids[:, 1:]), hidden[:, :-1], cos[:, 1:], sin[:, 1:], mtp_cache.layers[0], None, start=1) + self.nt_buf.copy_(nt) + self.h_buf.copy_(hidden[:, -1:]) + out = [int(nt.item())] + if out[-1] in eos: + return out + step = self._get_step(mtp_depth) + n_steps = 0 + n_accepted = 0 + # pipelined: step k+1 is queued before step k's result is read, so the host-side + # bookkeeping overlaps the GPU instead of leaving it idle at every sync. The one + # extra step launched after the stop condition is harmless (its state is discarded). + host = [torch.empty_like(self.out_buf, device="cpu", pin_memory=True) for _ in range(2)] + events = [torch.cuda.Event(), torch.cuda.Event()] + slot = 0 + step() + host[slot].copy_(self.out_buf, non_blocking=True) + events[slot].record() + done = False + while not done: + nxt = slot ^ 1 + step() + host[nxt].copy_(self.out_buf, non_blocking=True) + events[nxt].record() + events[slot].synchronize() + res = host[slot].tolist() + n_steps += 1 + if mtp_depth > 0: + a = res[0] + n_accepted += a + new = res[1 : a + 2] + else: + new = res[1:2] + for tok in new: + out.append(tok) + if tok in eos or len(out) >= max_new_tokens: + done = True + break + if progress is not None: + progress(len(out)) + slot = nxt + torch.cuda.current_stream().synchronize() + self.last_stats = {"steps": n_steps, "tokens": len(out), "accepted_drafts": n_accepted} + return out + + +def build_sequence_embeds(model: Qwen38ForConditionalGeneration, input_ids: torch.Tensor, media_embeds: Optional[torch.Tensor], placeholder_ids): + """Token embeddings with the media placeholder positions replaced by the vision + embeddings. input_ids [T]; media_embeds [N, H] (rows in placeholder order).""" + emb = model.get_input_embeddings()(input_ids[None]) + if media_embeds is not None: + mask = torch.zeros_like(input_ids, dtype=torch.bool) + for pid in placeholder_ids: + mask |= input_ids == pid + n = int(mask.sum()) + if n != media_embeds.shape[0]: + raise ValueError(f"media placeholder count {n} != media embeddings {media_embeds.shape[0]}") + emb = emb.masked_scatter(mask[None, :, None].expand_as(emb), media_embeds.to(emb.dtype)) + return emb diff --git a/extensions_built_in/llm_models/ui.tsx b/extensions_built_in/llm_models/ui.tsx index 0954053..75580a4 100644 --- a/extensions_built_in/llm_models/ui.tsx +++ b/extensions_built_in/llm_models/ui.tsx @@ -81,4 +81,77 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [ ), }, + { + name: "qwen38", + label: "Qwen3.8 27B", + group: "llm", + defaults: { + // default updates when [selected, unselected] in the UI + "config.process[0].model.name_or_path": [ + "Comfy-Org/Qwen3.8-27B/text_encoders/qwen3.8_27b_w4a8.safetensors", + defaultNameOrPath, + ], + "config.process[0].model.quantize": [true, false], + "config.process[0].model.quantize_te": [false, false], + "config.process[0].model.low_vram": [false, false], + // the ComfyUI file ships W4A8 (int4 codebook weights, int8 activations); requesting it keeps the layers as-is + "config.process[0].model.qtype": ["convrotw4a8", "qfloat8"], + "config.process[0].train.unload_text_encoder": [false, false], + "config.process[0].train.noise_scheduler": ["flowmatch", "flowmatch"], + "config.process[0].train.batch_size": [1, 1], + "config.process[0].sample": [ + defaultQwen25OmniSampleConfig, + defaultSampleConfig, + ], + // media is encoded on the GPU per step; the cache is optional + "config.process[0].datasets[x].cache_latents_to_disk": [false, true], + "config.process[0].datasets[x].resolution": [[1024], [512, 768, 1024]], + // the caption is the training target; a blank one trains nothing + "config.process[0].datasets[x].caption_dropout_rate": [0, 0.05], + "config.process[0].model.model_kwargs": [ + { instruction: "Describe this in detail." }, + {}, + ], + }, + disableSections: [ + "network.conv", + "trigger_word", + "train.diff_output_preservation", + "train.blank_prompt_preservation", + "train.unload_text_encoder", + "slider", + ], + additionalSections: [ + "model.model_kwargs.instruction", + "sample.ctrl_img", + "datasets.num_frames", + ], + modelNotes: ( +
+

+ Qwen3.8 27B as a text-generating vision-language model: image or video + in, text out. Each dataset item is a media file (jpg, png, webp, mp4, + ...) with a caption file next to it; the caption is the text the model + learns to produce for that media. Video files are used when Num Frames + is above 1 and are seen as frames only (no audio). +

+

+ The default weights are the ComfyUI W4A8 file (4-bit codebook weights, + 8-bit activations, about 17 GB) and are run in that format: a fused + int4 decode kernel for generation, int8 tensor-core GEMMs for training. + Generation uses the model's MTP draft head for speculative decoding + (model_kwargs.mtp_depth, default 3) and CUDA-graph decode + steps. +

+

+ The instruction in model_kwargs.instruction is the user + turn for every training item; use the same wording when captioning + with the trained LoRA. Thinking mode is off by default + (model_kwargs.enable_thinking). Batch size is 1; the LoRA + trains the text stack only. Watch loss/ce: next-token + cross-entropy on the caption. +

+
+ ), + }, ]; diff --git a/toolkit/inference_lora.py b/toolkit/inference_lora.py index d88a338..8e48ffe 100644 --- a/toolkit/inference_lora.py +++ b/toolkit/inference_lora.py @@ -303,20 +303,37 @@ class InferenceLoRA: x = inputs[0] if inputs else None if x is None: return output - key = (x.device, x.dtype) + key = (x.device, x.dtype, lora.strength) if key not in cache: cache.clear() + # A^T stored [in, r] contiguous: cuBLAS runs x @ A^T 2.5x faster in that + # layout for a few rows; B^T carries scale * strength so the hook is two GEMMs + A = entry.A.to(device=x.device, dtype=x.dtype).t().contiguous() if entry.A is not None and entry.A.dim() == 2 else None + B = entry.B.to(device=x.device, dtype=x.dtype) if entry.B is not None and entry.B.dim() == 2 else None + BT = (B.t() * (entry.scale * lora.strength)).contiguous() if B is not None else None cache[key] = ( - entry.A.to(device=x.device, dtype=x.dtype) if entry.A is not None and entry.A.dim() == 2 else None, - entry.B.to(device=x.device, dtype=x.dtype) if entry.B is not None and entry.B.dim() == 2 else None, + A, + BT, entry.diff.to(device=x.device, dtype=x.dtype) if entry.diff is not None and entry.diff.dim() == 2 else None, entry.diff_b.to(device=x.device, dtype=x.dtype) if entry.diff_b is not None else None, {n: t.to(device=x.device, dtype=x.dtype) for n, t in entry.kron.items()} if entry.kron_is_linear else None, ) - A, B, diff, diff_b, kron = cache[key] + A, BT, diff, diff_b, kron = cache[key] + out = output[0] if isinstance(output, tuple) else output + if A is not None and BT is not None and kron is None and diff is None and diff_b is None \ + and out.dtype == x.dtype and out.shape[-1] == BT.shape[1]: + # lora-only fast path: out + (x A^T) B^T as one addmm, in place on the + # fresh module output when no graph is being recorded + xa = x.reshape(-1, x.shape[-1]) @ A + out2 = out.reshape(-1, out.shape[-1]) + if out.is_contiguous() and out.data_ptr() != x.data_ptr() and not torch.is_grad_enabled(): + out2.addmm_(xa, BT) + return output + res = torch.addmm(out2, xa, BT).view(out.shape) + return (res,) + tuple(output[1:]) if isinstance(output, tuple) else res add = None - if A is not None and B is not None: - add = ((x @ A.t()) @ B.t()) * (entry.scale * lora.strength) + if A is not None and BT is not None: + add = (x @ A) @ BT if kron is not None: dk = _kron_apply(x, kron) * (entry.scale * lora.strength) add = dk if add is None else add + dk diff --git a/toolkit/util/comfy_quant_import.py b/toolkit/util/comfy_quant_import.py index 4634f02..3fe90f5 100644 --- a/toolkit/util/comfy_quant_import.py +++ b/toolkit/util/comfy_quant_import.py @@ -34,6 +34,7 @@ from toolkit.util.nvfp4_quant import ( unswizzle_nvfp4_scales, ) from toolkit.util.ostris_quant import OstrisLinear, get_ostris_quantizer +from toolkit.util.convrot_quant import rotate def parse_comfy_quant_blob(blob: torch.Tensor) -> dict: @@ -71,6 +72,35 @@ class Int8Embedding(torch.nn.Module): return out.to(input_ids.device).reshape(*input_ids.shape, self.embedding_dim) +class ConvRotInt8Embedding(torch.nn.Module): + """Per-row int8 embedding whose rows are stored in the ConvRot-rotated basis (the + storage of a comfy convrot int8 linear). Rows are gathered, dequantized and rotated + back per lookup, so neither the full table nor an unrotated copy materializes.""" + + def __init__(self, qweight: torch.Tensor, scales: torch.Tensor, rot_size: int, dtype: torch.dtype): + super().__init__() + self.num_embeddings, self.embedding_dim = qweight.shape + self.output_dtype = dtype + self.rot_size = int(rot_size) + self.register_buffer("qweight", qweight.contiguous(), persistent=False) + self.register_buffer( + "scales", + scales.detach().float().reshape(-1).contiguous().view(torch.uint8), + persistent=False, + ) + + @property + def weight(self): + return self.forward(torch.arange(self.num_embeddings, device=self.qweight.device)) + + def forward(self, input_ids: torch.Tensor) -> torch.Tensor: + flat = input_ids.reshape(-1).to(self.qweight.device) + rows = self.qweight.index_select(0, flat).float() + scales = self.scales.view(torch.float32).index_select(0, flat) + out = rotate(rows * scales.unsqueeze(1), self.rot_size).to(self.output_dtype) + return out.to(input_ids.device).reshape(*input_ids.shape, self.embedding_dim) + + @torch.no_grad() def split_fused_quantized_keys( state_dict: Dict[str, torch.Tensor], @@ -217,6 +247,7 @@ def import_comfy_quantized_layers( state_dict: Dict[str, torch.Tensor], orig_dtype: torch.dtype = torch.bfloat16, key_map=None, + rotated_embedding_rows: bool = False, ) -> Tuple[Dict[str, torch.Tensor], int]: """Convert every module a ``comfy_quant`` marker points at and attach its quantized tensors. Consumes the quantized entries from ``state_dict`` and @@ -225,6 +256,9 @@ def import_comfy_quantized_layers( ``key_map`` optionally maps a checkpoint prefix to the module path in ``root`` (e.g. comfy text encoder keys onto transformers module paths). + ``rotated_embedding_rows`` keeps a convrot int8 embedding table as int8 rows + (:class:`ConvRotInt8Embedding`, rotated back per lookup) instead of + materializing the full dequantized table. """ state_dict = dict(state_dict) converted = 0 @@ -247,7 +281,15 @@ def import_comfy_quantized_layers( ) parent_path, _, attr = module_path.rpartition(".") parent = root.get_submodule(parent_path) if parent_path else root - if conf.get("convrot"): + if conf.get("convrot") and rotated_embedding_rows: + setattr( + parent, + attr, + ConvRotInt8Embedding( + weight, weight_scale, int(conf.get("convrot_groupsize", 256)), orig_dtype + ), + ) + elif conf.get("convrot"): # rows are stored rotated like a linear's weight: attach to a stand-in linear # (same shape) so the convrot backend undoes the rotation, then keep the table # as a plain embedding (per-row lookups cannot apply the group rotation) @@ -320,10 +362,27 @@ def import_comfy_quantized_layers( pts=state_dict.pop(f"{prefix}.weight_scale_2"), pre_scale=state_dict.pop(f"{prefix}.pre_quant_scale", None), ) + elif fmt == "asym_w4a8_int8": + # int4 codebook codes + fp8 group scales + fp32 channel scales on a convrot + # basis: the exact storage of the toolkit's convrotw4a8 backend + from toolkit.util.convrot_w4a8_quant import ConvRotW4A8Quantizer + + if f"{prefix}.weight_correction" in state_dict: + raise ValueError(f"asymmetric W4A8 (weight_correction) on {prefix} is not supported") + quantizer = get_ostris_quantizer("convrotw4a8") + ConvRotW4A8Quantizer.attach_( + module, + qdata=weight, + srel=state_dict.pop(f"{prefix}.weight_s_rel"), + s_channel=state_dict.pop(f"{prefix}.weight_s_channel"), + codebook=state_dict.pop(f"{prefix}.weight_codebook", None), + rot_size=int(conf.get("convrot_groupsize", 256)) if conf.get("convrot") else 1, + group_size=int(conf.get("group_size", 16)), + ) else: raise ValueError( f"Unsupported comfy quant format {fmt!r} on {prefix} " - "(supported: int8_tensorwise, nvfp4, float8_e4m3fn)" + "(supported: int8_tensorwise, nvfp4, float8_e4m3fn, asym_w4a8_int8)" ) # drop unused calibration extras if present diff --git a/toolkit/util/convrot_quant.py b/toolkit/util/convrot_quant.py index 6e116d6..acd4d04 100644 --- a/toolkit/util/convrot_quant.py +++ b/toolkit/util/convrot_quant.py @@ -51,7 +51,7 @@ import torch.nn.functional as F from toolkit.print import print_acc from toolkit.util.ostris_quant import OstrisQuantizer -CONVROT_QTYPES = ("convrot4", "convrot8", "convrotbitnet", "convrotcomfyw4a4") + tuple( +CONVROT_QTYPES = ("convrot4", "convrot8", "convrotbitnet", "convrotcomfyw4a4", "convrotw4a8") + tuple( f"convrotint{b}" for b in range(2, 9) ) @@ -65,6 +65,11 @@ def get_convrot_quantizer(qtype: str): return ConvRotBitNetQuantizer(rot_size=256) if qtype == "convrotcomfyw4a4": return ConvRotComfyW4A4Quantizer() + if qtype == "convrotw4a8": + # comfy_kitchen asym_w4a8_int8 storage (toolkit/util/convrot_w4a8_quant.py) + from toolkit.util.convrot_w4a8_quant import ConvRotW4A8Quantizer + + return ConvRotW4A8Quantizer(rot_size=256) if qtype.startswith("convrotint"): bits = int(qtype[len("convrotint") :]) if 2 <= bits <= 8: diff --git a/toolkit/util/convrot_w4a8_quant.py b/toolkit/util/convrot_w4a8_quant.py new file mode 100644 index 0000000..1543921 --- /dev/null +++ b/toolkit/util/convrot_w4a8_quant.py @@ -0,0 +1,657 @@ +""" +ConvRot W4A8 codebook backend (qtype ``convrotw4a8``): ComfyUI's ``asym_w4a8_int8`` +storage (comfy_kitchen AsymW4A8Int8Layout) as an OstrisQuantizer. + +Weights are ConvRot-rotated (block regular Hadamard, default 256), every group of +``group_size`` (16) columns gets an fp8-e4m3 relative scale, and each value is the +nearest of 16 Lloyd-Max codebook levels — an int4 code, two per byte (even column = +low nibble). One fp32 scale per output channel maps the group grid onto the int8 grid +the GEMM consumes: + + w_int8[n, k] = clamp(rint(codebook[code[n, k]] * s_rel[n, k // 16]), -127, 127) + W_rot[n, k] = w_int8[n, k] * s_channel[n] + +After decoding the codes to that int8 grid the layer IS a convrot8 layer: per-token +int8 activations in the rotated basis against per-channel int8 weights through +torch._int_mm, straight-through backward on the dequantized weight. This backend +subclasses ConvRotInt8Quantizer and owns only the storage: + + - prefill / training (any M): a Triton kernel decodes int4 -> int8 grid per + forward (a transient the size of an int8 weight), then the parent's _int_mm path. + The training op saves only the packed codes and decodes again in the backward. + - decode (M <= 8): a fused Triton GEMV reads the packed nibbles, fp8 group scales + and codebook directly (int32 multiply-accumulate, no tensor cores), so weight + traffic per token is 4.5 bits/weight — half of convrot8. Bit-identical output. + +Buffers on the module (persistent=False): + cw4_qdata int8 [out, in/2] packed codes + cw4_srel uint8 [out, in/group] fp8 e4m3 group scales (byte view) + cw4_scales uint8 [out*4] fp32 per-channel scales (byte view) + cw4_codebook uint8 [16*4] fp32 codebook (byte view) + cw4_rot_size, cw4_group_size ints +""" + +from typing import Optional + +import torch + +from toolkit.util.convrot_quant import ( + ConvRotInt8Quantizer, + _import_triton, + _int8_act_quant_padded, + _int8_epilogue, + _triton_available, + rotate, +) + +# Lloyd-Max-optimal 16 levels for a group-normalized Gaussian (comfy_kitchen's fixed +# table). ConvRot makes the rotated groups Gaussian, so this table fits every layer. +FIXED_CODEBOOK = ( + -0.980602, -0.794529, -0.638165, -0.500986, -0.377321, -0.263187, -0.155210, -0.050720, + 0.052541, 0.156985, 0.265284, 0.379533, 0.502636, 0.638953, 0.794876, 0.980671, +) +# ALS group-scale refinement passes (comfy_kitchen: iteration 1 captures nearly all the gain) +_ALS_ITERS = 2 +# rows quantized per chunk: caps the fp32 working set of the level grid ([N, groups, 16]) +_QUANT_ROW_ELEM_BUDGET = 1 << 22 + + +# ---------------------------------------------------------------- eager reference + +# int8 grid rounding in both kernels: one saturating cvt.rni.sat.s8 (round-nearest-even, +# saturate to [-128, 127]) instead of rint + clamp(-127, 127). The two differ only for a +# code whose group value rounds to <= -127.5, which the fp8 scale rounding makes possible +# but never happened on any Qwen3.8 tensor measured; prefill and decode share the grid. +W4A8_SAT = True + + +def _pack_codes(codes: torch.Tensor) -> torch.Tensor: + """unsigned int codes [N, K] in 0..15 -> int8 [N, K/2], even column = low nibble.""" + c = codes.to(torch.int32) + return ((c[:, 0::2] & 0xF) | ((c[:, 1::2] & 0xF) << 4)).to(torch.uint8).view(torch.int8).contiguous() + + +def _unpack_codes(qdata: torch.Tensor) -> torch.Tensor: + """int8 [N, K/2] packed -> int64 codes [N, K].""" + b = qdata.view(torch.uint8).to(torch.int64) + n, kh = b.shape + codes = torch.empty(n, kh * 2, dtype=torch.int64, device=qdata.device) + codes[:, 0::2] = b & 0xF + codes[:, 1::2] = (b >> 4) & 0xF + return codes + + +def unpack_w4a8_eager(qdata: torch.Tensor, srel: torch.Tensor, codebook: torch.Tensor, group_size: int) -> torch.Tensor: + """Reference decode of packed codes to the int8 grid (comfy_kitchen eager math).""" + n = qdata.shape[0] + k = srel.shape[1] * group_size + values = codebook.float().to(qdata.device)[_unpack_codes(qdata)] + values = values.view(n, k // group_size, group_size) * srel.float().unsqueeze(-1) + lo = -128.0 if W4A8_SAT else -127.0 + return values.view(n, k).round().clamp_(lo, 127).to(torch.int8) + + +def _assign_grid(target: torch.Tensor, levels: torch.Tensor, stochastic: bool = False) -> torch.Tensor: + """Nearest level per weight. target [N, G, g] = w_rot / s_channel; levels [N, G, 16] + non-decreasing. stochastic: pick the upper bracket level with probability + proportional to the position inside the bracket (unbiased; a merged LoRA delta + below the int4 step survives instead of rounding away).""" + n, groups, gsize = target.shape + lv = levels.reshape(n * groups, -1).contiguous() + last = lv.shape[1] - 1 + tg = target.reshape(n * groups, gsize).contiguous() + pos = torch.searchsorted(lv, tg) + if stochastic: + lo = (pos - 1).clamp_(0, last - 1) + level_lo = torch.gather(lv, 1, lo) + width = torch.gather(lv, 1, lo + 1) - level_lo + thr = torch.rand_like(tg) * width + level_lo + return (lo + (tg > thr)).clamp_(0, last).view(n, groups, gsize) + lo = (pos - 1).clamp(0, last) + hi = pos.clamp(0, last) + dlo = (tg - torch.gather(lv, 1, lo)).abs_() + dhi = (tg - torch.gather(lv, 1, hi)).abs_() + return torch.where(dhi < dlo, hi, lo).view(n, groups, gsize) + + +def _assign_codes(normalized: torch.Tensor, codebook: torch.Tensor) -> torch.Tensor: + """Nearest codebook index (sorted codebook, binary search + neighbour compare).""" + last = codebook.numel() - 1 + pos = torch.searchsorted(codebook, normalized.contiguous()) + lo = (pos - 1).clamp(0, last) + hi = pos.clamp(0, last) + dlo = (normalized - codebook[lo]).abs_() + dhi = (normalized - codebook[hi]).abs_() + return torch.where(dhi < dlo, hi, lo) + + +def grid_levels(codebook: torch.Tensor, srel: torch.Tensor) -> torch.Tensor: + """The int8 value each codebook level decodes to, per group: [N, G, 16].""" + return (codebook.float().view(1, 1, -1) * srel.float().unsqueeze(-1)).round_().clamp_(-127, 127) + + +@torch.no_grad() +def _quantize_rotated_rows(w_rot: torch.Tensor, codebook: torch.Tensor, group_size: int): + """One row chunk of comfy_kitchen's symmetric codebook W4A8 quantization. + Returns (packed int8 [n, K/2], s_rel fp8 [n, G], s_channel fp32 [n]).""" + n, k = w_rot.shape + groups = k // group_size + gw = w_rot.float().view(n, groups, group_size) + gs = gw.abs().amax(dim=-1, keepdim=True).clamp(min=1e-8) + q = _assign_codes(gw / gs, codebook) + for _ in range(_ALS_ITERS): + lv = codebook[q] + gs = ((gw * lv).sum(-1, keepdim=True) / (lv * lv).sum(-1, keepdim=True).clamp(min=1e-8)).clamp(min=1e-8) + q = _assign_codes(gw / gs, codebook) + shifted = codebook[q] * gs + s_channel = (shifted.abs().amax(dim=(1, 2)) / 127.0).clamp(min=1e-8) + srel = (gs.squeeze(-1) / s_channel.unsqueeze(1)).to(torch.float8_e4m3fn).contiguous() + # final assignment through the int8 grid the kernels decode to + codes = _assign_grid(gw / s_channel.view(-1, 1, 1), grid_levels(codebook, srel)) + return _pack_codes(codes.view(n, k)), srel, s_channel.float().contiguous() + + +@torch.no_grad() +def quantize_w4a8_rotated(w_rot: torch.Tensor, group_size: int = 16, codebook: Optional[torch.Tensor] = None): + """Quantize an already-rotated fp weight [N, K] into W4A8 storage, row-chunked. + Returns (packed, s_rel fp8, s_channel fp32, codebook fp32).""" + n, k = w_rot.shape + if k % group_size or k % 16: + raise ValueError(f"W4A8 needs K divisible by 16 and group_size={group_size}, got K={k}") + if codebook is None: + codebook = torch.tensor(FIXED_CODEBOOK, dtype=torch.float32, device=w_rot.device) + codebook = codebook.to(device=w_rot.device, dtype=torch.float32).contiguous() + block = max(1, _QUANT_ROW_ELEM_BUDGET // max(k, 1)) + parts = [_quantize_rotated_rows(w_rot[r0 : r0 + block], codebook, group_size) for r0 in range(0, n, block)] + packed = torch.cat([p[0] for p in parts], dim=0) + srel = torch.cat([p[1] for p in parts], dim=0) + sch = torch.cat([p[2] for p in parts], dim=0) + return packed, srel, sch, codebook + + +# ---------------------------------------------------------------- triton kernels + +_kernels = None + + +def _get_kernels(): + global _kernels + if _kernels is not None: + return _kernels + triton, tl, libdevice = _import_triton() + + @triton.jit + def _fp8_rep8(s_u8): + # fp8 e4m3 byte -> fp32, then each group scale repeated over its 8 packed bytes + s = s_u8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + s = tl.interleave(s, s) + s = tl.interleave(s, s) + return tl.interleave(s, s) + + @triton.jit + def _sat_s8(v): + # round-to-nearest-even + saturate to int8 in one PTX cvt (-128 reachable; see SAT) + return tl.inline_asm_elementwise( + "cvt.rni.sat.s8.f32 $0, $1;", "=r,r", [v], dtype=tl.int32, is_pure=True, pack=1 + ) + + @triton.jit + def w4a8_unpack_kernel( + q_ptr, s_ptr, cb_ptr, o_ptr, + N, KH, G, + BLOCK_N: tl.constexpr, BLOCK_KH: tl.constexpr, SAT: tl.constexpr, + ): + # one [BLOCK_N, BLOCK_KH] tile of packed bytes -> [BLOCK_N, 2*BLOCK_KH] int8 grid. + # group_size is 16 (8 bytes per group): BLOCK_KH is a multiple of 8. + pid_n = tl.program_id(0) + pid_k = tl.program_id(1) + rows = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + bcols = pid_k * BLOCK_KH + tl.arange(0, BLOCK_KH) + gcols = pid_k * (BLOCK_KH // 8) + tl.arange(0, BLOCK_KH // 8) + mask_r = rows < N + m = mask_r[:, None] & (bcols[None, :] < KH) + byte = tl.load(q_ptr + rows[:, None] * KH + bcols[None, :], mask=m, other=0).to(tl.int32) & 0xFF + low = byte & 0xF + high = (byte >> 4) & 0xF + s_u8 = tl.load(s_ptr + rows[:, None] * G + gcols[None, :], mask=mask_r[:, None] & (gcols[None, :] < G), other=0) + s8 = _fp8_rep8(s_u8) + if SAT: + q_low = _sat_s8(tl.load(cb_ptr + low) * s8).to(tl.int8) + q_high = _sat_s8(tl.load(cb_ptr + high) * s8).to(tl.int8) + else: + q_low = tl.clamp(libdevice.rint(tl.load(cb_ptr + low) * s8), -127.0, 127.0).to(tl.int8) + q_high = tl.clamp(libdevice.rint(tl.load(cb_ptr + high) * s8), -127.0, 127.0).to(tl.int8) + K = KH * 2 + tl.store(o_ptr + rows[:, None] * K + (2 * bcols)[None, :], q_low, mask=m) + tl.store(o_ptr + rows[:, None] * K + (2 * bcols + 1)[None, :], q_high, mask=m) + + @triton.jit + def w4a8_gemv_kernel( + xq_ptr, xs_ptr, q_ptr, s_ptr, cb_ptr, ws_ptr, b_ptr, o_ptr, + K, N, + MP: tl.constexpr, + HAS_BIAS: tl.constexpr, + BLOCK_N: tl.constexpr, + BLOCK_K: tl.constexpr, + SAT: tl.constexpr, + ): + # decode GEMV without tensor cores: the int8-grid weight tile is decoded in + # registers and multiply-accumulated (int32 IMAD, exact and order independent) + # against the pre-quantized int8 activation rows, so the output is bit-identical + # to the _int_mm path. Feeding tl.dot instead costs a shared-memory layout + # round trip per tile and ran 3x slower on Blackwell. + # xq [MP, K] int8 (rows >= M are zero), xs [MP] fp32 row scales, out [MP, N]. + pid = tl.program_id(0) + offs_n = pid * BLOCK_N + tl.arange(0, BLOCK_N) + mask_n = offs_n < N + offs_m = tl.arange(0, MP) + KH = K // 2 + G = K // 16 + # the products accumulate over K in a [MP, BN, BK/2] register tile; one cross-lane + # reduction at the end instead of one per K tile (~25% faster) + acc_t = tl.zeros((MP, BLOCK_N, BLOCK_K // 2), tl.int32) + for k0 in range(0, K, BLOCK_K): + offs_b = k0 // 2 + tl.arange(0, BLOCK_K // 2) + m_b = mask_n[:, None] & (offs_b[None, :] < KH) + byte = tl.load(q_ptr + offs_n[:, None] * KH + offs_b[None, :], mask=m_b, other=0).to(tl.int32) & 0xFF + low = byte & 0xF + high = (byte >> 4) & 0xF + # group_size 16 = 8 packed bytes per fp8 group scale + s_u8 = tl.load(s_ptr + offs_n[:, None] * G + (offs_b // 8)[None, :], mask=m_b, other=0) + s = s_u8.to(tl.float8e4nv, bitcast=True).to(tl.float32) + vl = tl.load(cb_ptr + low) * s + vh = tl.load(cb_ptr + high) * s + if SAT: + ql = _sat_s8(vl) + qh = _sat_s8(vh) + else: + ql = tl.clamp(libdevice.rint(vl), -127.0, 127.0).to(tl.int32) + qh = tl.clamp(libdevice.rint(vh), -127.0, 127.0).to(tl.int32) + offs_ke = k0 + 2 * tl.arange(0, BLOCK_K // 2) + m_x = offs_ke[None, :] < K + xe = tl.load(xq_ptr + offs_m[:, None] * K + offs_ke[None, :], mask=m_x, other=0).to(tl.int32) + xo = tl.load(xq_ptr + offs_m[:, None] * K + (offs_ke + 1)[None, :], mask=m_x, other=0).to(tl.int32) + acc_t += ql[None, :, :] * xe[:, None, :] + qh[None, :, :] * xo[:, None, :] + acc = tl.sum(acc_t, axis=2) # [MP, BN] + ws = tl.load(ws_ptr + offs_n, mask=mask_n, other=0.0) + xs = tl.load(xs_ptr + offs_m) + out = acc.to(tl.float32) * (xs[:, None] * ws[None, :]) + if HAS_BIAS: + out += tl.load(b_ptr + offs_n, mask=mask_n, other=0.0).to(tl.float32)[None, :] + tl.store(o_ptr + offs_m[:, None] * N + offs_n[None, :], out.to(o_ptr.dtype.element_ty), mask=mask_n[None, :]) + + @triton.jit + def w4a8_act_quant_kernel(x_ptr, q_ptr, s_ptr, M, K, QMAX: tl.constexpr, MP: tl.constexpr, BLOCK_K: tl.constexpr): + # per-row symmetric int8 quant of x [M, K] into exactly MP rows (rows >= M: zeros, + # scale 1) -- the gemv's padded operand without any fill kernels. Same arithmetic + # as int8_act_quant_kernel: max is exact under any blocking. + row = tl.program_id(0) + valid = row < M + base = row * K + acc = tl.zeros((BLOCK_K,), tl.float32) + for k0 in range(0, K, BLOCK_K): + offs = k0 + tl.arange(0, BLOCK_K) + v = tl.load(x_ptr + base + offs, mask=(offs < K) & valid, other=0.0).to(tl.float32) + acc = tl.maximum(acc, tl.abs(v)) + amax = tl.max(acc, axis=0) + scale = tl.where(amax > 0, amax / QMAX, 1.0) + for k0 in range(0, K, BLOCK_K): + offs = k0 + tl.arange(0, BLOCK_K) + mask = offs < K + v = tl.load(x_ptr + base + offs, mask=mask & valid, other=0.0).to(tl.float32) + q = libdevice.rint(v / scale) + q = tl.minimum(tl.maximum(q, -1.0 * QMAX), 1.0 * QMAX) + tl.store(q_ptr + row * K + offs, q.to(tl.int8), mask=mask) + tl.store(s_ptr + row, scale) + + _kernels = (w4a8_unpack_kernel, w4a8_gemv_kernel, w4a8_act_quant_kernel) + return _kernels + + +def _unpack_impl(qdata: torch.Tensor, srel_u8: torch.Tensor, cb_u8: torch.Tensor, group_size: int) -> torch.Tensor: + n, kh = qdata.shape + k = kh * 2 + if group_size != 16: + return unpack_w4a8_eager(qdata, srel_u8.view(torch.float8_e4m3fn), cb_u8.view(torch.float32), group_size) + out = torch.empty(n, k, dtype=torch.int8, device=qdata.device) + kernel, _, _ = _get_kernels() + block_n, block_kh = 32, 128 + grid = (-(-n // block_n), -(-kh // block_kh)) + kernel[grid](qdata, srel_u8, cb_u8.view(torch.float32), out, n, kh, k // 16, + BLOCK_N=block_n, BLOCK_KH=block_kh, SAT=W4A8_SAT, num_warps=4) + return out + + +@torch.library.custom_op("ostris::convrot_w4a8_unpack", mutates_args=()) +def _w4a8_unpack_op(qdata: torch.Tensor, srel_u8: torch.Tensor, cb_u8: torch.Tensor, group_size: int) -> torch.Tensor: + if not (_triton_available() and qdata.is_cuda): + return unpack_w4a8_eager(qdata, srel_u8.view(torch.float8_e4m3fn), cb_u8.view(torch.float32), group_size) + return _unpack_impl(qdata, srel_u8, cb_u8, group_size) + + +@_w4a8_unpack_op.register_fake +def _w4a8_unpack_fake(qdata, srel_u8, cb_u8, group_size): + return torch.empty(qdata.shape[0], qdata.shape[1] * 2, dtype=torch.int8, device=qdata.device) + + +# decode GEMV: M <= W4A8_GEMV_MAX_M rows. the activation rows are int8-quantized by +# the shared convrot act-quant op (bit-identical to the GEMM path) and padded to MP. +W4A8_GEMV_MAX_M = 8 + + +def _gemv_mp(m: int) -> int: + return 1 << max(0, (m - 1).bit_length()) # 1, 2, 4, 8 + + +@torch.library.custom_op("ostris::convrot_w4a8_act_quant", mutates_args=()) +def _w4a8_act_quant_op(x2d: torch.Tensor, act_qmax: int) -> list[torch.Tensor]: + """Rotated activation rows -> (int8 [MP, K], fp32 scales [MP]), MP = next pow2 of M.""" + m, k = x2d.shape + mp = _gemv_mp(m) + x2d = x2d.contiguous() + q = torch.empty(mp, k, device=x2d.device, dtype=torch.int8) + s = torch.empty(mp, device=x2d.device, dtype=torch.float32) + _, _, kernel = _get_kernels() + block_k = min(2048, 1 << (k - 1).bit_length()) + kernel[(mp,)](x2d, q, s, m, k, QMAX=act_qmax, MP=mp, BLOCK_K=block_k, num_warps=8) + return [q, s] + + +@_w4a8_act_quant_op.register_fake +def _w4a8_act_quant_fake(x2d, act_qmax): + m, k = x2d.shape + mp = _gemv_mp(m) + return [torch.empty(mp, k, device=x2d.device, dtype=torch.int8), torch.empty(mp, device=x2d.device, dtype=torch.float32)] + + +@torch.library.custom_op("ostris::convrot_w4a8_gemv_prepared", mutates_args=()) +def _w4a8_gemv_prepared_op( + xq: torch.Tensor, + xs: torch.Tensor, + qdata: torch.Tensor, + srel_u8: torch.Tensor, + cb_u8: torch.Tensor, + scales_u8: torch.Tensor, + bias: Optional[torch.Tensor], + m: int, + out_dtype: str, +) -> torch.Tensor: + """GEMV on already-quantized activation rows (xq [MP, K] int8, xs [MP] fp32).""" + mp, k = xq.shape + n = qdata.shape[0] + out = torch.empty(mp, n, device=xq.device, dtype=getattr(torch, out_dtype)) + ws = scales_u8.view(torch.float32) + _, kernel, _ = _get_kernels() + # swept on RTX PRO 6000 (Blackwell) over the qwen3.8 shapes with HBM-resident weights: + # 8-row tiles (more programs in flight for the N<=12k projections); the K tile and + # warp count follow M to keep the [MP, BN, BK/2] accumulator in registers + if mp <= 2: + block_n, block_k, warps = 8, 512, 4 + elif mp <= 4: + block_n, block_k, warps = 8, 256, 4 + else: + block_n, block_k, warps = 16, 256, 8 + kernel[(-(-n // block_n),)]( + xq, xs, qdata, srel_u8, cb_u8.view(torch.float32), ws, + bias if bias is not None else ws, out, + k, n, + MP=mp, HAS_BIAS=bias is not None, BLOCK_N=block_n, BLOCK_K=block_k, SAT=W4A8_SAT, + num_warps=warps, num_stages=2, + ) + return out[:m] + + +@_w4a8_gemv_prepared_op.register_fake +def _w4a8_gemv_prepared_fake(xq, xs, qdata, srel_u8, cb_u8, scales_u8, bias, m, out_dtype): + return torch.empty(m, qdata.shape[0], device=xq.device, dtype=getattr(torch, out_dtype)) + + +def _w4a8_gemv_op(x2d, qdata, srel_u8, cb_u8, scales_u8, bias, act_qmax, out_dtype): + xq, xs = _w4a8_act_quant_op(x2d, act_qmax) + return _w4a8_gemv_prepared_op(xq, xs, qdata, srel_u8, cb_u8, scales_u8, bias, x2d.shape[0], out_dtype) + + +# training-path linear: int8 tensor-core forward on the decoded grid (bit-identical to +# inference), straight-through backward on the dequantized weight. autograd keeps only +# the packed codes and decodes again in the backward, so no layer holds an unpacked +# int8 weight between forward and backward. +@torch.library.custom_op("ostris::convrot_w4a8_linear_ste", mutates_args=()) +def _w4a8_linear_ste_op( + x2d: torch.Tensor, + qdata: torch.Tensor, + srel_u8: torch.Tensor, + cb_u8: torch.Tensor, + scales_u8: torch.Tensor, + bias: Optional[torch.Tensor], + group_size: int, + act_qmax: int, + out_dtype: str, +) -> torch.Tensor: + m = x2d.shape[0] + w8 = _w4a8_unpack_op(qdata, srel_u8, cb_u8, group_size) + aq, a_s = _int8_act_quant_padded(x2d, act_qmax) + i32 = torch._int_mm(aq, w8.t()) + return _int8_epilogue(i32[:m], a_s[:m], scales_u8.view(torch.float32), bias, getattr(torch, out_dtype)) + + +@_w4a8_linear_ste_op.register_fake +def _w4a8_linear_ste_fake(x2d, qdata, srel_u8, cb_u8, scales_u8, bias, group_size, act_qmax, out_dtype): + return torch.empty(x2d.shape[0], qdata.shape[0], device=x2d.device, dtype=getattr(torch, out_dtype)) + + +def _w4a8_linear_ste_setup(ctx, inputs, output): + x2d, qdata, srel_u8, cb_u8, scales_u8, bias, group_size, act_qmax, out_dtype = inputs + ctx.group_size = group_size + ctx.save_for_backward(qdata, srel_u8, cb_u8, scales_u8) + + +def _w4a8_linear_ste_backward(ctx, grad): + qdata, srel_u8, cb_u8, scales_u8 = ctx.saved_tensors + w8 = _w4a8_unpack_op(qdata, srel_u8, cb_u8, ctx.group_size) + w = w8.to(grad.dtype) * scales_u8.view(torch.float32).to(grad.dtype).unsqueeze(1) + return grad @ w, None, None, None, None, None, None, None, None + + +_w4a8_linear_ste_op.register_autograd(_w4a8_linear_ste_backward, setup_context=_w4a8_linear_ste_setup) + + +# ---------------------------------------------------------------- the backend + + +class ConvRotW4A8Quantizer(ConvRotInt8Quantizer): + """ConvRot W4A8 codebook backend: convrot8 numerics on an int4 codebook weight grid. + One instance per qtype, shareable across modules.""" + + wants_fp32_weight = True + + def __init__(self, rot_size: int = 256, group_size: int = 16): + super().__init__(rot_size=rot_size) + self.group_size = group_size + + def can_quantize(self, module: torch.nn.Linear) -> bool: + if module.in_features % self.group_size: + return False + return super().can_quantize(module) + + @staticmethod + def attach_( + module: torch.nn.Module, + qdata: torch.Tensor, + srel: torch.Tensor, + s_channel: torch.Tensor, + codebook: Optional[torch.Tensor], + rot_size: int, + group_size: int = 16, + ) -> None: + """Attach packed W4A8 tensors (comfy_kitchen storage) to a linear.""" + if codebook is None: + # symmetric uniform codes: value = code - 8 + codebook = torch.arange(-8, 8, dtype=torch.float32) + if srel.dtype != torch.float8_e4m3fn: + srel = srel.view(torch.float8_e4m3fn) if srel.dtype == torch.uint8 else srel.to(torch.float8_e4m3fn) + n, kh = qdata.shape + if (n, kh * 2 // group_size) != tuple(srel.shape): + raise ValueError(f"W4A8 scale shape {tuple(srel.shape)} does not match packed weight {tuple(qdata.shape)}") + module.register_buffer("cw4_qdata", qdata.contiguous(), persistent=False) + module.register_buffer("cw4_srel", srel.contiguous().view(torch.uint8), persistent=False) + module.register_buffer( + "cw4_scales", s_channel.detach().float().reshape(-1).contiguous().view(torch.uint8), persistent=False + ) + module.register_buffer( + "cw4_codebook", codebook.detach().float().reshape(-1).contiguous().view(torch.uint8), persistent=False + ) + module.cw4_rot_size = int(rot_size) + module.cw4_group_size = int(group_size) + + def quantize_(self, module: torch.nn.Linear, weight_fp32: torch.Tensor) -> None: + rot = self._rot_for(module.in_features) + packed, srel, sch, cb = quantize_w4a8_rotated(rotate(weight_fp32.float(), rot), self.group_size) + self.attach_(module, packed, srel, sch, cb, rot, self.group_size) + + # storage accessors used by the parent's forward / dequantize paths + def _rot(self, module) -> int: + return module.cw4_rot_size + + def _qdata(self, module) -> torch.Tensor: + return _w4a8_unpack_op(module.cw4_qdata, module.cw4_srel, module.cw4_codebook, module.cw4_group_size) + + def _scales_u8(self, module) -> torch.Tensor: + return module.cw4_scales + + def _gemv_args(self, module): + # the parent's packed gemv does not know the codebook grid; decode runs below + return None + + def _codebook(self, module) -> torch.Tensor: + return module.cw4_codebook.view(torch.float32) + + def _srel(self, module) -> torch.Tensor: + return module.cw4_srel.view(torch.float8_e4m3fn) + + def _linear_ste(self, module, x2d: torch.Tensor, out_dtype: str) -> torch.Tensor: + return _w4a8_linear_ste_op( + x2d, module.cw4_qdata, module.cw4_srel, module.cw4_codebook, module.cw4_scales, + module.bias, module.cw4_group_size, self.act_qmax, out_dtype, + ) + + def fake_quant_rotated_weight(self, module, w_rot: torch.Tensor) -> torch.Tensor: + """dequant(quant(w_rot)) on the STORED group/channel scales and codebook. fp32.""" + n, k = w_rot.shape + g = module.cw4_group_size + sch = self._scales(module) + levels = grid_levels(self._codebook(module), self._srel(module)) + codes = _assign_grid(w_rot.float().view(n, k // g, g) / sch.view(-1, 1, 1), levels) + return (torch.gather(levels, 2, codes).view(n, k) * sch.unsqueeze(1)) + + def requantize_(self, module, fp_weight: torch.Tensor) -> None: + w = fp_weight.to(device=module.cw4_qdata.device, dtype=torch.float32) + packed, srel, sch, cb = quantize_w4a8_rotated( + rotate(w, module.cw4_rot_size), module.cw4_group_size, self._codebook(module) + ) + # in place: fused views over several modules' storage stay coherent + module.cw4_qdata.copy_(packed) + module.cw4_srel.copy_(srel.view(torch.uint8)) + module.cw4_scales.copy_(sch.view(torch.uint8)) + + @torch.no_grad() + def requantize_codes_(self, module, fp_weight: torch.Tensor, stochastic: bool = False) -> None: + """Re-assign only the codes on the module's STORED scales (the grid a QAT run + trains against). stochastic: unbiased rounding between the bracketing levels.""" + w = fp_weight.to(device=module.cw4_qdata.device, dtype=torch.float32) + n, k = w.shape + g = module.cw4_group_size + w_rot = rotate(w, module.cw4_rot_size) + levels = grid_levels(self._codebook(module), self._srel(module)) + codes = _assign_grid(w_rot.view(n, k // g, g) / self._scales(module).view(-1, 1, 1), levels, stochastic) + module.cw4_qdata.copy_(_pack_codes(codes.view(n, k))) + + def gemv_ok(self, module, x: torch.Tensor) -> bool: + """Whether the fused decode GEMV handles this (module, input).""" + return ( + not x.requires_grad + and x.numel() // module.in_features <= W4A8_GEMV_MAX_M + and x.is_cuda + and module.cw4_group_size == 16 + and module.in_features % 16 == 0 + and _triton_available() + ) + + def prepare_act(self, module, x: torch.Tensor): + """Rotate + int8-quantize the activation rows once; shareable by every W4A8 + linear that consumes the same input (same in_features / rotation).""" + return _w4a8_act_quant_op(rotate(x, module.cw4_rot_size).reshape(-1, module.in_features), self.act_qmax) + + def gemv_prepared(self, module, xq: torch.Tensor, xs: torch.Tensor, shape, dtype: torch.dtype) -> torch.Tensor: + m = 1 + for d in shape[:-1]: + m *= d + out = _w4a8_gemv_prepared_op( + xq, xs, module.cw4_qdata, module.cw4_srel, module.cw4_codebook, module.cw4_scales, + module.bias, m, str(dtype).split(".")[-1], + ) + return out.reshape(*shape[:-1], module.out_features) + + def forward(self, module, x: torch.Tensor) -> torch.Tensor: + if self.gemv_ok(module, x): + xq, xs = self.prepare_act(module, x) + return self.gemv_prepared(module, xq, xs, x.shape, x.dtype) + return super().forward(module, x) + + +def module_is_plain(mod) -> bool: + """True when calling mod(x) is just its class forward: no instance-level forward + (training LoRA wrapper) and no forward hooks (inference LoRA sidechains).""" + return "forward" not in mod.__dict__ and not mod._forward_hooks and not mod._forward_pre_hooks + + +def module_hooks_only(mod) -> bool: + """Class forward plus plain forward hooks (what inference LoRA sidechains register): + the output can be computed by a fast path and the hooks applied by apply_forward_hooks.""" + from torch.nn.modules import module as _m + + return ( + "forward" not in mod.__dict__ + and not mod._forward_pre_hooks + and not mod._forward_hooks_with_kwargs + and not _m._global_forward_hooks + and not _m._global_forward_pre_hooks + ) + + +def apply_forward_hooks(mod, x: torch.Tensor, out: torch.Tensor) -> torch.Tensor: + """Run mod's forward hooks on an output computed outside mod.__call__.""" + for hook in mod._forward_hooks.values(): + r = hook(mod, (x,), out) + if r is not None: + out = r + return out + + +def w4a8_grouped_linear(modules, x: torch.Tensor): + """Outputs of several linears fed the same input. W4A8 OstrisLinears that share the + rotation and are not wrapped (e.g. by an active LoRA) share one rotate+quant of x on + the decode path; everything else just runs module(x).""" + outs = [None] * len(modules) + shared = None + for i, mod in enumerate(modules): + q = getattr(mod, "ostris_quantizer", None) + if ( + isinstance(q, ConvRotW4A8Quantizer) + and module_hooks_only(mod) + and q.gemv_ok(mod, x) + ): + if shared is None: + shared = (mod.cw4_rot_size, mod.in_features, q.prepare_act(mod, x)) + if shared[0] == mod.cw4_rot_size and shared[1] == mod.in_features: + out = q.gemv_prepared(mod, *shared[2], x.shape, x.dtype) + outs[i] = apply_forward_hooks(mod, x, out) if mod._forward_hooks else out + continue + outs[i] = mod(x) + return outs + + +def dequantize_w4a8(qdata, srel, s_channel, codebook, rot_size: int, group_size: int = 16, dtype=torch.float32): + """Full weight in the original basis from comfy_kitchen W4A8 tensors (reference path).""" + w8 = unpack_w4a8_eager(qdata, srel, codebook, group_size) + w_rot = w8.float() * s_channel.float().to(w8.device).unsqueeze(1) + return rotate(w_rot, rot_size).to(dtype) diff --git a/ui/src/app/jobs/new/options.tsx b/ui/src/app/jobs/new/options.tsx index 29fcbf7..4907c4f 100644 --- a/ui/src/app/jobs/new/options.tsx +++ b/ui/src/app/jobs/new/options.tsx @@ -151,6 +151,7 @@ export const quantizationOptions: SelectOption[] = [ { value: 'uint4', label: '4 bit' }, { value: 'uint3', label: '3 bit' }, { value: 'uint2', label: '2 bit' }, + { value: 'convrotw4a8', label: '4bit convrot W4A8 (int4 weights, int8 activations)' }, ]; export const defaultQtype = 'qfloat8'; diff --git a/ui/src/components/CaptionSimpleJob.tsx b/ui/src/components/CaptionSimpleJob.tsx index 4c0ee6d..13ee888 100644 --- a/ui/src/components/CaptionSimpleJob.tsx +++ b/ui/src/components/CaptionSimpleJob.tsx @@ -20,7 +20,9 @@ import { groupedCaptionerTypes, maxNewTokensOptions, maxResOptions, + mtpDepthOptions, quantizationOptions, + videoFpsOptions, } from '@/helpers/captionOptions'; type Props = { @@ -243,6 +245,21 @@ const CaptionSimpleJob: React.FC = ({ jobConfig, setJobConfig, gpuIDs, se /> )} + {additionalSections.includes('caption.video_fps') && ( +
+ { + const num = parseFloat(value); + if (!isNaN(num)) { + setJobConfig(num, 'config.process[0].caption.video_fps'); + } + }} + options={videoFpsOptions} + /> +
+ )} {additionalSections.includes('caption.batch_size') && (
= ({ jobConfig, setJobConfig, gpuIDs, se onChange={value => setJobConfig(value, 'config.process[0].caption.thinking')} /> )} + {additionalSections.includes('caption.mtp') && ( + <> + setJobConfig(value, 'config.process[0].caption.mtp')} + /> + {(jobConfig.config.process[0].caption.mtp ?? true) && ( +
+ { + const intVal = parseInt(value); + if (!isNaN(intVal)) { + setJobConfig(intVal, 'config.process[0].caption.mtp_depth'); + } + }} + options={mtpDepthOptions} + /> +
+ )} + + )} + {additionalSections.includes('caption.cuda_graph') && ( + setJobConfig(value, 'config.process[0].caption.cuda_graph')} + /> + )} + {additionalSections.includes('caption.vision_bf16') && ( + setJobConfig(value, 'config.process[0].caption.vision_bf16')} + /> + )} {additionalSections.includes('caption.layer_offloading') && ( <> = ({ jobConfig, setJobConfig, gpuIDs, se />
)} + {additionalSections.includes('caption.system_prompt') && ( +
+ { + setJobConfig(value?.trim() === '' ? undefined : value, 'config.process[0].caption.system_prompt'); + }} + placeholder="Optional system message" + /> +
+ )}