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>
+    ),
+  },
 ];