Added initial support for fine tuning and captioning with Qwen 3.8 27B

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

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,
 ]

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,
 ]

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 = "<<AITK_INSTRUCTION>>"
+        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 "</think>" in text:
+            text = text.split("</think>")[-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 <sample>.txt by GenerateImageConfig.save_image
+        return self.caption(gen_config.ctrl_img, gen_config.prompt)

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)

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

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

extensions_built_in/llm_models/ui.tsx

@@ -81,4 +81,77 @@ export const AI_TOOLKIT_UI_MODELS: ModelArch[] = [
       </div>
     ),
   },
+  {
+    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: (
+      <div className="space-y-2">
+        <p>
+          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).
+        </p>
+        <p>
+          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&apos;s MTP draft head for speculative decoding
+          (<code>model_kwargs.mtp_depth</code>, default 3) and CUDA-graph decode
+          steps.
+        </p>
+        <p>
+          The instruction in <code>model_kwargs.instruction</code> is the user
+          turn for every training item; use the same wording when captioning
+          with the trained LoRA. Thinking mode is off by default
+          (<code>model_kwargs.enable_thinking</code>). Batch size is 1; the LoRA
+          trains the text stack only. Watch <code>loss/ce</code>: next-token
+          cross-entropy on the caption.
+        </p>
+      </div>
+    ),
+  },
 ];

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,
+                            A,
-                            entry.B.to(device=x.device, dtype=x.dtype) if entry.B is not None and entry.B.dim() == 2 else None,
+                            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:
+                    if A is not None and BT is not None:
-                        add = ((x @ A.t()) @ B.t()) * (entry.scale * lora.strength)
+                        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

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

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:

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)

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';

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<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
               />
             </div>
           )}
+          {additionalSections.includes('caption.video_fps') && (
+            <div className="mt-4">
+              <SelectInput
+                label="Video FPS"
+                value={`${jobConfig.config.process[0].caption.video_fps ?? 2}`}
+                onChange={value => {
+                  const num = parseFloat(value);
+                  if (!isNaN(num)) {
+                    setJobConfig(num, 'config.process[0].caption.video_fps');
+                  }
+                }}
+                options={videoFpsOptions}
+              />
+            </div>
+          )}
           {additionalSections.includes('caption.batch_size') && (
             <div className="mt-4">
               <SelectInput
@@ -297,6 +314,44 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
                 onChange={value => setJobConfig(value, 'config.process[0].caption.thinking')}
               />
             )}
+            {additionalSections.includes('caption.mtp') && (
+              <>
+                <Checkbox
+                  label="MTP Speculative Decoding"
+                  checked={jobConfig.config.process[0].caption.mtp ?? true}
+                  onChange={value => setJobConfig(value, 'config.process[0].caption.mtp')}
+                />
+                {(jobConfig.config.process[0].caption.mtp ?? true) && (
+                  <div className="pt-2">
+                    <SelectInput
+                      label="MTP Draft Depth"
+                      value={`${jobConfig.config.process[0].caption.mtp_depth ?? 3}`}
+                      onChange={value => {
+                        const intVal = parseInt(value);
+                        if (!isNaN(intVal)) {
+                          setJobConfig(intVal, 'config.process[0].caption.mtp_depth');
+                        }
+                      }}
+                      options={mtpDepthOptions}
+                    />
+                  </div>
+                )}
+              </>
+            )}
+            {additionalSections.includes('caption.cuda_graph') && (
+              <Checkbox
+                label="CUDA Graph Decode"
+                checked={jobConfig.config.process[0].caption.cuda_graph ?? true}
+                onChange={value => setJobConfig(value, 'config.process[0].caption.cuda_graph')}
+              />
+            )}
+            {additionalSections.includes('caption.vision_bf16') && (
+              <Checkbox
+                label="Vision Tower bf16"
+                checked={jobConfig.config.process[0].caption.vision_bf16 ?? true}
+                onChange={value => setJobConfig(value, 'config.process[0].caption.vision_bf16')}
+              />
+            )}
             {additionalSections.includes('caption.layer_offloading') && (
               <>
                 <Checkbox
@@ -346,6 +401,18 @@ const CaptionSimpleJob: React.FC<Props> = ({ jobConfig, setJobConfig, gpuIDs, se
               />
             </div>
           )}
+          {additionalSections.includes('caption.system_prompt') && (
+            <div className="mb-4">
+              <TextInput
+                label="System Prompt"
+                value={jobConfig.config.process[0].caption.system_prompt || ''}
+                onChange={value => {
+                  setJobConfig(value?.trim() === '' ? undefined : value, 'config.process[0].caption.system_prompt');
+                }}
+                placeholder="Optional system message"
+              />
+            </div>
+          )}
           <TextAreaInput
             label="Caption Prompt"
             value={jobConfig.config.process[0].caption.caption_prompt || ''}

ui/src/helpers/captionOptions.ts

@@ -1,7 +1,7 @@
 import { CloudLora, GroupedSelectOption, SelectOption } from "@/types";
 
-type CaptionGroup = 'image' | 'music' | 'video' | 'image/video/sound';
+type CaptionGroup = 'image' | 'music' | 'video' | 'image/video' | 'image/video/sound';
-type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.extract_vocals_before_transcribe' | 'caption.keep_timestamps' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading';
+type AdditionalSections = 'caption.model_name_or_path2' | 'caption.caption_prompt' | 'caption.max_res' | 'caption.max_new_tokens' | 'caption.fixed_caption' | 'caption.caption_format' | 'caption.extract_vocals_before_transcribe' | 'caption.keep_timestamps' | 'caption.thinking' | 'caption.batch_size' | 'caption.layer_offloading' | 'caption.mtp' | 'caption.cuda_graph' | 'caption.vision_bf16' | 'caption.video_fps' | 'caption.system_prompt';
 
 export interface CaptionOption {
     name: string;
@@ -26,6 +26,8 @@ export interface CaptionOption {
 
 const defaultNameOrPath = '';
 
+export const defaultQtype = 'convrot8';
+
 const extensionsAudio = ['mp3', 'wav', 'flac', 'ogg'];
 const extensionsImage = ['jpg', 'jpeg', 'png', 'bmp', 'webp'];
 const extensionsVideo = ['mp4', 'mov', 'webm', 'mkv', 'avi'];
@@ -247,6 +249,51 @@ export const captionerTypes: CaptionOption[] = [
             'caption.batch_size',
         ],
     },
+    {
+        name: 'Qwen38Captioner',
+        label: 'Qwen3.8 27B',
+        group: 'image/video',
+        supportsLoras: true,
+        defaults: {
+            'config.process[0].caption.loras': [[], undefined],
+            'config.process[0].caption.model_name_or_path': ['Comfy-Org/Qwen3.8-27B/text_encoders/qwen3.8_27b_w4a8.safetensors', defaultNameOrPath],
+            'config.process[0].caption.extensions': [[...extensionsVideo, ...extensionsImage], defaultExtensions],
+            'config.process[0].caption.caption_prompt': [defaultImageCaptionPrompt, undefined],
+            'config.process[0].caption.max_res': [1024, undefined],
+            'config.process[0].caption.max_new_tokens': [512, undefined],
+            'config.process[0].caption.mtp': [true, undefined],
+            'config.process[0].caption.mtp_depth': [3, undefined],
+            'config.process[0].caption.cuda_graph': [true, undefined],
+            'config.process[0].caption.vision_bf16': [true, undefined],
+            'config.process[0].caption.video_fps': [2, undefined],
+            'config.process[0].caption.thinking': [false, undefined],
+            // the ComfyUI file is already W4A8; quantize only applies to the HF repo
+            'config.process[0].caption.quantize': [true, true],
+            'config.process[0].caption.qtype': ['convrotw4a8', defaultQtype],
+            'config.process[0].caption.compile': [false, false],
+        },
+        name_or_path_options: [
+            { value: 'Comfy-Org/Qwen3.8-27B/text_encoders/qwen3.8_27b_w4a8.safetensors', label: 'Comfy-Org/Qwen3.8-27B (W4A8)' },
+            { value: 'Qwen/Qwen3.8-27B', label: 'Qwen/Qwen3.8-27B (bf16, quantized on load)' },
+        ],
+        captionPrompts: {
+            'General (image)': defaultImageCaptionPrompt,
+            'General (video)': defaultVideoCaptionPrompt,
+            'MiniMax H4 T2V': minimaxT2VCaptionPrompt,
+            'MiniMax H4 Image': minimaxImageCaptionPrompt,
+        },
+        additionalSections: [
+            'caption.caption_prompt',
+            'caption.system_prompt',
+            'caption.max_res',
+            'caption.max_new_tokens',
+            'caption.video_fps',
+            'caption.thinking',
+            'caption.mtp',
+            'caption.cuda_graph',
+            'caption.vision_bf16',
+        ],
+    },
     {
         name: 'Ideogram4Captioner',
         label: 'Ideogram 4 Captioner',
@@ -310,6 +357,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)' },
 ];
 
 // AceStepCaptioner output layouts (caption.caption_format)
@@ -347,4 +395,17 @@ export const maxNewTokensOptions: SelectOption[] = [
     { value: '8192', label: '8192' },
 ];
 
-export const defaultQtype = 'convrot8';
+export const videoFpsOptions: SelectOption[] = [
+    { value: '1', label: '1' },
+    { value: '2', label: '2 (default)' },
+    { value: '4', label: '4' },
+    { value: '8', label: '8' },
+];
+
+export const mtpDepthOptions: SelectOption[] = [
+    { value: '1', label: '1' },
+    { value: '2', label: '2' },
+    { value: '3', label: '3 (default)' },
+    { value: '4', label: '4' },
+    { value: '5', label: '5' },
+];

ui/src/types.ts

@@ -349,6 +349,13 @@ export interface CaptionProcessConfig {
     layer_offloading?: boolean;
     layer_offloading_percent?: number;
     loras?: CaptionLora[];
+    // Qwen3.8 captioner
+    mtp?: boolean;
+    mtp_depth?: number;
+    cuda_graph?: boolean;
+    vision_bf16?: boolean;
+    video_fps?: number;
+    system_prompt?: string;
   }
 }