From 391dda40c85df1a7bd4647e2650ed8c3aef2cf75 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:24:41 +0000 Subject: [PATCH 1/7] feat(chameleon): Thor SM110 dynamic-FP8 frontend Standalone Chameleon-7B (image+text) prefill/decode frontend for Jetson AGX Thor: - flash_rt/models/chameleon/pipeline_thor.py: 32-layer Chameleon forward with runtime dynamic per-tensor FP8 (fused quantize kernels), cuBLASLt per-shape autotune, selective L31 ffn_down clamp, optional AWQ V-proj and NVFP4 FFN tiers, CUDA-graph capture with re-embed before replay, and incremental KV-cache decode over fmha_fp16_causal_br. - Vendored Meta Chameleon VQ-GAN tokenizer (flash_rt/models/chameleon/vqgan, Meta Chameleon License headers retained; see the package docstring) with an eager default path and an opt-in TensorRT engine backend (hardware/thor/vqgan_trt_backend.py). - ChameleonTorchFrontendThor (frontends/torch/chameleon_thor.py): checkpoint_dir is a required argument with a clear error when missing; declarative weight spec in _chameleon_thor_spec.py. - hardware/thor/attn_backend_chameleon.py: CUTLASS causal FMHA backend with optional FA4 fast path, loading libfmha_fp16_causal.so from the package directory. --- .../frontends/torch/_chameleon_thor_spec.py | 92 ++ flash_rt/frontends/torch/chameleon_thor.py | 901 +++++++++++++ .../hardware/thor/attn_backend_chameleon.py | 362 ++++++ flash_rt/hardware/thor/vqgan_trt_backend.py | 187 +++ flash_rt/models/chameleon/__init__.py | 25 + flash_rt/models/chameleon/pipeline_thor.py | 1126 +++++++++++++++++ flash_rt/models/chameleon/vqgan/__init__.py | 2 + .../models/chameleon/vqgan/image_tokenizer.py | 132 ++ flash_rt/models/chameleon/vqgan/vocab.py | 107 ++ flash_rt/models/chameleon/vqgan/vqgan.py | 634 ++++++++++ 10 files changed, 3568 insertions(+) create mode 100644 flash_rt/frontends/torch/_chameleon_thor_spec.py create mode 100644 flash_rt/frontends/torch/chameleon_thor.py create mode 100644 flash_rt/hardware/thor/attn_backend_chameleon.py create mode 100644 flash_rt/hardware/thor/vqgan_trt_backend.py create mode 100644 flash_rt/models/chameleon/__init__.py create mode 100644 flash_rt/models/chameleon/pipeline_thor.py create mode 100644 flash_rt/models/chameleon/vqgan/__init__.py create mode 100644 flash_rt/models/chameleon/vqgan/image_tokenizer.py create mode 100644 flash_rt/models/chameleon/vqgan/vocab.py create mode 100644 flash_rt/models/chameleon/vqgan/vqgan.py diff --git a/flash_rt/frontends/torch/_chameleon_thor_spec.py b/flash_rt/frontends/torch/_chameleon_thor_spec.py new file mode 100644 index 00000000..832369e8 --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_thor_spec.py @@ -0,0 +1,92 @@ +"""Declarative weight spec for standalone Chameleon-7B on Thor. + +Standard Chameleon 32-layer backbone layout: +attention_bias=false, mlp_bias=false, per-head Q/K norm, SwiGLU FFN. +Per-head QK Norm prevents norm_fuse, so QKV is fused with ``Cat`` and the +gate/up pair with ``FusedGateUp``. +""" + +from __future__ import annotations + +from flash_rt.executors.weight_loader import Item, LayerBlock, ModelWeightSpec +from flash_rt.executors.torch_weights import ( + Attr, + Cat, + FusedGateUp, + Quant, + T, + TensorList, + ToFp16, +) + + +def _llm_block(*, use_fp8: bool = True) -> LayerBlock: + """Chameleon-7B LLM — 32 layers, MHA (num_kv_heads=32, no interleave). + + 4 quantized GEMMs per layer (qkv, o, gu, d) → 32 × 4 = 128 scales + appended to ``target._llm_w_scales``. + """ + qkv_tx = [T(), Quant()] if use_fp8 else [T()] + o_tx = [ToFp16(), T(), Quant()] if use_fp8 else [ToFp16(), T()] + gu_tx = [T(), Quant()] if use_fp8 else [T()] + d_tx = [ToFp16(), T(), Quant()] if use_fp8 else [ToFp16(), T()] + scale_into = "_llm_w_scales" if use_fp8 else None + + lp = "model.layers.{i}" + items = [ + # ── Fused QKV (no bias; per-head QK norm as separate items) ── + Item("qkv_w", + Cat([f"{lp}.self_attn.q_proj.weight", + f"{lp}.self_attn.k_proj.weight", + f"{lp}.self_attn.v_proj.weight"], dim=0), + qkv_tx, + TensorList("_llm_qkv_w"), scale_into=scale_into), + # ── O projection (no bias) ── + Item("o_w", f"{lp}.self_attn.o_proj.weight", + o_tx, + TensorList("_llm_o_w"), scale_into=scale_into), + # ── Fused GateUp ── + Item("gu_w", + FusedGateUp(gate=f"{lp}.mlp.gate_proj.weight", + up=f"{lp}.mlp.up_proj.weight"), + gu_tx, + TensorList("_llm_gu_w"), scale_into=scale_into), + # ── Down projection ── + Item("d_w", f"{lp}.mlp.down_proj.weight", + d_tx, + TensorList("_llm_d_w"), scale_into=scale_into), + # ── Layer norms ── + Item("input_ln_w", f"{lp}.input_layernorm.weight", + [ToFp16()], TensorList("_llm_input_ln_w")), + Item("post_ln_w", f"{lp}.post_attention_layernorm.weight", + [ToFp16()], TensorList("_llm_post_ln_w")), + # ── Per-head Q/K Norm ── + Item("q_norm_w", f"{lp}.self_attn.q_norm.weight", + [ToFp16()], TensorList("_llm_q_norm_w")), + Item("q_norm_b", f"{lp}.self_attn.q_norm.bias", + [ToFp16()], TensorList("_llm_q_norm_b")), + Item("k_norm_w", f"{lp}.self_attn.k_norm.weight", + [ToFp16()], TensorList("_llm_k_norm_w")), + Item("k_norm_b", f"{lp}.self_attn.k_norm.bias", + [ToFp16()], TensorList("_llm_k_norm_b")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="llm") + + +def build_spec(*, use_fp8: bool = True) -> ModelWeightSpec: + """Build the standalone Chameleon-7B Thor weight spec.""" + return ModelWeightSpec( + framework="torch", + blocks=[_llm_block(use_fp8=use_fp8)], + singletons=[ + Item("embed_w", "model.embed_tokens.weight", + [ToFp16()], Attr("_llm_embed_w")), + Item("norm_w", "model.norm.weight", + [ToFp16()], Attr("_llm_norm_w")), + Item("lm_head_w", "lm_head.weight", + [ToFp16()], Attr("_llm_lm_head_w")), + ], + ) + + +__all__ = ["build_spec"] diff --git a/flash_rt/frontends/torch/chameleon_thor.py b/flash_rt/frontends/torch/chameleon_thor.py new file mode 100644 index 00000000..2ff9ec9a --- /dev/null +++ b/flash_rt/frontends/torch/chameleon_thor.py @@ -0,0 +1,901 @@ +"""Standalone Chameleon-7B frontend for Jetson Thor (SM110). + +Direct-use VLM/LLM frontend: text + real images -> Chameleon prefill logits +and incremental KV-cache greedy generation. It uses the Chameleon +Thor dynamic-FP8 pipeline and causal FMHA attention backend. +""" + +from __future__ import annotations + +import ctypes +import json +import logging +import math +import os +import pathlib +from typing import List, Optional + +import numpy as np +import PIL +from PIL import Image as _PILImage +import torch + +import flash_rt.flash_rt_kernels as fvk +try: + import flash_rt.flash_rt_fp4 as fvk_fp4 +except Exception: + fvk_fp4 = None +from flash_rt.hardware.thor.attn_backend_chameleon import ( + ThorChameleonAttnBackend, + make_chameleon_attention_spec, +) +from flash_rt.models.chameleon.pipeline_thor import ( + chameleon_decode_step, + chameleon_forward, + chameleon_forward_calibrate, + chameleon_forward_fp16, +) + +logger = logging.getLogger(__name__) + +fp16 = torch.float16 +fp8 = torch.float8_e4m3fn +_cudart = ctypes.CDLL("libcudart.so") + +D_LLM = 4096 +NH_LLM = 32 +HD_LLM = 128 +L_LLM = 32 +DFF_LLM = 11008 +VOCAB_SIZE = 65536 +ROPE_THETA = 10000.0 + +# Chameleon image/text special ids from the shipped vocabulary_map. +PAD_ID = 1 +EOS_ID = 2 +IMG_START_ID = 8197 # +IMG_END_ID = 8196 # +NEWLINE_ID = 8803 +GRID_TOK_BASE = 8804 +PATCH_SIZE = 32 + + +class ChameleonTorchFrontendThor: + """Standalone Chameleon-7B Thor prefill frontend.""" + + def __init__( + self, + checkpoint_dir: str, + *, + use_fp8: bool = True, + use_cuda_graph: bool = True, + max_seq: int = 4096, + target_size: int = 512, + tokenizer_path: Optional[str] = None, + vqgan_path: Optional[str] = None, + use_trt_vqgan: bool = False, + trt_vqgan_engine_dir: Optional[str] = None, + use_autotune: bool = True, + ffn_clamp_layers: Optional[List[int]] = None, + fp4_ffn_layers: Optional[List[int]] = None, + use_fa4_attn: Optional[bool] = None, + ) -> None: + self.checkpoint_dir = pathlib.Path(checkpoint_dir).expanduser().resolve() + if not self.checkpoint_dir.exists(): + raise FileNotFoundError(f"checkpoint_dir not found: {self.checkpoint_dir}") + self._use_fp8 = bool(use_fp8) + self._use_cuda_graph = bool(use_cuda_graph) + self._max_pos = int(max_seq) + self.target_size = int(target_size) + self.tokenizer_path = tokenizer_path + self.vqgan_path = vqgan_path + self._use_trt_vqgan = bool(use_trt_vqgan) + self._trt_vqgan_engine_dir = trt_vqgan_engine_dir + self._trt_vqgan_backend = None + self._vqgan_backend = "eager" + self._trt_stream = torch.cuda.Stream() if self._use_trt_vqgan else None + self._infer_graph = None + self._captured_Se = None + self._last_input_ids: Optional[list[int]] = None + self.Se: Optional[int] = None + self._real_len: int = 0 + self._use_autotune = bool(use_autotune) + self._autotuned_se: set = set() + self._ffn_clamp_layers = self._resolve_ffn_clamp_layers(ffn_clamp_layers) + self._fp4_ffn_layers = self._resolve_fp4_ffn_layers(fp4_ffn_layers) + if use_fa4_attn is None: + use_fa4_attn = os.environ.get("FLASHRT_CHAMELEON_FA4_ATTN", "0") in ("1", "true", "on") + self._use_fa4_attn = bool(use_fa4_attn) + self._stream = torch.cuda.current_stream().cuda_stream + + with open(self.checkpoint_dir / "config.json") as f: + self.config_json = json.load(f) + self._validate_config() + self._build_image_token_mask() + with open(self.checkpoint_dir / "model.safetensors.index.json") as f: + self.weight_index = json.load(f) + + self._load_weights() + self._build_rope_tables() + self._load_tokenizer() + self._load_vqgan() + self._allocate_buffers() + + from flash_rt.core.context import FvkContext + self._ctx = FvkContext() + self._gemm = self._ctx.gemm + self._build_attention_backend() + if self._use_fa4_attn: + self._attn.set_fa4_attn(self._bufs["xn"], self._kv_cache) + self._logits_buf.zero_() + + logger.info( + "ChameleonTorchFrontendThor init: max_seq=%d target_size=%d fp8=%s graph=%s", + self._Se_max, self.target_size, self._use_fp8, self._use_cuda_graph, + ) + + def _resolve_ffn_clamp_layers(self, layers): + spec = os.environ.get("FLASHRT_CHAMELEON_FFN_CLAMP_LAYERS") \ + if layers is None else layers + if spec is None: + return frozenset({31}) + if isinstance(spec, str): + text = spec.strip().lower() + if text == "all": + return None + if text in ("", "none", "off", "false"): + return frozenset() + out = set() + for chunk in text.split(","): + chunk = chunk.strip() + if not chunk: + continue + if "-" in chunk: + a, b = chunk.split("-", 1) + out.update(range(int(a), int(b) + 1)) + else: + out.add(int(chunk)) + return frozenset(i for i in out if 0 <= i < L_LLM) + return frozenset(int(i) for i in spec if 0 <= int(i) < L_LLM) + + def _resolve_fp4_ffn_layers(self, layers): + spec = os.environ.get("FLASHRT_CHAMELEON_FP4_LAYERS") \ + if layers is None else layers + if spec is None: + return frozenset() + parsed = self._parse_layer_spec(spec) + return frozenset() if parsed is None else parsed + + def _parse_layer_spec(self, spec): + if isinstance(spec, str): + text = spec.strip().lower() + if text == "all": + return frozenset(range(L_LLM)) + if text in ("", "none", "off", "false"): + return frozenset() + out = set() + for chunk in text.split(","): + chunk = chunk.strip() + if not chunk: + continue + if "-" in chunk: + a, b = chunk.split("-", 1) + out.update(range(int(a), int(b) + 1)) + else: + out.add(int(chunk)) + return frozenset(i for i in out if 0 <= i < L_LLM) + return frozenset(int(i) for i in spec if 0 <= int(i) < L_LLM) + + def _validate_config(self) -> None: + c = self.config_json + for k, want in (("hidden_size", D_LLM), + ("num_attention_heads", NH_LLM), + ("num_hidden_layers", L_LLM), + ("intermediate_size", DFF_LLM), + ("vocab_size", VOCAB_SIZE)): + if int(c[k]) != want: + raise ValueError(f"config {k}={c[k]}, expected {want}") + if bool(c.get("attention_bias", False)): + raise ValueError("standard Chameleon Thor path expects attention_bias=false") + if bool(c.get("mlp_bias", False)): + raise ValueError("standard Chameleon Thor path expects mlp_bias=false") + + def _build_image_token_mask(self) -> None: + """Image-codebook vocab ids to suppress for text generation. + + Mirrors vendor ``ChameleonForConditionalGeneration``'s + ``mask_image_logits``: without this, greedy decode on a text + prompt can emit VQGAN codebook ids (garbage BPE decode) because + those ids are heavily represented in training and the raw + logit distribution favors them for many contexts. + """ + vocab_map = self.config_json.get("vocabulary_map", {}) + image_tokens = sorted( + v for k, v in vocab_map.items() if k.startswith("IMGIMG")) + self._mask_image_logits = bool(self.config_json.get("mask_image_logits", False)) + if image_tokens: + self._image_token_ids = torch.tensor( + image_tokens, dtype=torch.long, device="cuda") + else: + self._image_token_ids = None + self._mask_image_logits = False + + def _load_weights(self) -> None: + from flash_rt.executors.torch_weights import MultiSafetensorsSource, WeightLoader + from flash_rt.frontends.torch._chameleon_thor_spec import build_spec + + shard_paths = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + if not shard_paths: + raise FileNotFoundError(f"No safetensors shards in {self.checkpoint_dir}") + src = MultiSafetensorsSource([str(p) for p in shard_paths], device="cuda") + WeightLoader(source=src, target=self, spec=build_spec(use_fp8=self._use_fp8)).run() + + self._split_fused_llm_weights() + self._lm_head_w_t = self._llm_lm_head_w.t().contiguous() + if self._use_fp8: + self._setup_fp8_weight_scales() + self._init_fp4_weight_lists() + if self._use_fp8 and self._fp4_ffn_layers: + self._quantize_fp4_weights() + + def _split_fused_llm_weights(self) -> None: + self._q_w, self._k_w, self._v_w = [], [], [] + self._gate_w, self._up_w = [], [] + for li in range(L_LLM): + qkv = self._llm_qkv_w[li] + self._q_w.append(qkv[:, :D_LLM].contiguous()) + self._k_w.append(qkv[:, D_LLM:2 * D_LLM].contiguous()) + self._v_w.append(qkv[:, 2 * D_LLM:].contiguous()) + + gu = self._llm_gu_w[li] + self._gate_w.append(gu[:, :DFF_LLM].contiguous()) + self._up_w.append(gu[:, DFF_LLM:].contiguous()) + + self._llm_q_norm_w[li] = self._llm_q_norm_w[li].reshape(-1).contiguous() + self._llm_q_norm_b[li] = self._llm_q_norm_b[li].reshape(-1).contiguous() + self._llm_k_norm_w[li] = self._llm_k_norm_w[li].reshape(-1).contiguous() + self._llm_k_norm_b[li] = self._llm_k_norm_b[li].reshape(-1).contiguous() + + def _setup_fp8_weight_scales(self) -> None: + self._llm_w_dev = torch.tensor(self._llm_w_scales, dtype=torch.float32, device="cuda") + base = self._llm_w_dev.data_ptr() + self._d_w_qkv_ptrs = [(base + (li * 4 + 0) * 4) for li in range(L_LLM)] + self._d_w_o_ptrs = [(base + (li * 4 + 1) * 4) for li in range(L_LLM)] + self._d_w_gu_ptrs = [(base + (li * 4 + 2) * 4) for li in range(L_LLM)] + self._d_w_d_ptrs = [(base + (li * 4 + 3) * 4) for li in range(L_LLM)] + + def _init_fp4_weight_lists(self) -> None: + self._gu_w_fp4 = [None] * L_LLM + self._gu_sfb = [None] * L_LLM + self._d_w_fp4 = [None] * L_LLM + self._d_sfb = [None] * L_LLM + + def _quantize_fp4_weights(self) -> None: + if fvk_fp4 is None: + logger.warning("flash_rt_fp4 unavailable; FP4 FFN disabled") + self._fp4_ffn_layers = frozenset() + return + try: + if not fvk_fp4.has_nvfp4(): + logger.warning("NVFP4 unavailable on this device; FP4 FFN disabled") + self._fp4_ffn_layers = frozenset() + return + except AttributeError: + pass + + from flash_rt.executors.torch_weights import MultiSafetensorsSource + + shard_paths = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + src = MultiSafetensorsSource([str(p) for p in shard_paths], device="cuda") + + def _quant(w_fp16: torch.Tensor): + N, K = w_fp16.shape + packed = torch.empty(N, K // 2, dtype=torch.uint8, device="cuda") + sfb_bytes = fvk_fp4.sfa_size_bytes(N, K, True) + sfb = torch.empty(sfb_bytes, dtype=torch.uint8, device="cuda") + rc = fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + w_fp16.data_ptr(), packed.data_ptr(), sfb.data_ptr(), + N, K, True, 0) + if rc != 0: + raise RuntimeError(f"quantize_fp4_dynamic_sfa_fp16 failed rc={rc}") + return packed, sfb + + packed_layers = set() + for li in range(L_LLM): + if li not in self._fp4_ffn_layers: + continue + gate_key = f"model.layers.{li}.mlp.gate_proj.weight" + up_key = f"model.layers.{li}.mlp.up_proj.weight" + down_key = f"model.layers.{li}.mlp.down_proj.weight" + gate_w = src.get(gate_key).to(fp16).contiguous() + up_w = src.get(up_key).to(fp16).contiguous() + down_w = src.get(down_key).to(fp16).contiguous() + gu_w = torch.cat([gate_w, up_w], dim=0).contiguous() + self._gu_w_fp4[li], self._gu_sfb[li] = _quant(gu_w) + self._d_w_fp4[li], self._d_sfb[li] = _quant(down_w) + packed_layers.add(li) + del gate_w, up_w, down_w, gu_w + torch.cuda.synchronize() + self._fp4_ffn_layers = frozenset(packed_layers) + logger.info("FP4 FFN weights quantized for standalone Chameleon layers: %s", + sorted(self._fp4_ffn_layers)) + + def _build_rope_tables(self) -> None: + inv_freq = 1.0 / (ROPE_THETA ** ( + torch.arange(0, HD_LLM, 2, dtype=torch.float32) / HD_LLM)) + pos = torch.arange(self._max_pos, dtype=torch.float32) + freqs = torch.outer(pos, inv_freq) + emb = torch.cat([freqs, freqs], dim=-1) + self._rope_cos = emb.cos().to(fp16).cuda().contiguous() + self._rope_sin = emb.sin().to(fp16).cuda().contiguous() + + def _load_tokenizer(self) -> None: + from transformers import AutoTokenizer + + candidates: list[pathlib.Path] = [] + if self.tokenizer_path: + candidates.append(pathlib.Path(self.tokenizer_path)) + env = os.environ.get("FLASHRT_CHAMELEON_TOKENIZER_DIR") + if env: + candidates.append(pathlib.Path(env)) + candidates.append(self.checkpoint_dir) + + for p in candidates: + if p.exists() and (p / "tokenizer.json").exists(): + self.tokenizer = AutoTokenizer.from_pretrained(str(p), use_fast=True) + self.tokenizer_dir = p + logger.info("Tokenizer loaded from %s", p) + return + raise FileNotFoundError("Chameleon tokenizer not found") + + def _load_vqgan(self) -> None: + roots: list[pathlib.Path] = [] + if self.vqgan_path: + roots.append(pathlib.Path(self.vqgan_path)) + env = os.environ.get("FLASHRT_VQGAN_DIR") + if env: + roots.append(pathlib.Path(env)) + roots.extend([ + self.checkpoint_dir / "original_tokenizers", + self.checkpoint_dir, + ]) + + yaml_path = ckpt_path = None + for root in roots: + for d in (root, root / "original_tokenizers", root / "tokenizer"): + if (d / "vqgan.yaml").exists() and (d / "vqgan.ckpt").exists(): + yaml_path = d / "vqgan.yaml" + ckpt_path = d / "vqgan.ckpt" + break + if yaml_path is not None: + break + if yaml_path is None or ckpt_path is None: + raise FileNotFoundError("Chameleon VQGAN assets not found") + + # Chameleon-7B VQGAN reference implementation (Meta Chameleon + # license), vendored into the repo at flash_rt/models/chameleon/vqgan. + from flash_rt.models.chameleon import vqgan as chameleon_vae_ori # type: ignore + + self._vqgan_tokenizer = chameleon_vae_ori.ImageTokenizer( + cfg_path=str(yaml_path), ckpt_path=str(ckpt_path), device="cuda") + + text_tok = pathlib.Path(yaml_path).parent / "text_tokenizer.json" + if not text_tok.exists(): + text_tok = self.checkpoint_dir / "original_tokenizers" / "text_tokenizer.json" + with open(text_tok, encoding="utf8") as f: + vocab_json = json.load(f) + vocab_info = chameleon_vae_ori.VocabInfo(vocab_json["model"]["vocab"]) + self._vqgan_translation = chameleon_vae_ori.VocabTranslation(vocab_info, device="cuda") + logger.info("VQGAN loaded: yaml=%s ckpt=%s", yaml_path, ckpt_path) + + @property + def vqgan_backend(self) -> str: + return self._vqgan_backend + + @property + def fa4_attn_active(self) -> bool: + return bool(getattr(self._attn, "_fa4_enabled", False)) + + def _ensure_trt_vqgan_loaded(self) -> bool: + """Lazy-instantiate optional TensorRT VQGAN acceleration. + + Generic standalone Chameleon defaults to eager Chameleon VQGAN + tokenization. TensorRT is an explicit opt-in framework acceleration + path for deployments that have compatible engines under the FlashRT + engine cache (or a caller-provided engine directory). If unavailable, + the frontend falls back to eager tokenization. + """ + if not self._use_trt_vqgan: + return False + if self._trt_vqgan_backend is not None: + return self._trt_vqgan_backend.is_available() + try: + from flash_rt.hardware.thor.vqgan_trt_backend import VQGANTRTBackend + except ImportError as e: + logger.warning("TRT VQGAN import failed: %s", e) + self._use_trt_vqgan = False + self._vqgan_backend = "eager" + return False + if torch.backends.cudnn.allow_tf32: + logger.info("TRT VQGAN: disabling cuDNN TF32 globally.") + torch.backends.cudnn.allow_tf32 = False + torch.backends.cuda.matmul.allow_tf32 = False + engine_dir = pathlib.Path(self._trt_vqgan_engine_dir) \ + if self._trt_vqgan_engine_dir else None + self._trt_vqgan_backend = VQGANTRTBackend(engine_dir=engine_dir) + available = self._trt_vqgan_backend.is_available() + if not available: + logger.warning( + "TRT VQGAN backend unavailable at %s; falling back to " + "eager PyTorch encode.", + engine_dir or VQGANTRTBackend.ENGINE_DIR) + self._use_trt_vqgan = False + self._vqgan_backend = "eager" + else: + self._vqgan_backend = "trt" + return available + + def _preprocess_image_for_trt(self, pil_image, out_hw): + """PIL uint8 -> CUDA [1,3,H,W] float32 [-1,1] (vendor-matched).""" + H_out, W_out = out_hw + if pil_image.size != (W_out, H_out): + pil_image = pil_image.resize((W_out, H_out), resample=PIL.Image.BICUBIC) + np_img = np.array(pil_image.convert("RGB")) / 255.0 + np_img = np_img * 2.0 - 1.0 + t = torch.from_numpy(np_img).permute(2, 0, 1).unsqueeze(0) + return t.to(dtype=torch.float32, device="cuda").contiguous() + + def _vqgan_encode(self, image) -> list[int]: + if isinstance(image, np.ndarray): + image = _PILImage.fromarray(image.astype(np.uint8)).convert("RGB") + elif not isinstance(image, PIL.Image.Image): + raise TypeError(f"image must be PIL.Image or np.ndarray, got {type(image)!r}") + + # ── TRT fast path ── + if self._ensure_trt_vqgan_loaded(): + H_eng = W_eng = int(self.target_size) + if self._trt_vqgan_backend.supports_resolution(H_eng, W_eng): + with torch.cuda.stream(self._trt_stream): + img = self._preprocess_image_for_trt(image, (H_eng, W_eng)) + indices = self._trt_vqgan_backend.encode(img) + if indices is not None: + latent_ids = indices.view(-1) + global_ids = self._vqgan_translation.convert_img2bp2( + latent_ids).view(-1) + h_lat = H_eng // 16 + w_lat = W_eng // 16 + h_grids = H_eng // PATCH_SIZE + w_grids = W_eng // PATCH_SIZE + grid = global_ids.view(h_lat, w_lat) + newline_col = torch.full( + (h_lat, 1), NEWLINE_ID, dtype=grid.dtype, device=grid.device) + with_nl = torch.cat([grid, newline_col], dim=1).flatten().tolist() + return [IMG_START_ID, GRID_TOK_BASE + h_grids, GRID_TOK_BASE + w_grids, + *with_nl, IMG_END_ID] + + # ── Eager PyTorch fallback ── + # data_lerobot's aspect-preserving center crop when importable; + # otherwise a self-contained re-implementation of the same + # selection (max-coverage crop from the aspect-varied size list), + # so the eager path never depends on external packages. + try: + from data_lerobot.item_processor import ( # type: ignore + var_center_crop, generate_crop_size_list) + except ImportError: + def generate_crop_size_list(num_patches, patch_size, max_ratio=4.0): + crop_size_list = [] + wp, hp = num_patches, 1 + while wp > 0: + if max(wp, hp) / min(wp, hp) <= max_ratio: + crop_size_list.append((wp * patch_size, hp * patch_size)) + if (hp + 1) * wp <= num_patches: + hp += 1 + else: + wp -= 1 + return crop_size_list + + def var_center_crop(image, crop_size_list=None): + crop_size_list = crop_size_list or [(image.size[0], image.size[1])] + w, h = image.size + rem = [min(cw / w, ch / h) / max(cw / w, ch / h) + for cw, ch in crop_size_list] + best = sorted(zip(rem, crop_size_list), reverse=True)[0][1] + left = max(0, (w - best[0]) // 2) + top = max(0, (h - best[1]) // 2) + return image.crop((left, top, left + best[0], top + best[1])) + + crop_size_list = generate_crop_size_list( + (self.target_size // PATCH_SIZE) ** 2, PATCH_SIZE) + cropped = var_center_crop(image, crop_size_list=crop_size_list) + latent_ids = self._vqgan_tokenizer.img_tokens_from_pil(cropped) + global_ids = self._vqgan_translation.convert_img2bp2(latent_ids).view(-1) + + w_grids = cropped.size[0] // PATCH_SIZE + h_grids = cropped.size[1] // PATCH_SIZE + w_lat = cropped.size[0] // 16 + h_lat = cropped.size[1] // 16 + grid = global_ids.view(h_lat, w_lat) + newline_col = torch.full((h_lat, 1), NEWLINE_ID, dtype=grid.dtype, device=grid.device) + with_nl = torch.cat([grid, newline_col], dim=1).flatten().tolist() + return [IMG_START_ID, GRID_TOK_BASE + h_grids, GRID_TOK_BASE + w_grids, + *with_nl, IMG_END_ID] + + def _allocate_buffers(self) -> None: + Se = self._max_pos + D = D_LLM + Dff = DFF_LLM + self._Se_max = Se + self._bufs: dict[str, torch.Tensor] = {} + + def _alloc(name, shape, dtype=fp16): + self._bufs[name] = torch.zeros(shape, dtype=dtype, device="cuda") + + _alloc("x", (Se, D)) + _alloc("xn", (Se, D)) + _alloc("xn_fp8", (Se, D), fp8) + _alloc("o_proj_out", (Se, D)) + _alloc("hidden_all", (Se, D)) + _alloc("gate_out", (Se, Dff)) + _alloc("up_out", (Se, Dff)) + _alloc("gu_fp8", (Se, Dff), fp8) + _alloc("zero_bias_d", (D,)) + _alloc("zero_bias_dff", (Dff,)) + _alloc("act_fp4", (Se * D // 2,), torch.uint8) + _alloc("act_sfa", (Se * D // 16 * 2,), torch.uint8) + _alloc("ffn_act_fp4", (Se * Dff // 2,), torch.uint8) + _alloc("ffn_act_sfa", (Se * Dff // 16 * 2,), torch.uint8) + _alloc("gu_merged", (Se, 2 * Dff)) + _alloc("dyn_act_scales", (L_LLM * 4,), torch.float32) + _alloc("last_logits", (VOCAB_SIZE,)) + + self._llm_calib_scales = torch.ones(L_LLM * 4, dtype=torch.float32, device="cuda") + self._kv_cache = torch.zeros(L_LLM, 2, Se, D, dtype=fp16, device="cuda") + self._kv_layer_stride = 2 * Se * D * 2 + logits_sz = max(NH_LLM * Se * Se, 4) + self._logits_buf = torch.zeros(logits_sz, dtype=fp16, device="cuda") + + def _build_attention_backend(self) -> None: + spec = make_chameleon_attention_spec(seq_max=self._Se_max) + kv_base = self._kv_cache.data_ptr() + se_d_bytes = self._Se_max * D_LLM * 2 + chameleon_slots = { + "Q_O": self._bufs["xn"].data_ptr(), + "Kc": kv_base, + "Vc": kv_base + se_d_bytes, + "logits": self._logits_buf.data_ptr(), + "layer_stride": self._kv_layer_stride, + "scale": 1.0 / math.sqrt(HD_LLM), + } + self._attn = ThorChameleonAttnBackend( + spec, self._ctx, chameleon_slots=chameleon_slots) + + def _build_llm_weights(self) -> dict: + w = { + "input_ln_w": [x.data_ptr() for x in self._llm_input_ln_w], + "post_ln_w": [x.data_ptr() for x in self._llm_post_ln_w], + "q_w": [x.data_ptr() for x in self._q_w], + "k_w": [x.data_ptr() for x in self._k_w], + "v_w": [x.data_ptr() for x in self._v_w], + "o_w": [x.data_ptr() for x in self._llm_o_w], + "gate_w": [x.data_ptr() for x in self._gate_w], + "up_w": [x.data_ptr() for x in self._up_w], + "d_w": [x.data_ptr() for x in self._llm_d_w], + "q_norm_w": [x.data_ptr() for x in self._llm_q_norm_w], + "q_norm_b": [x.data_ptr() for x in self._llm_q_norm_b], + "k_norm_w": [x.data_ptr() for x in self._llm_k_norm_w], + "k_norm_b": [x.data_ptr() for x in self._llm_k_norm_b], + "o_b": [self._bufs["zero_bias_d"].data_ptr()] * L_LLM, + "final_norm_w": self._llm_norm_w.data_ptr(), + "rope_cos": self._rope_cos.data_ptr(), + "rope_sin": self._rope_sin.data_ptr(), + } + if self._use_fp8: + w.update({ + "w_scales_flat": self._llm_w_dev.data_ptr(), + "d_w_qkv": self._d_w_qkv_ptrs, + "d_w_o": self._d_w_o_ptrs, + "d_w_gu": self._d_w_gu_ptrs, + "d_w_d": self._d_w_d_ptrs, + "alpha_host": [1.0] * (L_LLM * 4), + "gu_w_fp4": [ + x.data_ptr() if x is not None else 0 for x in self._gu_w_fp4], + "gu_sfb": [ + x.data_ptr() if x is not None else 0 for x in self._gu_sfb], + "d_w_fp4": [ + x.data_ptr() if x is not None else 0 for x in self._d_w_fp4], + "d_sfb": [ + x.data_ptr() if x is not None else 0 for x in self._d_sfb], + }) + return w + + def _build_llm_bufs(self) -> dict: + return {k: self._bufs[k].data_ptr() for k in ( + "x", "xn", "xn_fp8", "o_proj_out", "hidden_all", + "gate_out", "up_out", "gu_fp8", "zero_bias_d", "zero_bias_dff", + "act_fp4", "act_sfa", "ffn_act_fp4", "ffn_act_sfa", "gu_merged", + "dyn_act_scales", + )} + + def _build_llm_scales_dev(self) -> dict: + calib_base = self._llm_calib_scales.data_ptr() + dyn_base = self._bufs["dyn_act_scales"].data_ptr() + return { + "act_qkv": [(calib_base + (li * 4 + 0) * 4) for li in range(L_LLM)], + "act_o": [(calib_base + (li * 4 + 1) * 4) for li in range(L_LLM)], + "act_gu": [(calib_base + (li * 4 + 2) * 4) for li in range(L_LLM)], + "act_down": [(calib_base + (li * 4 + 3) * 4) for li in range(L_LLM)], + "dyn_act_qkv": [(dyn_base + (li * 4 + 0) * 4) for li in range(L_LLM)], + "dyn_act_o": [(dyn_base + (li * 4 + 1) * 4) for li in range(L_LLM)], + "dyn_act_gu": [(dyn_base + (li * 4 + 2) * 4) for li in range(L_LLM)], + "dyn_act_down": [(dyn_base + (li * 4 + 3) * 4) for li in range(L_LLM)], + } + + def encode_prompt(self, text: str, images: Optional[list] = None) -> list[int]: + images = images or [] + chunks = text.split("") + if len(chunks) - 1 not in (0, len(images)): + raise ValueError("number of placeholders must be 0 or match images") + ids: list[int] = [] + bos = getattr(self.tokenizer, "bos_token_id", None) + if bos is not None: + ids.append(int(bos)) + for i, chunk in enumerate(chunks): + if chunk: + ids.extend(self.tokenizer.encode(chunk, add_special_tokens=False)) + if i < len(chunks) - 1: + ids.extend(self._vqgan_encode(images[i])) + if len(chunks) == 1: + for image in images: + ids.extend(self._vqgan_encode(image)) + return ids + + def _embed_ids(self, input_ids: list[int]) -> None: + Se = len(input_ids) + ids_t = torch.tensor(input_ids, dtype=torch.long, device="cuda") + emb = torch.nn.functional.embedding(ids_t, self._llm_embed_w) + self._bufs["x"].zero_() + _cudart.cudaMemcpyAsync( + ctypes.c_void_p(self._bufs["x"].data_ptr()), + ctypes.c_void_p(emb.data_ptr()), + Se * D_LLM * 2, + 3, + ctypes.c_void_p(self._stream), + ) + + def set_prompt(self, text: str, images: Optional[list] = None) -> list[int]: + input_ids = self.encode_prompt(text, images) + if len(input_ids) > self._Se_max: + raise ValueError(f"sequence length {len(input_ids)} exceeds max_seq={self._Se_max}") + self._real_len = len(input_ids) + rem = len(input_ids) % 16 + if rem: + input_ids.extend([PAD_ID] * (16 - rem)) + self.Se = len(input_ids) + self._last_input_ids = input_ids + if self._use_autotune: + self._autotune_gemms(self.Se) + self._embed_ids(input_ids) + if self._use_cuda_graph: + self._capture_graph(self.Se) + return input_ids + + def _autotune_gemms(self, Se: int, num_algos: int = 16) -> None: + """Per-shape cuBLASLt algo autotune (motus/hyvla pattern). + + Chameleon's per-layer GEMMs collapse to 3 distinct (M,N,K) shapes + at a given Se (q/k/v/o share one shape, gate/up share another, + down is the third) plus the M=1 lm_head projection. Tuning each + shape once mutates the GemmRunner's internal algo cache (keyed on + (M,N,K)); every subsequent real call with that shape — including + inside a captured CUDA graph — picks up the tuned algo for free. + Dummy buffers are only used for timing, not correctness. + """ + if Se in self._autotuned_se: + return + D, Dff = D_LLM, DFF_LLM + dev = "cuda" + if self._use_fp8 and hasattr(self._gemm, "autotune_fp8_nn_dev_fp16"): + shapes = [(Se, D, D), (Se, Dff, D), (Se, D, Dff)] + shapes += [(1, D, D), (1, Dff, D), (1, D, Dff)] # decode + for (M, N, K) in dict.fromkeys(shapes): + A = torch.empty(M, K, dtype=torch.uint8, device=dev) + B = torch.empty(K, N, dtype=torch.uint8, device=dev) + Dbuf = torch.empty(M, N, dtype=fp16, device=dev) + sa = torch.ones(1, dtype=torch.float32, device=dev) + sb = torch.ones(1, dtype=torch.float32, device=dev) + self._gemm.autotune_fp8_nn_dev_fp16( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), + M, N, K, sa.data_ptr(), sb.data_ptr(), num_algos) + elif hasattr(self._gemm, "autotune_fp16_nn"): + shapes = [(Se, D, D), (Se, Dff, D), (Se, D, Dff)] + shapes += [(1, D, D), (1, Dff, D), (1, D, Dff)] # decode + for (M, N, K) in dict.fromkeys(shapes): + A = torch.empty(M, K, dtype=fp16, device=dev) + B = torch.empty(K, N, dtype=fp16, device=dev) + Dbuf = torch.empty(M, N, dtype=fp16, device=dev) + self._gemm.autotune_fp16_nn( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), M, N, K, num_algos) + if hasattr(self._gemm, "autotune_fp16_nn"): + A = torch.empty(1, D, dtype=fp16, device=dev) + B = torch.empty(D, VOCAB_SIZE, dtype=fp16, device=dev) + Dbuf = torch.empty(1, VOCAB_SIZE, dtype=fp16, device=dev) + self._gemm.autotune_fp16_nn( + A.data_ptr(), B.data_ptr(), Dbuf.data_ptr(), 1, VOCAB_SIZE, D, num_algos) + torch.cuda.synchronize() + self._autotuned_se.add(Se) + + def _capture_graph(self, Se: int) -> None: + if self._infer_graph is not None and self._captured_Se == Se: + return + capture_stream = torch.cuda.Stream() + prev_stream_id = self._stream + with torch.cuda.stream(capture_stream): + self._stream = capture_stream.cuda_stream + self._run_backbone(Se) + torch.cuda.synchronize() + # The warmup forward mutates x into the final residual stream; the + # capture pass (and every later replay) must start from clean + # embeddings — see _replay_backbone. + self._embed_ids(self._last_input_ids) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=capture_stream): + self._run_backbone(Se) + self._stream = prev_stream_id + self._infer_graph = graph + self._captured_Se = Se + + def _run_backbone(self, Se: int) -> None: + dims = {"Se": Se, "D": D_LLM, "Dff": DFF_LLM, + "L": L_LLM, "H": NH_LLM, "Hd": HD_LLM} + if self._use_fp8: + chameleon_forward( + self._gemm, fvk, self._build_llm_bufs(), self._build_llm_weights(), + dims, self._build_llm_scales_dev(), attn=self._attn, + stream=self._stream, dynamic_fp8_layers=frozenset(range(L_LLM)), + fp4_ffn_layers=self._fp4_ffn_layers, ffn_down_clamp_value=60000.0, + ffn_clamp_layers=self._ffn_clamp_layers, + ) + else: + chameleon_forward_fp16( + self._gemm, fvk, self._build_llm_bufs(), self._build_llm_weights(), + dims, attn=self._attn, stream=self._stream, + ) + + def _project_last(self, last_idx: Optional[int] = None) -> None: + if last_idx is None: + last_idx = self._real_len - 1 + last_hidden_ptr = self._bufs["hidden_all"].data_ptr() + last_idx * D_LLM * 2 + self._gemm.fp16_nn( + last_hidden_ptr, self._lm_head_w_t.data_ptr(), + self._bufs["last_logits"].data_ptr(), 1, VOCAB_SIZE, D_LLM, + int(self._stream), + ) + if self._mask_image_logits and self._image_token_ids is not None: + self._bufs["last_logits"].index_fill_(0, self._image_token_ids, -65504.0) + + def _run_forward(self, Se: int) -> None: + self._run_backbone(Se) + self._project_last() + + def _replay_backbone(self) -> None: + """Replay the captured prefill graph over clean embeddings. + + Every backbone run mutates ``x`` in place into the final residual + stream, so a replay must re-embed first or it recomputes over the + previous run's residuals. + """ + self._embed_ids(self._last_input_ids) + self._infer_graph.replay() + + def prefill(self, text: str, images: Optional[list] = None) -> dict: + input_ids = self.set_prompt(text, images) + if self._use_cuda_graph and self._infer_graph is not None: + self._replay_backbone() + self._project_last() + else: + self._run_forward(self.Se) + torch.cuda.synchronize() + return { + "input_ids": input_ids, + "Se": self.Se, + "vqgan_backend": self._vqgan_backend, + "fa4_attn": self.fa4_attn_active, + "logits": self._bufs["last_logits"].detach().float().cpu(), + "hidden": self._bufs["hidden_all"][:self.Se].detach().float().cpu(), + } + + def generate_greedy(self, text: str, images: Optional[list] = None, + max_new_tokens: int = 16, + eos_token_id: Optional[int] = None) -> dict: + """Greedy generation with incremental KV-cache decode. + + One prefill over the prompt, then single-token decode steps + (``chameleon_decode_step``) — O(n) instead of the historical O(n^2) + full recompute. Decode runs eagerly (no CUDA graph) because the + position is a host scalar baked into RoPE offsets and cache rows. + """ + if not self._use_fp8: + raise NotImplementedError( + "incremental decode requires the dynamic-FP8 path " + "(use_fp8=True)") + if self._fp4_ffn_layers: + raise NotImplementedError( + "incremental decode does not model the NVFP4 FFN tiers " + "(fp4_ffn_layers); use the prefill-only path") + + self.set_prompt(text, images) + generated = list(self._last_input_ids[:self._real_len]) + eos = EOS_ID if eos_token_id is None else int(eos_token_id) + budget = int(max_new_tokens) + if budget <= 0: + return {"input_ids": generated, + "text": self.tokenizer.decode(generated)} + + # Prefill (graph or eager) + first token from the last prompt row. + if self._use_cuda_graph and self._infer_graph is not None: + self._replay_backbone() + else: + self._run_backbone(self.Se) + self._project_last() + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + + # Decode-step plumbing: all pointers are stable across steps, so + # build the dicts once; only `pos` changes per token. + dims = {"Se": 1, "D": D_LLM, "Dff": DFF_LLM, + "L": L_LLM, "H": NH_LLM, "Hd": HD_LLM} + bufs = self._build_llm_bufs() + weights = self._build_llm_weights() + scales_dev = self._build_llm_scales_dev() + pos = self._real_len + for _ in range(budget - 1): + if next_id == eos or pos >= self._Se_max: + break + # Decode state: single-token embedding in residual row 0. + self._bufs["x"][:1].copy_(self._llm_embed_w[next_id]) + chameleon_decode_step( + self._gemm, fvk, bufs, weights, dims, scales_dev, + attn=self._attn, pos=pos, stream=int(self._stream), + ffn_down_clamp_value=60000.0, + ffn_clamp_layers=self._ffn_clamp_layers, + ) + self._project_last(last_idx=0) + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + pos += 1 + return {"input_ids": generated, "text": self.tokenizer.decode(generated)} + + def _generate_greedy_recompute(self, text: str, + images: Optional[list] = None, + max_new_tokens: int = 16) -> dict: + """Legacy O(n^2) full-recompute greedy path (oracle for decode). + + Always eager: it shares the KV cache with the incremental path, + and graph capture/replay of every growing-Se forward would + overwrite cache rows for pad-16 filler positions. + """ + ids = self.encode_prompt(text, images) + generated = list(ids) + for _ in range(int(max_new_tokens)): + if len(generated) >= self._Se_max: + break + self._real_len = len(generated) + padded = list(generated) + rem = len(padded) % 16 + if rem: + padded.extend([PAD_ID] * (16 - rem)) + self.Se = len(padded) + self._last_input_ids = padded + if self._use_autotune: + self._autotune_gemms(self.Se) + self._embed_ids(padded) + self._run_forward(self.Se) + torch.cuda.synchronize() + next_id = int(torch.argmax(self._bufs["last_logits"]).item()) + generated.append(next_id) + return {"input_ids": generated, "text": self.tokenizer.decode(generated)} + + +__all__ = ["ChameleonTorchFrontendThor"] diff --git a/flash_rt/hardware/thor/attn_backend_chameleon.py b/flash_rt/hardware/thor/attn_backend_chameleon.py new file mode 100644 index 00000000..f1484158 --- /dev/null +++ b/flash_rt/hardware/thor/attn_backend_chameleon.py @@ -0,0 +1,362 @@ +"""FlashRT — Thor SM110 attention backend for standalone Chameleon-7B. + +One site: + +* **chameleon** — 32-layer Chameleon-7B LLM self-attention (NH=32, HD=128, + causal). Q/O aliased into the frontend's xn buffer. Dispatch order: + FA4 (FlashAttention-4, CuTe-DSL, optional fast path) → CUTLASS causal + FMHA (``libfmha_fp16_causal.so``) → cuBLAS decomposed causal MHA + (``attention_mha_causal_fp16``). + +Prefill (``run``) uses the top-left causal alignment (SQ == SK); single-query +incremental decode (``run_decode``) uses the bottom-right alignment via +``fmha_fp16_causal_br`` (or FA4 / plain cuBLAS MHA, both bottom-right +equivalent at SQ == 1). +""" + +from __future__ import annotations + +import ctypes +import logging +import pathlib + +from flash_rt.hardware.backend import AttentionBackendBase, AttentionSpec + +logger = logging.getLogger(__name__) + +# ── CUTLASS causal FMHA (SM100/110) dynamic loading ── +_fmha_causal_fn = None +_fmha_causal_br_fn = None # bottom-right aligned variant (decode, SQ < SK) + + +def _load_fmha_causal_library() -> bool: + """Load libfmha_fp16_causal.so and resolve the fmha_fp16_causal symbol.""" + global _fmha_causal_fn, _fmha_causal_br_fn + if _fmha_causal_fn is not None: + return True + search_paths = [ + pathlib.Path(__file__).parent.parent.parent / "libfmha_fp16_causal.so", + pathlib.Path(__file__).parent.parent.parent.parent / "build" / "libfmha_fp16_causal.so", + ] + argtypes = [ + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_void_p, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_int, ctypes.c_int, ctypes.c_int, + ctypes.c_void_p, + ] + for p in search_paths: + if p.exists(): + try: + lib = ctypes.CDLL(str(p)) + fn = lib.fmha_fp16_causal + fn.restype = ctypes.c_int + fn.argtypes = argtypes + _fmha_causal_fn = fn + # Bottom-right aligned causal (decode). Older .so builds may + # lack it; run_decode then falls through to the cuBLAS tier. + try: + fn_br = lib.fmha_fp16_causal_br + fn_br.restype = ctypes.c_int + fn_br.argtypes = argtypes + _fmha_causal_br_fn = fn_br + except AttributeError: + logger.warning( + "libfmha_fp16_causal.so lacks fmha_fp16_causal_br " + "(rebuild for CUTLASS decode path)") + logger.info("CUTLASS causal FMHA loaded from %s", p) + return True + except OSError as e: + logger.warning("Failed to load causal FMHA from %s: %s", p, e) + logger.warning("CUTLASS causal FMHA not found — will fall back to cuBLAS MHA") + return False + + +class ThorChameleonAttnBackend(AttentionBackendBase): + """Standalone Chameleon-7B attention backend on Thor (SM110). + + Single ``chameleon`` site: 32-layer causal MHA, Q/O aliased, per-layer + KV cache with ``layer_stride`` bytes between consecutive layers. + """ + + def __init__(self, spec: AttentionSpec, ctx, *, chameleon_slots: dict) -> None: + super().__init__(spec) + + expected_sites = {"chameleon"} + got = set(spec.sites.keys()) + if got != expected_sites: + raise ValueError( + f"ThorChameleonAttnBackend expects sites {expected_sites}, " + f"got {got}") + + self._ctx_cpp = ctx.cpp if hasattr(ctx, "cpp") else ctx + self._slots = {"chameleon": dict(chameleon_slots)} + self._require_keys("chameleon", + ("Q_O", "Kc", "Vc", "logits", "layer_stride", "scale")) + + self._per_layer_kv = {} + s = self._slots["chameleon"] + nL = spec.site("chameleon").num_layers + stride = int(s["layer_stride"]) + Kc = int(s["Kc"]) + Vc = int(s["Vc"]) + self._per_layer_kv["chameleon"] = [ + (Kc + l * stride, Vc + l * stride) for l in range(nL) + ] + + self._fvk = None + self._has_causal_fmha = _load_fmha_causal_library() + + # FA4 (FlashAttention-4, CuTe-DSL) fast path state. Populated by the + # frontend with torch-tensor references to the live Q (xn) and KV + # cache buffers, so run() can slice views (metadata-only, capture + # safe) and dispatch to FA4 instead of the CUTLASS FMHA kernel. + self._fa4_enabled = False + self._fa4_q_tensor = None + self._fa4_kv_cache = None + + def set_fa4_attn(self, q_tensor, kv_cache) -> None: + """Enable the FA4 (FlashAttention-4, CuTe-DSL) fast path. + + q_tensor must alias the chameleon Q_O slot buffer (the frontend's xn + buffer, which Q GEMMs write into and O is read from). kv_cache must + alias the per-layer K/V buffer with layers at [li, 0, :Se] (K) and + [li, 1, :Se] (V). Both are torch tensors so run() can build + capture-safe views. + """ + from flash_rt.hardware.thor import fa4_backend + if not fa4_backend.is_available(): + logger.warning("FA4 unavailable (%s); keeping CUTLASS FMHA", + fa4_backend.status()) + return + if q_tensor is None or kv_cache is None: + raise ValueError("FA4 requires the Q (xn) and KV cache tensor refs") + self._fa4_q_tensor = q_tensor + self._fa4_kv_cache = kv_cache + self._fa4_enabled = True + logger.info("FA4 causal FMHA enabled for chameleon site") + + def disable_fa4_attn(self) -> None: + """Revert chameleon site to the CUTLASS FMHA path.""" + self._fa4_enabled = False + + def _require_keys(self, site, keys): + slot = self._slots[site] + for k in keys: + if k not in slot: + raise ValueError(f"{site}_slots missing required key {k!r}") + if k in ("Q_O", "Kc", "Vc", "logits"): + if int(slot[k]) == 0: + raise ValueError( + f"{site}_slots[{k!r}] is a null device pointer") + + def _fvk_mod(self): + if self._fvk is None: + import flash_rt.flash_rt_kernels as fvk + self._fvk = fvk + return self._fvk + + def get_slot_ptrs(self, site, layer_idx): + """Return {Q, K, V, O} device pointer ints for (site, layer_idx).""" + if site not in self._slots: + raise KeyError(f"unknown site {site!r}") + nL = self._spec.site(site).num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + q_o = int(self._slots[site]["Q_O"]) + return {"Q": q_o, "K": K_ptr, "V": V_ptr, "O": q_o} + + def kv_row_ptrs(self, site, layer_idx, pos): + """Return (K_row, V_row) device pointer ints for token position pos. + + Rows are fp16 [NH*HD] within the layer's cache segment, so decode + can GEMM the new token's K/V directly into the cache. + """ + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + spec = self._spec.site(site) + if not (0 <= layer_idx < spec.num_layers): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if not (0 <= pos < spec.max_kv_seq): + raise IndexError(f"pos {pos} out of range for site {site!r}") + row_bytes = spec.num_kv_heads * spec.head_dim * 2 # fp16 + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + return K_ptr + pos * row_bytes, V_ptr + pos * row_bytes + + def run(self, site, layer_idx, q_seq, *, kv_seq=None, stream=0, + state_nk=None, cross_attn=False): + """Dispatch causal attention for (site, layer_idx).""" + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + fvk = self._fvk_mod() + site_spec = self._spec.site(site) + nL = site_spec.num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if kv_seq is None: + kv_seq = q_seq + if q_seq == 1 and kv_seq > q_seq: + raise ValueError( + "run() got a decode shape (q_seq=1 < kv_seq); the top-left " + "causal mask misaligns here — call run_decode() instead") + + s = self._slots[site] + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + NH = site_spec.num_q_heads + HD = site_spec.head_dim + + # ── FA4 (FlashAttention-4, CuTe-DSL) fast path ── + if self._fa4_enabled and self._fa4_q_tensor is not None: + try: + from flash_rt.hardware.thor import fa4_backend + fa4 = fa4_backend.fa4_func() + if fa4 is not None: + import torch + qv = self._fa4_q_tensor[:q_seq].view(1, q_seq, NH, HD) + kv_cache = self._fa4_kv_cache + k_view = kv_cache[layer_idx, 0, :kv_seq].view(1, kv_seq, NH, HD) + v_view = kv_cache[layer_idx, 1, :kv_seq].view(1, kv_seq, NH, HD) + with torch.no_grad(): + o = fa4(qv.contiguous(), k_view.contiguous(), + v_view.contiguous(), causal=True, pack_gqa=True) + if isinstance(o, tuple): + o = o[0] + o_flat = o.reshape(q_seq, NH * HD) + # Q_O slot aliases xn (self._fa4_q_tensor): overwrite + # Q with O in place via a metadata-only view. + self._fa4_q_tensor[:q_seq].view(q_seq, NH * HD).copy_(o_flat) + return int(s["Q_O"]) + except Exception as e: # pragma: no cover - defensive fallback + logger.warning("FA4 failed at layer %d (%s); falling back to CUTLASS", + layer_idx, e) + self._fa4_enabled = False + + # ── CUTLASS Causal FMHA (SM100/110) — preferred ── + if self._has_causal_fmha and _fmha_causal_fn is not None: + ret = _fmha_causal_fn( + ctypes.c_void_p(int(s["Q_O"])), + ctypes.c_void_p(K_ptr), + ctypes.c_void_p(V_ptr), + ctypes.c_void_p(int(s["Q_O"])), + 1, q_seq, kv_seq, NH, NH, HD, + ctypes.c_void_p(stream), + ) + if ret != 0: + logger.warning( + "CUTLASS causal FMHA returned %d, falling back to cuBLAS", + ret) + fvk.attention_mha_causal_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + q_seq, kv_seq, NH, HD, + float(s["scale"]), stream, + ) + else: + # ── Fallback: cuBLAS decomposed causal MHA ── + fvk.attention_mha_causal_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + q_seq, kv_seq, NH, HD, + float(s["scale"]), stream, + ) + return int(s["Q_O"]) + + def run_decode(self, site, layer_idx, kv_len, *, stream=0): + """Dispatch single-query (decode) attention for (site, layer_idx). + + Q is the single row in the Q_O slot; K/V are the first kv_len rows + of the layer's cache segment. The causal mask must be bottom-right + aligned (query at position kv_len-1 attends all kv_len keys). + """ + if site != "chameleon": + raise KeyError(f"unknown site {site!r}") + fvk = self._fvk_mod() + site_spec = self._spec.site(site) + nL = site_spec.num_layers + if not (0 <= layer_idx < nL): + raise IndexError( + f"layer_idx {layer_idx} out of range for site {site!r}") + if not (1 <= kv_len <= site_spec.max_kv_seq): + raise IndexError(f"kv_len {kv_len} out of range for site {site!r}") + + s = self._slots[site] + K_ptr, V_ptr = self._per_layer_kv[site][layer_idx] + NH = site_spec.num_q_heads + HD = site_spec.head_dim + + # ── FA4 fast path: causal is bottom-right (offset_k = sk - sq) ── + if self._fa4_enabled and self._fa4_q_tensor is not None: + try: + from flash_rt.hardware.thor import fa4_backend + fa4 = fa4_backend.fa4_func() + if fa4 is not None: + import torch + qv = self._fa4_q_tensor[:1].view(1, 1, NH, HD) + kv_cache = self._fa4_kv_cache + k_view = kv_cache[layer_idx, 0, :kv_len].view(1, kv_len, NH, HD) + v_view = kv_cache[layer_idx, 1, :kv_len].view(1, kv_len, NH, HD) + with torch.no_grad(): + o = fa4(qv.contiguous(), k_view.contiguous(), + v_view.contiguous(), causal=True, pack_gqa=True) + if isinstance(o, tuple): + o = o[0] + self._fa4_q_tensor[:1].view(1, NH * HD).copy_( + o.reshape(1, NH * HD)) + return int(s["Q_O"]) + except Exception as e: # pragma: no cover - defensive fallback + logger.warning("FA4 decode failed at layer %d (%s); falling back", + layer_idx, e) + self._fa4_enabled = False + + # ── CUTLASS bottom-right causal FMHA ── + if self._has_causal_fmha and _fmha_causal_br_fn is not None: + ret = _fmha_causal_br_fn( + ctypes.c_void_p(int(s["Q_O"])), + ctypes.c_void_p(K_ptr), + ctypes.c_void_p(V_ptr), + ctypes.c_void_p(int(s["Q_O"])), + 1, 1, kv_len, NH, NH, HD, + ctypes.c_void_p(stream), + ) + if ret != 0: + logger.warning( + "CUTLASS causal FMHA (br) returned %d at layer %d, " + "falling back to cuBLAS", ret, layer_idx) + fvk.attention_mha_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + 1, kv_len, NH, HD, + float(s["scale"]), stream, + ) + else: + # ── cuBLAS non-causal MHA: with q_seq==1 the bottom-right + # causal mask is the identity, so this is equivalent. ── + fvk.attention_mha_fp16( + self._ctx_cpp, + int(s["Q_O"]), K_ptr, V_ptr, + int(s["logits"]), int(s["Q_O"]), + 1, kv_len, NH, HD, + float(s["scale"]), stream, + ) + return int(s["Q_O"]) + + +def make_chameleon_attention_spec(*, seq_max: int) -> AttentionSpec: + """Build the standalone Chameleon-7B Thor AttentionSpec (1 site).""" + spec = AttentionSpec() + spec.add_site( + "chameleon", + num_layers=32, num_q_heads=32, num_kv_heads=32, head_dim=128, + max_q_seq=int(seq_max), max_kv_seq=int(seq_max), + causal=True, + ) + return spec + + +__all__ = ["ThorChameleonAttnBackend", "make_chameleon_attention_spec"] diff --git a/flash_rt/hardware/thor/vqgan_trt_backend.py b/flash_rt/hardware/thor/vqgan_trt_backend.py new file mode 100644 index 00000000..344610c9 --- /dev/null +++ b/flash_rt/hardware/thor/vqgan_trt_backend.py @@ -0,0 +1,187 @@ +"""TensorRT VQ-GAN encoder backend for Jetson Thor. + +Manages multiple fixed-shape TRT engines (one per resolution). +Lazily loads/deserializes engines on first use of each resolution. +Falls back gracefully if engines are unavailable. +""" + +import json +import logging +from pathlib import Path +from typing import Optional + +import torch + +logger = logging.getLogger(__name__) + +_TRT_AVAILABLE = False +try: + import tensorrt as trt + _TRT_AVAILABLE = True +except ImportError: + pass + + +class VQGANTRTBackend: + ENGINE_DIR = Path.home() / ".flash_rt" / "trt_engines" / "vqgan" + + def __init__(self, engine_dir: Optional[Path] = None): + self._engine_dir = Path(engine_dir) if engine_dir else self.ENGINE_DIR + self._manifest = None + self._engines = {} + self._contexts = {} + self._buffers = {} + self._available = None + self._trt_logger = None + + def is_available(self) -> bool: + if self._available is not None: + return self._available + if not _TRT_AVAILABLE: + logger.warning("tensorrt not importable; TRT VQGAN disabled") + self._available = False + return False + manifest_path = self._engine_dir / "manifest.json" + if not manifest_path.exists(): + logger.warning("No TRT VQGAN manifest at %s", manifest_path) + self._available = False + return False + with open(manifest_path) as f: + self._manifest = json.load(f) + has_engines = len(self._manifest.get("engines", {})) > 0 + if has_engines: + logger.info("TRT VQGAN backend: %d engines at %s (TRT %s)", + len(self._manifest["engines"]), self._engine_dir, + self._manifest.get("trt_version", "?")) + self._available = has_engines + return has_engines + + def supports_resolution(self, height: int, width: int) -> bool: + if not self._available: + return False + key = f"{height}x{width}" + return key in self._manifest.get("engines", {}) + + def encode(self, image_tensor: torch.Tensor) -> Optional[torch.Tensor]: + """Run TRT VQGAN encoder. + + Args: + image_tensor: [1, 3, H, W] float32 on CUDA, range [-1, 1] + + Returns: + [1, H//16, W//16] int64 codebook indices on CUDA, or None on failure. + """ + _, _, h, w = image_tensor.shape + key = f"{h}x{w}" + + if key not in self._contexts: + if not self._load_engine(h, w): + return None + + ctx = self._contexts[key] + inp_buf, out_buf, out_dtype_is_int32 = self._buffers[key] + + inp_buf.copy_(image_tensor) + ctx.set_tensor_address("image", inp_buf.data_ptr()) + ctx.set_tensor_address("indices", out_buf.data_ptr()) + + stream = torch.cuda.current_stream() + ok = ctx.execute_async_v3(stream_handle=stream.cuda_stream) + if not ok: + logger.error("TRT execute_async_v3 failed for resolution %s", key) + return None + + if out_dtype_is_int32: + return out_buf.to(torch.int64) + return out_buf.clone() + + def encode_batch(self, image_tensor: torch.Tensor) -> Optional[torch.Tensor]: + """Run TRT VQGAN encoder on a multi-view batch. + + Args: + image_tensor: [B, 3, H, W] float32 on CUDA, range [-1, 1] + where B == engine batch (from per-engine manifest entry). + + Returns: + [B, H//16, W//16] int64 codebook indices, or None on failure / + batch mismatch. + """ + b, _, h, w = image_tensor.shape + key = f"{h}x{w}" + + if key not in self._contexts: + if not self._load_engine(h, w): + return None + + engine_batch = self._manifest["engines"][key]["input_shape"][0] + if b != engine_batch: + return None # caller should fall back to per-view encode + + ctx = self._contexts[key] + inp_buf, out_buf, out_dtype_is_int32 = self._buffers[key] + + inp_buf.copy_(image_tensor) + ctx.set_tensor_address("image", inp_buf.data_ptr()) + ctx.set_tensor_address("indices", out_buf.data_ptr()) + + stream = torch.cuda.current_stream() + ok = ctx.execute_async_v3(stream_handle=stream.cuda_stream) + if not ok: + logger.error("TRT execute_async_v3 (batch=%d) failed at %s", b, key) + return None + + if out_dtype_is_int32: + return out_buf.to(torch.int64) + return out_buf.clone() + + def _load_engine(self, height: int, width: int) -> bool: + key = f"{height}x{width}" + meta = self._manifest["engines"].get(key) + if meta is None: + return False + + engine_path = self._engine_dir / meta["file"] + if not engine_path.exists(): + logger.error("Engine file missing: %s", engine_path) + return False + + if self._trt_logger is None: + self._trt_logger = trt.Logger(trt.Logger.WARNING) + if hasattr(trt, "init_libnvinfer_plugins"): + trt.init_libnvinfer_plugins(self._trt_logger, "") + + runtime = trt.Runtime(self._trt_logger) + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + + if engine is None: + logger.error("Failed to deserialize TRT engine: %s", engine_path) + return False + + context = engine.create_execution_context() + self._engines[key] = engine + self._contexts[key] = context + + self._allocate_buffers(key, height, width) + logger.info("TRT VQGAN engine loaded: %s (%s)", key, engine_path.name) + return True + + def _allocate_buffers(self, key: str, height: int, width: int): + # Per-engine batch (each res may have been built independently with + # a different batch); the top-level manifest["batch"] is only the + # last-built entry and unreliable when engines coexist. + batch = self._manifest["engines"][key]["input_shape"][0] + h_lat, w_lat = height // 16, width // 16 + device = torch.device("cuda") + + inp_buf = torch.empty(batch, 3, height, width, device=device, dtype=torch.float32) + + engine = self._engines[key] + out_dtype_trt = engine.get_tensor_dtype("indices") + out_dtype_is_int32 = (out_dtype_trt == trt.DataType.INT32) + if out_dtype_is_int32: + out_buf = torch.empty(batch, h_lat, w_lat, device=device, dtype=torch.int32) + else: + out_buf = torch.empty(batch, h_lat, w_lat, device=device, dtype=torch.int64) + + self._buffers[key] = (inp_buf, out_buf, out_dtype_is_int32) diff --git a/flash_rt/models/chameleon/__init__.py b/flash_rt/models/chameleon/__init__.py new file mode 100644 index 00000000..f9cef0a9 --- /dev/null +++ b/flash_rt/models/chameleon/__init__.py @@ -0,0 +1,25 @@ +"""FlashRT Chameleon-7B (VLM/LLM) model namespace. + +Chameleon-7B is a 32-layer early-fusion multimodal LLM (MHA 32x128 with +per-head QK LayerNorm + RoPE, attention_bias=false, mlp_bias=false, +SwiGLU FFN, 8192-token VQ-GAN image vocabulary). + +Two hardware paths: + +- Thor SM110 (``pipeline_thor.py``): dynamic per-tensor FP8 backbone with + fused norm/activation+quantize kernels, optional NVFP4 FFN layers, + causal CUTLASS SM100 FMHA (``libfmha_fp16_causal.so`` / + ``libfmha_fp8_causal.so``) and CUDA-graph decode. +- Orin SM87 (``pipeline_rtx.py``): INT8 W8A8 + INT4 W4A4 QuaRot-Hadamard + rotated weights, SM80 CUTLASS rowwise GEMMs, FA2 fp16 causal attention. + +Vendored code notice +-------------------- +``flash_rt/models/chameleon/vqgan`` contains the Chameleon VQ-GAN image +tokenizer reference implementation, vendored from Meta's Chameleon +repository (``chameleon.vae_ori``). Those files are Copyright (c) Meta +Platforms, Inc. and affiliates, licensed under the Chameleon License (see +the copyright headers in ``vqgan/*.py``). They are inference-only and are +used solely to decode generated image tokens; no training code is +included. +""" diff --git a/flash_rt/models/chameleon/pipeline_thor.py b/flash_rt/models/chameleon/pipeline_thor.py new file mode 100644 index 00000000..0741b91c --- /dev/null +++ b/flash_rt/models/chameleon/pipeline_thor.py @@ -0,0 +1,1126 @@ +"""FlashRT — standalone Chameleon-7B Thor SM110 pipeline forward functions. + +Chameleon-7B LLM (32-layer MHA, attention_bias=false, mlp_bias=false, +per-head QK LayerNorm + RoPE). Used by the standalone Chameleon Thor +frontend. + +All functions use raw-pointer interface (int pointers + Python primitives) +for CUDA Graph compatibility. No dynamic allocation, no torch ops, no sync. + +Functions: + chameleon_forward — Chameleon-7B LLM inference (dynamic FP8 default) + chameleon_forward_fp16 — pure-FP16 reference path + chameleon_forward_calibrate — FP8 static-scale calibration +""" + +from __future__ import annotations + +import math + +import flash_rt.flash_rt_kernels as _fvk +import torch + +from flash_rt.hardware.thor.shared_primitives import ( + _measure_scale_gpu, + _gpu_copy, + _gpu_sync, + _gpu_zero, +) + +try: + import flash_rt.flash_rt_fp4 as _fvk_fp4 +except Exception: + _fvk_fp4 = None + + +def _parse_fp4_layer_policy() -> frozenset: + """FP4 FFN layer policy from FLASHRT_RYNNVLA2_FP4_LAYERS env var. + + Values: + - unset / "": default = L0-L2 FP4, L3-L31 FP8 (safe, matches 001). + - "0-7": FP4 for L0..L7 inclusive. + - "0-14,20-31": FP4 for L0..L14 + L20..L31 (skip outlier L15-L19). + This is the SM120-sweep-validated aggressive setting (13ms savings on + RTX 5090; expect similar Thor gains). + - comma-separated list: "0,1,2,5" → FP4 on those specific layers. + + Returns the frozenset of FP8 layer indices (complement of FP4 set). + """ + import os as _os + val = _os.environ.get("FLASHRT_RYNNVLA2_FP4_LAYERS", "").strip() + if not val: + return frozenset(range(3, 32)) # default: L0-L2 FP4 + + fp4_layers = set() + for chunk in val.split(","): + chunk = chunk.strip() + if "-" in chunk: + a, b = chunk.split("-") + fp4_layers.update(range(int(a), int(b) + 1)) + else: + fp4_layers.add(int(chunk)) + fp4_layers = {li for li in fp4_layers if 0 <= li < 32} + return frozenset(li for li in range(32) if li not in fp4_layers) + + +# ── FP4 GEMM variant and layer policy ── +FP4_VARIANT = 9 +_FFN_FP8_LAYERS = _parse_fp4_layer_policy() + + +# ══════════════════════════════════════════════════════════════════ +# Chameleon-7B LLM (32 layers, mixed FP4/FP8 GEMMs) +# ══════════════════════════════════════════════════════════════════ + +def chameleon_forward( + gemm, fvk, bufs, weights, dims, scales_dev, + *, attn, stream: int = 0, + alpha_host=None, awq_v_proj=None, + ffn_down_clamp_value: float = 10000.0, + ffn_clamp_layers=None, + dynamic_fp8_layers: frozenset = frozenset(), + fp4_ffn_layers: frozenset = frozenset(), + probe=None, +) -> None: + """Chameleon-7B LLM forward pass (32 layers). + + Production precision: dynamic per-tensor FP8 on all layers, with the FFN of + ``fp4_ffn_layers`` optionally run in NVFP4 W4A16 (decoupled from attention). + + Output: hidden_all = RMSNorm(x) written to bufs['hidden_all'] as + full [Se, D] FP16 tensor (consumed by action_head_forward). + + ``ffn_down_clamp_value``: Clamp down_out (o_proj_out) to ±V after + the FFN before residual_2. Chameleon-7B L31 down_proj's FP32 + accumulator × alpha can push output beyond ±65504 producing inf. + Applied after both FP4 and FP8 FFN paths. Set to <= 0 to disable. + + ``ffn_clamp_layers``: Optional set/frozenset of layer indices where the + FFN clamps are applied. ``None`` preserves the historical behavior and + clamps every layer; callers can pass e.g. ``frozenset({31})`` to clamp + only the deep outlier layer and avoid redundant elementwise passes. + + ``dynamic_fp8_layers``: frozenset of layer indices to run with full FP8 + GEMMs but RUNTIME (per-forward) per-tensor activation scaling instead of + the static calibrated scale. Uses quantize_fp8_device_fp16 (GPU amax) + + fp8_nn_dev_fp16 (device-scale GEMM). Restores 512-res precision (static + scale is wrong for long sequences) at full FP8 speed. CUDA-Graph safe. + This is the default for all 32 layers; see docs/chameleon_thor_sm110.md. + + ``fp4_ffn_layers``: frozenset of layer indices whose FFN runs in NVFP4 + W4A16 while attention stays on the dynamic-FP8 path (decoupled). Requires + those layers' FP4 weights to be packed. 512-res default = L0-7 (~1.10x E2E, + cos >= 0.99 on 3/4 variants); empty at 256-res (FP4 breaks precision). See + docs/chameleon_thor_sm110.md (§6 FP4). + + Optional ``probe`` dict for layer-wise precision debugging:: + + { + 'layers': tuple[int, ...], # layer indices to snapshot post-residual-2 + 'bufs': tuple[int, ...], # device pointers, one per layer index, + # each must hold >= Se*D fp16 elements + 'final_buf': int, # device pointer for post-final-RMSNorm + # snapshot. Pass 0 to skip. + } + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + if alpha_host is None: + alpha_host = weights.get('alpha_host') + if awq_v_proj is None: + awq_v_proj = weights.get('awq_v_proj') + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + xn_fp8_ptr = int(bufs['xn_fp8']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + + # FP4 FFN buffers + act_fp4_ptr = int(bufs['act_fp4']) + act_sfa_ptr = int(bufs['act_sfa']) + ffn_act_fp4_ptr = int(bufs['ffn_act_fp4']) + ffn_act_sfa_ptr = int(bufs['ffn_act_sfa']) + gu_merged_ptr = int(bufs['gu_merged']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + # Layer-0 pre-attention RMSNorm + FP8 quantize + fvk.rms_norm_fp16( + x_ptr, int(weights['input_ln_w'][0]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, int(scales_dev["act_qkv"][0]), Se * D, int(stream), + ) + + for li in range(L): + clamp_this_layer = ( + ffn_down_clamp_value > 0.0 + and (ffn_clamp_layers is None or li in ffn_clamp_layers) + ) + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + K_ptr = int(slots["K"]) + V_ptr = int(slots["V"]) + O_ptr = int(slots["O"]) + + d_act_qkv = int(scales_dev["act_qkv"][li]) + d_act_o = int(scales_dev["act_o"][li]) + d_act_gu = int(scales_dev["act_gu"][li]) + d_act_d = int(scales_dev["act_down"][li]) + + d_w_qkv = int(weights['d_w_qkv'][li]) + d_w_o = int(weights['d_w_o'][li]) + d_w_gu = int(weights['d_w_gu'][li]) + d_w_d = int(weights['d_w_d'][li]) + + q_w_ptr = int(weights['q_w'][li]) + k_w_ptr = int(weights['k_w'][li]) + v_w_ptr = int(weights['v_w'][li]) + + + # ═══ Dynamic per-tensor FP8 branch ═══ + # Runtime amax scaling (quantize_fp8_device_fp16 -> fp8_nn_dev_fp16) + # instead of static-calibrated scale. The static per-tensor scale is + # wrong for long-sequence (512-res) variants where deep-layer + # activations shift; recomputing amax each forward restores precision + # at full FP8 speed (no FP16 GEMM). CUDA-Graph safe: the device scale + # pointer is dereferenced at replay, so it adapts per input. + if li in dynamic_fp8_layers: + # Self-sufficient entry: re-derive FP16 xn from residual stream. + # Fused RMSNorm + dynamic FP8 quantize: amax is measured inside + # the norm's own output-write pass, skipping the separate + # absmax_kernel read of xn (one fewer full pass over Se*D). + dyn_qkv = int(scales_dev['dyn_act_qkv'][li]) + dyn_o = int(scales_dev['dyn_act_o'][li]) + dyn_gu = int(scales_dev['dyn_act_gu'][li]) + dyn_d = int(scales_dev['dyn_act_down'][li]) + + fvk.rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, int(weights['input_ln_w'][li]), xn_ptr, xn_fp8_ptr, + dyn_qkv, Se, D, 1e-5, int(stream)) + + # QKV: device-scale GEMM (xn_fp8 already produced above). + # All three GEMMs read xn_fp8; Q writes into xn (aliased Q_O) + # AFTER the quantize, so the clobber is harmless. + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, v_w_ptr, V_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, k_w_ptr, K_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, q_w_ptr, Q_ptr, + Se, D, D, dyn_qkv, d_w_qkv, int(stream)) + + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream)) + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # O projection: dynamic. + fvk.quantize_fp8_device_fp16( + O_ptr, xn_fp8_ptr, dyn_o, Se * D, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['o_w'][li]), + o_proj_out_ptr, Se, D, D, dyn_o, d_w_o, + int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── FFN: decoupled NVFP4 (opt-in) or dynamic FP8 ── + # DECOUPLED mode keeps attention on dynamic FP8 (above) but runs the + # FFN in NVFP4 W4A16 for FP4-eligible layers — isolating FFN + # precision from attention precision. Falls back to dynamic FP8 FFN. + _use_fp4_ffn = ( + _fvk_fp4 is not None + and li in fp4_ffn_layers + and int(weights['gu_w_fp4'][li]) != 0 + ) + if _use_fp4_ffn: + # Post-attn RMSNorm -> FP16 xn (no FP8 quantize needed here; + # the FP4 path quantizes xn_ptr directly below). + fvk.rms_norm_fp16(x_ptr, int(weights['post_ln_w'][li]), xn_ptr, + Se, D, 1e-5, int(stream)) + # NOTE: intended for SHALLOW layers only. Unlike the FP8 branch + # below, there is no intermediate clamp on the SwiGLU output + # (gate_geglu fuses silu*mul + FP4 quantize). A deep layer whose + # gu exceeds fp16 max (~65504, e.g. L31) could overflow to + # inf/nan *before* the FP4 quantize; the final o_proj_out clamp + # cannot recover that. The shipped 512 default (L0-7) is safe. + # gate+up merged NVFP4 GEMM (dynamic per-block SFA). + _fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + xn_ptr, act_fp4_ptr, act_sfa_ptr, Se, D, False, int(stream)) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, act_fp4_ptr, act_sfa_ptr, + int(weights['gu_w_fp4'][li]), int(weights['gu_sfb'][li]), + gu_merged_ptr, Se, 2 * Dff, D, 1.0, 0.0, int(stream)) + # fused SwiGLU + quantize down-proj input to NVFP4. + _fvk_fp4.gate_geglu_fp4_sfa_v2_fp16( + gu_merged_ptr, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + Se, Dff, int(stream)) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + int(weights['d_w_fp4'][li]), int(weights['d_sfb'][li]), + o_proj_out_ptr, Se, D, Dff, 1.0, 0.0, int(stream)) + else: + # Post-attn residual add + RMSNorm + dynamic FP8 quantize, + # fused into one elementwise kernel: residual is fp16-rounded + # (same as residual_add_fp16), ssq is over the rounded values + # (same as rms_norm reading the fp16 residual), the amax for + # dyn_gu is folded into the xn write pass, and the residual is + # register-cached so xn is never re-read from global. + fvk.residual_add_rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_ptr, xn_fp8_ptr, dyn_gu, Se, D, 1e-5, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['gate_w'][li]), + int(bufs['gate_out']), Se, Dff, D, + dyn_gu, d_w_gu, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['up_w'][li]), + int(bufs['up_out']), Se, Dff, D, + dyn_gu, d_w_gu, int(stream)) + + # Down input gu = silu(gate)*up in FP16, fused with its own + # dynamic FP8 quantize on layers with no outlier clamp (amax + # measured inside the SwiGLU write pass, skipping a separate + # absmax_kernel read of the Se*Dff intermediate). Layers that + # need the outlier clamp (default: L31 only) keep the + # unfused gate_geglu -> clamp -> quantize sequence since the + # clamp must run BEFORE the amax/scale is computed. + if clamp_this_layer: + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Se * Dff, int(stream)) + fvk.clamp_inplace_fp16(int(bufs['gate_out']), + float(ffn_down_clamp_value), + Se * Dff, int(stream)) + fvk.quantize_fp8_device_fp16( + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Se * Dff, int(stream)) + else: + fvk.gate_geglu_quantize_dynamic_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Se * Dff, int(stream)) + gemm.fp8_nn_dev_fp16(int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, Se, D, Dff, dyn_d, d_w_d, + int(stream)) + + if clamp_this_layer: + fvk.clamp_inplace_fp16(o_proj_out_ptr, + float(ffn_down_clamp_value), + Se * D, int(stream)) + + # Residual 2 + next-layer prep. + if li < L - 1: + fvk.residual_add_rms_norm_fp16( + x_ptr, o_proj_out_ptr, int(weights['input_ln_w'][li + 1]), + xn_ptr, Se, D, 1e-5, int(stream)) + # A following STATIC FP8 layer consumes xn_fp8; produce it. + if (li + 1) not in dynamic_fp8_layers: + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, + int(scales_dev["act_qkv"][li + 1]), Se * D, int(stream)) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, + int(stream)) + + if probe is not None: + _pl = probe.get('layers') or () + if li in _pl: + _gpu_copy(int(probe['bufs'][_pl.index(li)]), + x_ptr, Se * D * 2, stream) + continue + # ═══ End dynamic per-tensor FP8 branch ═══ + + # ── QKV FP8 GEMMs ── + # 002 has no attention bias; pass zero-buffer for fp8_nn_bias epilogue + if alpha_host is not None: + alpha_qkv = float(alpha_host[li * 4 + 0]) + zero_bias_ptr = int(bufs['zero_bias_d']) + + # V projection: AWQ path OR shared-scale path. + # IMPORTANT ORDER: The chameleon attention backend aliases + # Q_ptr to xn_ptr (bufs['xn']). If Q is written first, xn is + # clobbered, and the AWQ V path reads Q's output instead of + # the RMSNormed xn. So run AWQ V-quantize + V-GEMM BEFORE Q + # (V's destination is the KV cache, separate buffer, safe). + if awq_v_proj is not None: + fvk.awq_quant_fp8_static_fp16( + xn_ptr, + int(awq_v_proj['inv_s_ptrs'][li]), + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['act_scale_ptrs'][li]), + Se, D, int(stream), + ) + v_alpha = float(awq_v_proj['alpha_host'][li]) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_ptrs'][li]), + V_ptr, zero_bias_ptr, + Se, D, D, v_alpha, int(stream), + ) + else: + gemm.fp8_nn_bias( + xn_fp8_ptr, v_w_ptr, V_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + + # Q and K: also use AWQ smoothed path (same xn_v_fp8, per-proj weights) + if awq_v_proj is not None and 'w_q_ptrs' in awq_v_proj: + q_alpha = float(awq_v_proj['alpha_q_host'][li]) + k_alpha = float(awq_v_proj['alpha_k_host'][li]) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_q_ptrs'][li]), + Q_ptr, zero_bias_ptr, + Se, D, D, q_alpha, int(stream), + ) + gemm.fp8_nn_bias( + int(awq_v_proj['xn_v_fp8']), + int(awq_v_proj['w_k_ptrs'][li]), + K_ptr, zero_bias_ptr, + Se, D, D, k_alpha, int(stream), + ) + else: + gemm.fp8_nn_bias( + xn_fp8_ptr, q_w_ptr, Q_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + xn_fp8_ptr, k_w_ptr, K_ptr, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + else: + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, q_w_ptr, Q_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, k_w_ptr, K_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, v_w_ptr, V_ptr, + Se, D, D, d_act_qkv, d_w_qkv, int(stream), + ) + + # ── Fused per-head QK LayerNorm + RoPE ── + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # ── MHA via attention backend ── + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # ── O projection ── + fvk.quantize_fp8_static_fp16( + O_ptr, xn_fp8_ptr, d_act_o, Se * D, int(stream), + ) + + if alpha_host is not None: + alpha_o = float(alpha_host[li * 4 + 1]) + gemm.fp8_nn_bias( + xn_fp8_ptr, int(weights['o_w'][li]), o_proj_out_ptr, + int(bufs['zero_bias_d']), + Se, D, D, alpha_o, int(stream), + ) + else: + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, int(weights['o_w'][li]), o_proj_out_ptr, + Se, D, D, d_act_o, d_w_o, int(stream), + ) + + # ── Post-attention: residual + RMSNorm (path depends on FFN precision) ── + if li in _FFN_FP8_LAYERS: + fvk.residual_add_rms_norm_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_fp8_ptr, Se, D, 1e-5, + d_act_gu, int(stream), + ) + else: + fvk.residual_add_fp16( + x_ptr, o_proj_out_ptr, Se * D, int(stream), + ) + fvk.rms_norm_fp16( + x_ptr, int(weights['post_ln_w'][li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # ── FFN (legacy static path; only for layers excluded from + # dynamic_fp8_layers): NVFP4 if li not in _FFN_FP8_LAYERS, else FP8 ── + gu_w_fp4_li = int(weights['gu_w_fp4'][li]) if li not in _FFN_FP8_LAYERS else 0 + fp4_available = ( + _fvk_fp4 is not None + and li not in _FFN_FP8_LAYERS + and gu_w_fp4_li != 0 + ) + if fp4_available: + # FP4 path + _fvk_fp4.quantize_fp4_dynamic_sfa_fp16( + xn_ptr, act_fp4_ptr, act_sfa_ptr, + Se, D, False, int(stream), + ) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, + act_fp4_ptr, act_sfa_ptr, + int(weights['gu_w_fp4'][li]), int(weights['gu_sfb'][li]), + gu_merged_ptr, Se, 2 * Dff, D, + 1.0, 0.0, int(stream), + ) + _fvk_fp4.gate_geglu_fp4_sfa_v2_fp16( + gu_merged_ptr, ffn_act_fp4_ptr, ffn_act_sfa_ptr, + Se, Dff, int(stream), + ) + _fvk_fp4.cutlass_fp4_gemm_variant( + FP4_VARIANT, + ffn_act_fp4_ptr, ffn_act_sfa_ptr, + int(weights['d_w_fp4'][li]), int(weights['d_sfb'][li]), + o_proj_out_ptr, Se, D, Dff, + 1.0, 0.0, int(stream), + ) + else: + # FP8 path (also used as fallback when FP4 is unavailable + # for a layer in the L0-2 range). + gate_w_ptr = int(weights['gate_w'][li]) + up_w_ptr = int(weights['up_w'][li]) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, gate_w_ptr, int(bufs['gate_out']), + Se, Dff, D, d_act_gu, d_w_gu, int(stream), + ) + gemm.fp8_nn_dev_fp16( + xn_fp8_ptr, up_w_ptr, int(bufs['up_out']), + Se, Dff, D, d_act_gu, d_w_gu, int(stream), + ) + # Down-proj: 2-tier adaptive dispatch. + # Tier 1 (standard FP8): no AWQ dict → fused silu_mul_split_fp8 + # Tier 2 (AWQ D smooth): li in awq_d_layers → per-K smooth + FP8 + _has_awq = (awq_v_proj is not None + and 'inv_s_D_ptrs' in awq_v_proj) + _in_awq_d = (_has_awq and li in awq_v_proj.get( + 'awq_d_layers', frozenset())) + + if _in_awq_d: + # Tier 2: AWQ D smoothed FP8 path (fused SwiGLU). + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Se * Dff, int(stream)) + fvk.awq_quant_fp8_static_fp16( + int(bufs['gate_out']), + int(awq_v_proj['inv_s_D_ptrs'][li]), + int(bufs['gu_fp8']), + int(awq_v_proj['act_scale_D_ptrs'][li]), + Se, Dff, int(stream)) + gemm.fp8_nn_bias( + int(bufs['gu_fp8']), + int(awq_v_proj['w_D_ptrs'][li]), + o_proj_out_ptr, int(bufs['zero_bias_d']), + Se, D, Dff, + float(awq_v_proj['alpha_D_host'][li]), + int(stream)) + else: + # Tier 1: standard per-tensor FP8 (no AWQ) + fvk.silu_mul_split_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gu_fp8']), + Se * Dff, d_act_d, int(stream), + ) + gemm.fp8_nn_dev_fp16( + int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, + Se, D, Dff, d_act_d, d_w_d, int(stream), + ) + + # ── Clamp down_out to fp16 range ── + # Chameleon-7B L31 down_proj's FP32 accumulator × alpha can push + # the FP16 output beyond ±65504 producing inf, which propagates + # through the final RMSNorm and destroys action precision. Mirror + # the RTX FP8 path's clamp (pipeline_rtx.py:829-833). + if clamp_this_layer: + fvk.clamp_inplace_fp16( + o_proj_out_ptr, float(ffn_down_clamp_value), + Se * D, int(stream), + ) + + # ── Fused: residual_2 + next-layer input_ln + FP8 quantize ── + # When AWQ V-proj is enabled we need the FP16 xn buffer populated + # for the NEXT layer's V-proj activation smoothing kernel (which + # reads FP16 xn, not FP8). Split the fused path into + # residual_add_rms_norm_fp16 + quantize_fp8_static_fp16 (one + # extra launch per layer) so xn_ptr stays valid. + if li < L - 1: + if awq_v_proj is not None: + fvk.residual_add_rms_norm_fp16( + x_ptr, o_proj_out_ptr, + int(weights['input_ln_w'][li + 1]), + xn_ptr, Se, D, 1e-5, int(stream), + ) + fvk.quantize_fp8_static_fp16( + xn_ptr, xn_fp8_ptr, + int(scales_dev["act_qkv"][li + 1]), + Se * D, int(stream), + ) + else: + fvk.residual_add_rms_norm_fp8_fp16( + x_ptr, o_proj_out_ptr, + int(weights['input_ln_w'][li + 1]), + xn_fp8_ptr, Se, D, 1e-5, + int(scales_dev["act_qkv"][li + 1]), int(stream), + ) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── Optional layer probe: snapshot post-residual-2 hidden state ── + if probe is not None: + probe_layers = probe.get('layers') or () + if li in probe_layers: + idx = probe_layers.index(li) + snap_ptr = int(probe['bufs'][idx]) + if snap_ptr != 0: + _gpu_copy(snap_ptr, x_ptr, Se * D * 2, stream) + + # ── Final RMSNorm → hidden_all [Se, D] FP16 ── + fvk.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), hidden_all_ptr, + Se, D, 1e-5, int(stream), + ) + + # ── Optional final probe ── + if probe is not None: + final_buf = int(probe.get('final_buf', 0)) + if final_buf != 0: + _gpu_copy(final_buf, hidden_all_ptr, Se * D * 2, stream) + + +def chameleon_decode_step( + gemm, fvk, bufs, weights, dims, scales_dev, + *, attn, pos: int, stream: int = 0, + ffn_down_clamp_value: float = 60000.0, + ffn_clamp_layers=frozenset({31}), +) -> None: + """Single-token (Se=1) incremental decode step at position ``pos``. + + Mirrors chameleon_forward's dynamic per-tensor FP8 branch op-for-op, + with: + * K/V GEMMs writing the KV cache row at ``pos`` + (``attn.kv_row_ptrs``) — history rows were filled by prefill; + * RoPE cos/sin taken at row ``pos`` via byte offsets; + * attention dispatched through ``attn.run_decode`` (bottom-right + causal mask) over kv_len=pos+1 keys; + * final RMSNorm written to ``hidden_all`` row 0. + + The token embedding must already be in ``bufs['x']`` row 0, and the + residual stream in ``bufs['x']`` is updated in place across layers. + CUDA-Graph safe (int pointers only, no allocations); ``pos`` is a host + scalar, so callers run this eagerly, outside graph capture. + """ + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + xn_fp8_ptr = int(bufs['xn_fp8']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + + cos_ptr = int(weights['rope_cos']) + pos * Hd * 2 + sin_ptr = int(weights['rope_sin']) + pos * Hd * 2 + + for li in range(L): + clamp_this_layer = ( + ffn_down_clamp_value > 0.0 + and (ffn_clamp_layers is None or li in ffn_clamp_layers) + ) + dyn_qkv = int(scales_dev['dyn_act_qkv'][li]) + dyn_o = int(scales_dev['dyn_act_o'][li]) + dyn_gu = int(scales_dev['dyn_act_gu'][li]) + dyn_d = int(scales_dev['dyn_act_down'][li]) + + d_w_o = int(weights['d_w_o'][li]) + d_w_gu = int(weights['d_w_gu'][li]) + d_w_d = int(weights['d_w_d'][li]) + + K_row_ptr, V_row_ptr = attn.kv_row_ptrs("chameleon", li, pos) + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + O_ptr = int(slots["O"]) + + # Fused RMSNorm + dynamic FP8 quantize from the residual stream. + fvk.rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, int(weights['input_ln_w'][li]), xn_ptr, xn_fp8_ptr, + dyn_qkv, 1, D, 1e-5, int(stream)) + + # QKV GEMMs (M=1): K/V land in the cache row at pos, Q in Q_O. + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['v_w'][li]), V_row_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['k_w'][li]), K_row_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['q_w'][li]), Q_ptr, + 1, D, D, dyn_qkv, int(weights['d_w_qkv'][li]), + int(stream)) + + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_row_ptr, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + 1, H, Hd, 1e-5, int(stream)) + attn.run_decode("chameleon", li, kv_len=pos + 1, stream=int(stream)) + + # O projection (M=1). + fvk.quantize_fp8_device_fp16( + O_ptr, xn_fp8_ptr, dyn_o, D, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['o_w'][li]), + o_proj_out_ptr, 1, D, D, dyn_o, d_w_o, + int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, D, int(stream)) + + # FFN: residual add + RMSNorm + dynamic FP8 quantize (fused). + fvk.residual_add_rms_norm_quantize_dynamic_fp8_fp16( + x_ptr, o_proj_out_ptr, int(weights['post_ln_w'][li]), + xn_ptr, xn_fp8_ptr, dyn_gu, 1, D, 1e-5, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['gate_w'][li]), + int(bufs['gate_out']), 1, Dff, D, + dyn_gu, d_w_gu, int(stream)) + gemm.fp8_nn_dev_fp16(xn_fp8_ptr, int(weights['up_w'][li]), + int(bufs['up_out']), 1, Dff, D, + dyn_gu, d_w_gu, int(stream)) + + if clamp_this_layer: + fvk.gate_geglu_fp16(int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), Dff, int(stream)) + fvk.clamp_inplace_fp16(int(bufs['gate_out']), + float(ffn_down_clamp_value), + Dff, int(stream)) + fvk.quantize_fp8_device_fp16( + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Dff, int(stream)) + else: + fvk.gate_geglu_quantize_dynamic_fp8_fp16( + int(bufs['gate_out']), int(bufs['up_out']), + int(bufs['gate_out']), int(bufs['gu_fp8']), dyn_d, + Dff, int(stream)) + gemm.fp8_nn_dev_fp16(int(bufs['gu_fp8']), int(weights['d_w'][li]), + o_proj_out_ptr, 1, D, Dff, dyn_d, d_w_d, + int(stream)) + + if clamp_this_layer: + fvk.clamp_inplace_fp16(o_proj_out_ptr, + float(ffn_down_clamp_value), + D, int(stream)) + + # Residual 2 (next layer re-derives xn from the residual stream). + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, D, int(stream)) + + # Final RMSNorm → hidden_all row 0. + fvk.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), hidden_all_ptr, + 1, D, 1e-5, int(stream), + ) + + +def chameleon_forward_fp16( + gemm, fvk, bufs, weights, dims, + *, attn, stream: int = 0, + ffn_gate_clamp_value: float = 10000.0, + probe=None, +) -> None: + """Chameleon-7B FP16-only forward (no FP8, no FP4, no AWQ). + + Ported from pipeline_rtx.chameleon_forward. All 32 layers run pure + FP16 GEMMs via ``gemm.fp16_nn`` — same on Thor as RTX. + + Precision-optimal path (cosine target ≥ 0.99 vs vendor bf16) at the + cost of ~2× the FP8 path latency in the LLM. Recommended when + downstream ActionHead is sensitive to accumulated FP8 error. + + Weights required (all FP16, KN row-major layout — spec built with + ``use_fp8=False``): + q_w[li], k_w[li], v_w[li], o_w[li] : (D, D) + gate_w[li], up_w[li] : (D, Dff) + d_w[li] / down_w[li] : (Dff, D) + input_ln_w[li], post_ln_w[li] : (D,) + q_norm_w/b[li], k_norm_w/b[li] : (1, HD) or (HD,) + final_norm_w : (D,) + rope_cos, rope_sin : (max_pos, HD) + + Output: hidden_all = RMSNorm(x_post_res_2, final_norm_w) written to + bufs['hidden_all'] as (Se, D) FP16 (consumed by action_head_forward). + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + o_proj_out_ptr = int(bufs['o_proj_out']) + hidden_all_ptr = int(bufs['hidden_all']) + # Reuse existing gate/up buffers (Se, Dff) fp16 + gate_ptr = int(bufs['gate_out']) + up_ptr = int(bufs['up_out']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + q_w = weights['q_w'] + k_w = weights['k_w'] + v_w = weights['v_w'] + o_w = weights['o_w'] + gate_w = weights['gate_w'] + up_w = weights['up_w'] + down_w = weights['d_w'] # frontend uses 'd_w' key for down projection + + input_ln_w = weights['input_ln_w'] + post_ln_w = weights['post_ln_w'] + q_norm_w = weights['q_norm_w'] + q_norm_b = weights['q_norm_b'] + k_norm_w = weights['k_norm_w'] + k_norm_b = weights['k_norm_b'] + final_norm_w = int(weights['final_norm_w']) + + for li in range(L): + slots = attn.get_slot_ptrs("chameleon", li) + Q_ptr = int(slots["Q"]) + K_ptr = int(slots["K"]) + V_ptr = int(slots["V"]) + + # input_layernorm (RMSNorm eps=1e-5) → xn + fvk.rms_norm_fp16( + x_ptr, int(input_ln_w[li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # Q / K / V GEMMs (no bias in 002). + # NOTE: Q_ptr aliases xn_ptr on Thor (chameleon slots["Q_O"] = + # bufs['xn'].data_ptr()). Because gemm.fp16_nn reads A (xn) and + # writes D (Q) at the SAME fp16 dtype and SAME buffer, cuBLAS + # in-place semantics are undefined → output corruption. + # Route Q through o_proj_out_ptr scratch, then copy into Q slot. + # (V and K write to separate KV cache buffers; safe direct.) + gemm.fp16_nn(xn_ptr, int(v_w[li]), V_ptr, Se, D, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(k_w[li]), K_ptr, Se, D, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(q_w[li]), o_proj_out_ptr, + Se, D, D, int(stream)) + _gpu_copy(Q_ptr, o_proj_out_ptr, Se * D * 2, stream) + + # Per-head QK LayerNorm + RoPE (in-place). + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(q_norm_w[li]), int(q_norm_b[li]), + int(k_norm_w[li]), int(k_norm_b[li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # Causal MHA (CUTLASS SM110 FMHA / cuBLAS fallback). + attn.run("chameleon", li, q_seq=Se, kv_seq=Se, stream=int(stream)) + + # O projection. + gemm.fp16_nn(Q_ptr, int(o_w[li]), o_proj_out_ptr, + Se, D, D, int(stream)) + + # Residual 1. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # post_attention_layernorm. + fvk.rms_norm_fp16( + x_ptr, int(post_ln_w[li]), xn_ptr, + Se, D, 1e-5, int(stream), + ) + + # FFN gate / up (SwiGLU). + gemm.fp16_nn(xn_ptr, int(gate_w[li]), gate_ptr, + Se, Dff, D, int(stream)) + gemm.fp16_nn(xn_ptr, int(up_w[li]), up_ptr, + Se, Dff, D, int(stream)) + fvk.gate_geglu_fp16(gate_ptr, up_ptr, gate_ptr, Se * Dff, int(stream)) + + # L31 gate*up overflow guard (fp16 max ≈ 65504, Chameleon L31 + # gate*up amax observed ≈ 48000). See pipeline_rtx.py:243-256. + if ffn_gate_clamp_value > 0.0: + fvk.clamp_inplace_fp16( + gate_ptr, float(ffn_gate_clamp_value), + Se * Dff, int(stream), + ) + + # down projection. + gemm.fp16_nn(gate_ptr, int(down_w[li]), o_proj_out_ptr, + Se, D, Dff, int(stream)) + + # Residual 2. + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # Optional probe: snapshot post-residual-2 hidden state. + if probe is not None: + probe_layers = probe.get('layers') or () + if li in probe_layers: + idx = probe_layers.index(li) + snap_ptr = int(probe['bufs'][idx]) + if snap_ptr != 0: + _gpu_copy(snap_ptr, x_ptr, Se * D * 2, stream) + + # Final RMSNorm → hidden_all. + fvk.rms_norm_fp16( + x_ptr, final_norm_w, hidden_all_ptr, + Se, D, 1e-5, int(stream), + ) + + if probe is not None: + final_buf = int(probe.get('final_buf', 0)) + if final_buf != 0: + _gpu_copy(final_buf, hidden_all_ptr, Se * D * 2, stream) + + +# ══════════════════════════════════════════════════════════════════ +# Chameleon-7B LLM calibration (FP16 + amax measurement) +# ══════════════════════════════════════════════════════════════════ + +def _d2h_float(d_ptr: int) -> float: + """Read a single float32 from device to host.""" + t = torch.empty(1, dtype=torch.float32, device='cuda') + import ctypes + ctypes.CDLL('libcudart.so').cudaMemcpy( + ctypes.c_void_p(t.data_ptr()), + ctypes.c_void_p(d_ptr), + 4, 2, # cudaMemcpyDeviceToDevice is 3, D2H is 2 + ) + return float(t.item()) + + +def _d2h_floats(d_ptr: int, n: int) -> list: + """Read n float32 values from device to host.""" + t = torch.empty(n, dtype=torch.float32, device='cuda') + import ctypes + ctypes.CDLL('libcudart.so').cudaMemcpy( + ctypes.c_void_p(t.data_ptr()), + ctypes.c_void_p(d_ptr), + n * 4, 2, + ) + return t.cpu().tolist() + + +def chameleon_forward_calibrate( + gemm, fvk_mod, bufs, weights, dims, + calib_scales_ptr, stream: int = 0, + attn_calib_scales_ptr: int = 0, +) -> None: + """Calibrate Chameleon-7B FP8 scales. + + 4 quantization points per layer × 32 layers = 128 scales. + Points: act_qkv, act_o, act_gu, act_down. + """ + Se = int(dims['Se']) + D = int(dims['D']) + Dff = int(dims['Dff']) + L = int(dims['L']) + H = int(dims['H']) + Hd = int(dims['Hd']) + + import numpy as np + + x_ptr = int(bufs['x']) + xn_ptr = int(bufs['xn']) + o_proj_out_ptr = int(bufs['o_proj_out']) + gate_out_ptr = int(bufs['gate_out']) + up_out_ptr = int(bufs['up_out']) + down_out_ptr = int(bufs['down_out']) + Q_ptr_buf = int(bufs['Q']) + K_ptr_buf = int(bufs['K']) + V_ptr_buf = int(bufs['V']) + O_ptr_buf = int(bufs['O']) + + calib_buf = int(bufs['calib_buf']) + d_scale = int(bufs['d_scale']) + fp8_scratch = int(bufs['fp8_scratch']) + norm_scratch = int(bufs['norm_scratch']) + + cos_ptr = int(weights['rope_cos']) + sin_ptr = int(weights['rope_sin']) + + w_scales_dev = int(weights['w_scales_flat']) + ws_host = _d2h_floats(w_scales_dev, L * 4) + + _gpu_zero(calib_buf, L * 4 * 4, stream) + + for li in range(L): + # ── 1. RMSNorm → measure amax (act_qkv scale) ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['input_ln_w'][li]), norm_scratch, + Se, D, 1e-5, int(stream), + ) + _measure_scale_gpu(fvk_mod, norm_scratch, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_qkv = _d2h_float(d_scale) + cs_qkv = calib_buf + (li * 4 + 0) * 4 + _gpu_copy(cs_qkv, d_scale, 4, stream) + + # ── 2. Quantize xn → FP8 ── + fvk_mod.quantize_fp8_static_fp16( + norm_scratch, int(bufs['xn_fp8']), cs_qkv, + Se * D, int(stream), + ) + + # ── 3. Q/K/V FP8 GEMMs (no bias for 002) ── + q_w_ptr = int(weights['q_w'][li]) + k_w_ptr = int(weights['k_w'][li]) + v_w_ptr = int(weights['v_w'][li]) + + alpha_qkv = float(np.float32(as_qkv) * np.float32(ws_host[li * 4 + 0])) + zero_bias_ptr = int(bufs['zero_bias_d']) + + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), q_w_ptr, Q_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), k_w_ptr, K_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), v_w_ptr, V_ptr_buf, zero_bias_ptr, + Se, D, D, alpha_qkv, int(stream), + ) + + # ── 4. Fused QK LayerNorm + RoPE ── + fvk_mod.qk_norm_rope_fused_fp16( + Q_ptr_buf, K_ptr_buf, + int(weights['q_norm_w'][li]), int(weights['q_norm_b'][li]), + int(weights['k_norm_w'][li]), int(weights['k_norm_b'][li]), + cos_ptr, sin_ptr, + Se, H, Hd, 1e-5, int(stream), + ) + + # ── 4b. Optional Q/K/V amax for FP8 attention ── + if attn_calib_scales_ptr: + n_qkv = Se * H * Hd + for i, ptr in enumerate((Q_ptr_buf, K_ptr_buf, V_ptr_buf)): + _measure_scale_gpu( + fvk_mod, ptr, n_qkv, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + _gpu_copy( + attn_calib_scales_ptr + (li * 3 + i) * 4, + d_scale, 4, stream, + ) + + # ── 5. Attention (cuBLAS — no FMHA during calibration) ── + attn_scale = 1.0 / math.sqrt(float(Hd)) + fvk_mod.attention_qkv_fp16( + bufs['ctx'], Q_ptr_buf, K_ptr_buf, V_ptr_buf, + int(bufs['logits']), O_ptr_buf, + Se, Se, H, Hd, attn_scale, int(stream), + ) + + # ── 6. O proj — measure amax → quantize → GEMM ── + _measure_scale_gpu(fvk_mod, O_ptr_buf, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_o = _d2h_float(d_scale) + cs_o = calib_buf + (li * 4 + 1) * 4 + _gpu_copy(cs_o, d_scale, 4, stream) + fvk_mod.quantize_fp8_static_fp16( + O_ptr_buf, int(bufs['xn_fp8']), cs_o, + Se * D, int(stream), + ) + alpha_o = float(np.float32(as_o) * np.float32(ws_host[li * 4 + 1])) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), int(weights['o_w'][li]), o_proj_out_ptr, + zero_bias_ptr, + Se, D, D, alpha_o, int(stream), + ) + + # ── 7. Residual 1 ── + fvk_mod.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── 8. Post-attn RMSNorm → measure amax (act_gu scale) ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['post_ln_w'][li]), norm_scratch, + Se, D, 1e-5, int(stream), + ) + _measure_scale_gpu(fvk_mod, norm_scratch, Se * D, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_gu = _d2h_float(d_scale) + cs_gu = calib_buf + (li * 4 + 2) * 4 + _gpu_copy(cs_gu, d_scale, 4, stream) + + # ── 9. Quantize → FP8 ── + fvk_mod.quantize_fp8_static_fp16( + norm_scratch, int(bufs['xn_fp8']), cs_gu, + Se * D, int(stream), + ) + + # ── 10. Gate + Up FP8 GEMMs (no bias) ── + gate_w_ptr = int(weights['gate_w'][li]) + up_w_ptr = int(weights['up_w'][li]) + alpha_gu = float(np.float32(as_gu) * np.float32(ws_host[li * 4 + 2])) + zero_bias_dff = int(bufs['zero_bias_dff']) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), gate_w_ptr, gate_out_ptr, + zero_bias_dff, + Se, Dff, D, alpha_gu, int(stream), + ) + gemm.fp8_nn_bias( + int(bufs['xn_fp8']), up_w_ptr, up_out_ptr, + zero_bias_dff, + Se, Dff, D, alpha_gu, int(stream), + ) + + # ── 11. SiLU(gate)*up → measure amax → FP8 ── + silu_scr = int(bufs['silu_scratch']) + _gpu_copy(silu_scr, gate_out_ptr, Se * Dff * 2, stream) + fvk_mod.gate_geglu_fp16(silu_scr, up_out_ptr, down_out_ptr, + Se * Dff, int(stream)) + _measure_scale_gpu(fvk_mod, down_out_ptr, Se * Dff, d_scale, fp8_scratch, stream) + _gpu_sync(stream) + as_d = _d2h_float(d_scale) + cs_d = calib_buf + (li * 4 + 3) * 4 + _gpu_copy(cs_d, d_scale, 4, stream) + fvk_mod.silu_mul_split_fp8_fp16( + gate_out_ptr, up_out_ptr, int(bufs['gu_fp8']), + Se * Dff, cs_d, int(stream), + ) + + # ── 12. Down FP8 GEMM ── + alpha_d = float(np.float32(as_d) * np.float32(ws_host[li * 4 + 3])) + gemm.fp8_nn_bias( + int(bufs['gu_fp8']), int(weights['d_w'][li]), o_proj_out_ptr, + zero_bias_ptr, + Se, D, Dff, alpha_d, int(stream), + ) + + # ── 13. Residual 2 ── + fvk_mod.residual_add_fp16(x_ptr, o_proj_out_ptr, Se * D, int(stream)) + + # ── Final RMSNorm ── + fvk_mod.rms_norm_fp16( + x_ptr, int(weights['final_norm_w']), int(bufs['xn']), + Se, D, 1e-5, int(stream), + ) + + # ── Copy calibrated scales to output ── + _gpu_copy(calib_scales_ptr, calib_buf, L * 4 * 4, stream) + _gpu_sync(stream) + + +__all__ = [ + "chameleon_forward", + "chameleon_forward_fp16", + "chameleon_forward_calibrate", +] diff --git a/flash_rt/models/chameleon/vqgan/__init__.py b/flash_rt/models/chameleon/vqgan/__init__.py new file mode 100644 index 00000000..02a06b7b --- /dev/null +++ b/flash_rt/models/chameleon/vqgan/__init__.py @@ -0,0 +1,2 @@ +from .image_tokenizer import ImageTokenizer +from .vocab import VocabInfo, VocabTranslation diff --git a/flash_rt/models/chameleon/vqgan/image_tokenizer.py b/flash_rt/models/chameleon/vqgan/image_tokenizer.py new file mode 100644 index 00000000..f857bab9 --- /dev/null +++ b/flash_rt/models/chameleon/vqgan/image_tokenizer.py @@ -0,0 +1,132 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates +# +# This source code is licensed under the Chameleon License found in the +# LICENSE file in the root directory of this source tree. + +import PIL +from PIL import Image +import numpy as np +import torch +import yaml + +from .vqgan import VQModel + + +class ImageTokenizer: + def __init__( + self, + cfg_path: str, + ckpt_path: str, + device: str | torch.device | None = None, + ): + with open(cfg_path) as f: + config = yaml.safe_load(f) + + params = config["model"]["params"] + if "lossconfig" in params: + del params["lossconfig"] + params["ckpt_path"] = ckpt_path + + self._vq_model = VQModel(**params) + self._vq_model.eval() + + if device is None: + devices = {p.device for p in self._vq_model.parameters()} + assert len(devices) == 1 + device = devices.pop() + else: + self._vq_model.to(device) + self._device = device + + dtypes = {p.dtype for p in self._vq_model.parameters()} + assert len(dtypes) == 1 + self._dtype = dtypes.pop() + + def _whiten_transparency(self, img: PIL.Image) -> PIL.Image: + # Check if it's already in RGB format. + if img.mode == "RGB": + return img + + vals_rgba = np.array(img.convert("RGBA")) + + # If there is no transparency layer, simple convert and return. + if not (vals_rgba[:, :, 3] < 255).any(): + return img.convert("RGB") + + # There is a transparency layer, blend it with a white background. + + # Calculate the alpha proportion for blending. + alpha = vals_rgba[:, :, 3] / 255.0 + # Blend with white background. + vals_rgb = (1 - alpha[:, :, np.newaxis]) * 255 + alpha[:, :, np.newaxis] * vals_rgba[:, :, :3] + return PIL.Image.fromarray(vals_rgb.astype("uint8"), "RGB") + + # def _vqgan_input_from(self, img: PIL.Image, target_image_size=512) -> torch.Tensor: + # # Resize with aspect ratio preservation. + # s = min(img.size) + # scale = target_image_size / s + # new_size = (round(scale * img.size[0]), round(scale * img.size[1])) + # img = img.resize(new_size, PIL.Image.LANCZOS) + # + # # Center crop. + # x0 = (img.width - target_image_size) // 2 + # y0 = (img.height - target_image_size) // 2 + # img = img.crop((x0, y0, x0 + target_image_size, y0 + target_image_size)) + # + # # Convert to tensor. + # np_img = np.array(img) / 255.0 # Normalize to [0, 1] + # np_img = np_img * 2 - 1 # Scale to [-1, 1] + # tensor_img = torch.from_numpy(np_img).permute(2, 0, 1).float() # (Channels, Height, Width) format. + # + # # Add batch dimension. + # return tensor_img.unsqueeze(0) + + def img_tokens_from_pil(self, img: PIL.Image) -> list[int]: + img = self._whiten_transparency(img) + # Convert to tensor. + np_img = np.array(img) / 255.0 # Normalize to [0, 1] + np_img = np_img * 2 - 1 # Scale to [-1, 1] + img = torch.from_numpy(np_img).permute(2, 0, 1).to(self._vq_model.encoder.conv_in.weight) + img = img.unsqueeze(0) + + _, _, [_, _, img_toks] = self._vq_model.encode(img) + return img_toks + + def _pil_from_chw_tensor(self, chw_tensor: torch.Tensor) -> PIL.Image: + # Ensure detachment and move tensor to CPU. + detached_chw_tensor = chw_tensor.detach().cpu() + + # Normalize tensor to [0, 1] range from [-1, 1] range. + normalized_chw_tensor = (torch.clamp(detached_chw_tensor, -1.0, 1.0) + 1.0) / 2.0 + + # Permute CHW tensor to HWC format and convert to NumPy array. + hwc_array = normalized_chw_tensor.permute(1, 2, 0).numpy() + + # Convert to an 8-bit unsigned integer format. + image_array_uint8 = (hwc_array * 255).astype(np.uint8) + + # Convert NumPy array to PIL Image. + pil_image = Image.fromarray(image_array_uint8) + + # Convert image to RGB if it is not already. + if pil_image.mode != "RGB": + pil_image = pil_image.convert("RGB") + + return pil_image + + def pil_from_img_toks(self, tokens: torch.Tensor, h_latent_dim=32, w_latent_dim=32) -> PIL.Image: + emb_dim = self._vq_model.quantize.embedding.weight.shape[-1] + codebook_entry = self._vq_model.quantize.get_codebook_entry(tokens, (1, h_latent_dim, w_latent_dim, emb_dim)) + pixels = self._vq_model.decode(codebook_entry) + return self._pil_from_chw_tensor(pixels[0]) + + def latent_embedding_from_pil(self, img: PIL.Image): + img = self._whiten_transparency(img) + + # Convert to tensor. + np_img = np.array(img) / 255.0 # Normalize to [0, 1] + np_img = np_img * 2 - 1 # Scale to [-1, 1] + img = torch.from_numpy(np_img).permute(2, 0, 1) # (Channels, Height, Width) format. + img = img.unsqueeze(0).to(self._vq_model.encoder.conv_in.weight) + latent_embedding, _, _ = self._vq_model.encode(img) + return latent_embedding diff --git a/flash_rt/models/chameleon/vqgan/vocab.py b/flash_rt/models/chameleon/vqgan/vocab.py new file mode 100644 index 00000000..16e39cc0 --- /dev/null +++ b/flash_rt/models/chameleon/vqgan/vocab.py @@ -0,0 +1,107 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the Chameleon License found in the +# LICENSE file in the root directory of this source tree. + +from functools import cached_property + +import torch + + +class VocabInfo: + def __init__(self, vocab_map: dict[str, int]): + self.name2val = vocab_map + + self.bos_id = vocab_map.get("") + self.eos_id = vocab_map.get("") + self.boi_id = vocab_map.get("") + self.eoi_id = vocab_map.get("") + self.pad_id = vocab_map.get("") + self.eot_id = vocab_map.get("") + + @property + def begin_sequence(self) -> int: + return self.bos_id + + @property + def end_sequence(self) -> int: + return self.eos_id + + @property + def begin_image(self) -> int: + return self.boi_id + + @property + def end_image(self) -> int: + return self.eoi_id + + @property + def padding(self) -> int: + return self.pad_id + + @property + def end_turn(self) -> int: + return self.eot_id + + @cached_property + def val2name(self) -> dict[int, str]: + return {v: k for k, v in self.name2val.items()} + + @cached_property + def all_tokens(self) -> list[int]: + return sorted(self.name2val.values()) + + @cached_property + def image_tokens(self) -> list[int]: + return sorted([val for name, val in self.name2val.items() if name.startswith("IMGIMG")]) + + @cached_property + def special_tokens(self) -> list[int]: + return sorted([val for name, val in self.name2val.items() if name.startswith("<") and name != "<"]) + + @cached_property + def text_tokens(self) -> list[int]: + return sorted(set(self.all_tokens) - set(self.image_tokens) - set(self.special_tokens)) + + +class VocabTranslation: + def __init__(self, vocab_info: VocabInfo, device: str | None = None): + self._vocab = vocab_info + self._device = device + + @cached_property + def bpe2img(self) -> dict[int, int]: + img_tkn_chr_mapping = {chr(ord("A") + i): str(i) for i in range(10)} + + def remap(old_name: str) -> str: + return "".join(img_tkn_chr_mapping.get(c, c) for c in old_name[len("IMGIMG") : -1]) + + return {tok: int(remap(self._vocab.val2name[tok])) for tok in self._vocab.image_tokens} + + @cached_property + def img2bpe(self) -> dict[int, int]: + return {v: k for k, v in self.bpe2img.items()} + + @cached_property + def bpe2img_search_tensors(self) -> tuple[torch.Tensor, torch.Tensor]: + sorted_bpe = torch.tensor(sorted(self.bpe2img.keys()), device=self._device) + sorted_img = torch.tensor(sorted(self.bpe2img.values()), device=self._device) + return sorted_bpe, sorted_img + + @cached_property + def img2bpe_mapping_tensor(self) -> torch.LongTensor: + mapping = torch.zeros( + max(self.img2bpe.keys()) + 1, + dtype=torch.int, + device=self._device, + ) + for k, v in self.img2bpe.items(): + mapping[k] = v + return mapping + + def convert_bpe2img(self, bpe_batch: torch.Tensor) -> torch.Tensor: + bpe_tok, img_tok = self.bpe2img_search_tensors + return img_tok[torch.searchsorted(bpe_tok, bpe_batch)] + + def convert_img2bp2(self, img_batch: torch.Tensor) -> torch.Tensor: + return self.img2bpe_mapping_tensor[img_batch] diff --git a/flash_rt/models/chameleon/vqgan/vqgan.py b/flash_rt/models/chameleon/vqgan/vqgan.py new file mode 100644 index 00000000..a78437a2 --- /dev/null +++ b/flash_rt/models/chameleon/vqgan/vqgan.py @@ -0,0 +1,634 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. + +# This source code is licensed under the Chameleon License found in the +# LICENSE file in the root directory of this source tree. + +""" +Contents of this file are taken from https://github.com/CompVis/taming-transformers/blob/3ba01b241669f5ade541ce990f7650a3b8f65318/taming/models/vqgan.py +[with minimal dependencies] + +This implementation is inference-only -- training steps and optimizer components +introduce significant additional dependencies +""" + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class VectorQuantizer2(nn.Module): + """ + Improved version over VectorQuantizer, can be used as a drop-in replacement. Mostly + avoids costly matrix multiplications and allows for post-hoc remapping of indices. + """ + + # NOTE: due to a bug the beta term was applied to the wrong term. for + # backwards compatibility we use the buggy version by default, but you can + # specify legacy=False to fix it. + def __init__( + self, + n_e, + e_dim, + beta, + remap=None, + unknown_index="random", + sane_index_shape=False, + legacy=True, + ): + super().__init__() + self.n_e = n_e + self.e_dim = e_dim + self.beta = beta + self.legacy = legacy + + self.embedding = nn.Embedding(self.n_e, self.e_dim) + self.embedding.weight.data.uniform_(-1.0 / self.n_e, 1.0 / self.n_e) + + self.remap = remap + if self.remap is not None: + self.register_buffer("used", torch.tensor(np.load(self.remap))) + self.re_embed = self.used.shape[0] + self.unknown_index = unknown_index # "random" or "extra" or integer + if self.unknown_index == "extra": + self.unknown_index = self.re_embed + self.re_embed = self.re_embed + 1 + print( + f"Remapping {self.n_e} indices to {self.re_embed} indices. " + f"Using {self.unknown_index} for unknown indices." + ) + else: + self.re_embed = n_e + + self.sane_index_shape = sane_index_shape + + def remap_to_used(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + match = (inds[:, :, None] == used[None, None, ...]).long() + new = match.argmax(-1) + unknown = match.sum(2) < 1 + if self.unknown_index == "random": + new[unknown] = torch.randint(0, self.re_embed, size=new[unknown].shape).to(device=new.device) + else: + new[unknown] = self.unknown_index + return new.reshape(ishape) + + def unmap_to_all(self, inds): + ishape = inds.shape + assert len(ishape) > 1 + inds = inds.reshape(ishape[0], -1) + used = self.used.to(inds) + if self.re_embed > self.used.shape[0]: # extra token + inds[inds >= self.used.shape[0]] = 0 # simply set to zero + back = torch.gather(used[None, :][inds.shape[0] * [0], :], 1, inds) + return back.reshape(ishape) + + def forward(self, z, temp=None, rescale_logits=False, return_logits=False): + assert temp is None or temp == 1.0, "Only for interface compatible with Gumbel" + assert rescale_logits is False, "Only for interface compatible with Gumbel" + assert return_logits is False, "Only for interface compatible with Gumbel" + # reshape z -> (batch, height, width, channel) and flatten + z = z.permute(0, 2, 3, 1).contiguous() + z_flattened = z.view(-1, self.e_dim) + # distances from z to embeddings e_j (z - e)^2 = z^2 + e^2 - 2 e * z + + d = ( + torch.sum(z_flattened**2, dim=1, keepdim=True) + + torch.sum(self.embedding.weight**2, dim=1) + - 2 * torch.einsum("bd,dn->bn", z_flattened, self.embedding.weight.transpose(0, 1)) + ) + + min_encoding_indices = torch.argmin(d, dim=1) + z_q = self.embedding(min_encoding_indices).view(z.shape) + perplexity = None + min_encodings = None + + # compute loss for embedding + if not self.legacy: + loss = self.beta * torch.mean((z_q.detach() - z) ** 2) + torch.mean((z_q - z.detach()) ** 2) + else: + loss = torch.mean((z_q.detach() - z) ** 2) + self.beta * torch.mean((z_q - z.detach()) ** 2) + + # preserve gradients + z_q = z + (z_q - z).detach() + + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + if self.remap is not None: + min_encoding_indices = min_encoding_indices.reshape(z.shape[0], -1) # add batch axis + min_encoding_indices = self.remap_to_used(min_encoding_indices) + min_encoding_indices = min_encoding_indices.reshape(-1, 1) # flatten + + if self.sane_index_shape: + min_encoding_indices = min_encoding_indices.reshape(z_q.shape[0], z_q.shape[2], z_q.shape[3]) + + return z_q, loss, (perplexity, min_encodings, min_encoding_indices) + + def get_codebook_entry(self, indices, shape): + # shape specifying (batch, height, width, channel) + if self.remap is not None: + indices = indices.reshape(shape[0], -1) # add batch axis + indices = self.unmap_to_all(indices) + indices = indices.reshape(-1) # flatten again + + # get quantized latent vectors + z_q = self.embedding(indices) + + if shape is not None: + z_q = z_q.view(shape) + # reshape back to match original input shape + z_q = z_q.permute(0, 3, 1, 2).contiguous() + + return z_q + + +# Alias +VectorQuantizer = VectorQuantizer2 + + +def nonlinearity(x): + # swish + return x * torch.sigmoid(x) + + +def Normalize(in_channels, num_groups=32): + return torch.nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + +class Upsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + x = F.interpolate(x, scale_factor=2.0, mode="nearest") + if self.with_conv: + x = self.conv(x) + return x + + +class Downsample(nn.Module): + def __init__(self, in_channels, with_conv): + super().__init__() + self.with_conv = with_conv + if self.with_conv: + # no asymmetric padding in torch conv, must do it ourselves + self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x): + if self.with_conv: + pad = (0, 1, 0, 1) + x = F.pad(x, pad, mode="constant", value=0) + x = self.conv(x) + else: + x = F.avg_pool2d(x, kernel_size=2, stride=2) + return x + + +class ResnetBlock(nn.Module): + def __init__( + self, + *, + in_channels, + out_channels=None, + conv_shortcut=False, + dropout, + temb_channels=512, + ): + super().__init__() + self.in_channels = in_channels + out_channels = in_channels if out_channels is None else out_channels + self.out_channels = out_channels + self.use_conv_shortcut = conv_shortcut + + self.norm1 = Normalize(in_channels) + self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + if temb_channels > 0: + self.temb_proj = torch.nn.Linear(temb_channels, out_channels) + self.norm2 = Normalize(out_channels) + self.dropout = torch.nn.Dropout(dropout) + self.conv2 = torch.nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + self.conv_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + else: + self.nin_shortcut = torch.nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x, temb): + h = x + h = self.norm1(h) + h = nonlinearity(h) + h = self.conv1(h) + + if temb is not None: + h = h + self.temb_proj(nonlinearity(temb))[:, :, None, None] + + h = self.norm2(h) + h = nonlinearity(h) + h = self.dropout(h) + h = self.conv2(h) + + if self.in_channels != self.out_channels: + if self.use_conv_shortcut: + x = self.conv_shortcut(x) + else: + x = self.nin_shortcut(x) + + return x + h + + +class AttnBlock(nn.Module): + def __init__(self, in_channels): + super().__init__() + self.in_channels = in_channels + + self.norm = Normalize(in_channels) + self.q = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = torch.nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x): + h_ = x + h_ = self.norm(h_) + q = self.q(h_) + k = self.k(h_) + v = self.v(h_) + + # compute attention + b, c, h, w = q.shape + q = q.reshape(b, c, h * w) + q = q.permute(0, 2, 1) # b,hw,c + k = k.reshape(b, c, h * w) # b,c,hw + w_ = torch.bmm(q, k) # b,hw,hw w[b,i,j]=sum_c q[b,i,c]k[b,c,j] + w_ = w_ * (int(c) ** (-0.5)) + w_ = F.softmax(w_, dim=2) + + # attend to values + v = v.reshape(b, c, h * w) + w_ = w_.permute(0, 2, 1) # b,hw,hw (first hw of k, second of q) + h_ = torch.bmm(v, w_) # b, c,hw (hw of q) h_[b,c,j] = sum_i v[b,c,i] w_[b,i,j] + h_ = h_.reshape(b, c, h, w) + + h_ = self.proj_out(h_) + + return x + h_ + + +def make_attn(in_channels, attn_type="vanilla"): + assert attn_type in ["vanilla", "linear", "none"], f"attn_type {attn_type} unknown" + # print(f"making attention of type '{attn_type}' with {in_channels} in_channels") + if attn_type == "vanilla": + return AttnBlock(in_channels) + elif attn_type == "none": + return nn.Identity(in_channels) + else: + raise ValueError("Unexpected attention type") + + +class Encoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + double_z=True, + use_linear_attn=False, + attn_type="vanilla", + **ignore_kwargs, + ): + super().__init__() + if use_linear_attn: + attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + + # downsampling + self.conv_in = torch.nn.Conv2d(in_channels, self.ch, kernel_size=3, stride=1, padding=1) + + curr_res = resolution + in_ch_mult = (1,) + tuple(ch_mult) + self.in_ch_mult = in_ch_mult + self.down = nn.ModuleList() + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = ch * in_ch_mult[i_level] + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks): + block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level != self.num_resolutions - 1: + down.downsample = Downsample(block_in, resamp_with_conv) + curr_res = curr_res // 2 + self.down.append(down) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d( + block_in, + 2 * z_channels if double_z else z_channels, + kernel_size=3, + stride=1, + padding=1, + ) + + def forward(self, x): + # timestep embedding + temb = None + + # downsampling + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1], temb) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level != self.num_resolutions - 1: + hs.append(self.down[i_level].downsample(hs[-1])) + + # middle + h = hs[-1] + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # end + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + return h + + +class Decoder(nn.Module): + def __init__( + self, + *, + ch, + out_ch, + ch_mult=(1, 2, 4, 8), + num_res_blocks, + attn_resolutions, + dropout=0.0, + resamp_with_conv=True, + in_channels, + resolution, + z_channels, + give_pre_end=False, + tanh_out=False, + use_linear_attn=False, + attn_type="vanilla", + **ignorekwargs, + ): + super().__init__() + if use_linear_attn: + attn_type = "linear" + self.ch = ch + self.temb_ch = 0 + self.num_resolutions = len(ch_mult) + self.num_res_blocks = num_res_blocks + self.resolution = resolution + self.in_channels = in_channels + self.give_pre_end = give_pre_end + self.tanh_out = tanh_out + + # compute in_ch_mult, block_in and curr_res at lowest res + block_in = ch * ch_mult[self.num_resolutions - 1] + curr_res = resolution // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + + # z to block_in + self.conv_in = torch.nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + # middle + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + self.mid.attn_1 = make_attn(block_in, attn_type=attn_type) + self.mid.block_2 = ResnetBlock( + in_channels=block_in, + out_channels=block_in, + temb_channels=self.temb_ch, + dropout=dropout, + ) + + # upsampling + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = ch * ch_mult[i_level] + for i_block in range(self.num_res_blocks + 1): + block.append( + ResnetBlock( + in_channels=block_in, + out_channels=block_out, + temb_channels=self.temb_ch, + dropout=dropout, + ) + ) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(make_attn(block_in, attn_type=attn_type)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level != 0: + up.upsample = Upsample(block_in, resamp_with_conv) + curr_res = curr_res * 2 + self.up.insert(0, up) # prepend to get consistent order + + # end + self.norm_out = Normalize(block_in) + self.conv_out = torch.nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z): + # assert z.shape[1:] == self.z_shape[1:] + self.last_z_shape = z.shape + + # timestep embedding + temb = None + + # z to block_in + h = self.conv_in(z) + + # middle + h = self.mid.block_1(h, temb) + h = self.mid.attn_1(h) + h = self.mid.block_2(h, temb) + + # upsampling + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h, temb) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level != 0: + h = self.up[i_level].upsample(h) + + # end + if self.give_pre_end: + return h + + h = self.norm_out(h) + h = nonlinearity(h) + h = self.conv_out(h) + if self.tanh_out: + h = torch.tanh(h) + return h + + +class VQModel(nn.Module): + def __init__( + self, + ddconfig, + n_embed, + embed_dim, + ckpt_path=None, + ignore_keys=[], + image_key="image", + colorize_nlabels=None, + monitor=None, + scheduler_config=None, + lr_g_factor=1.0, + remap=None, + sane_index_shape=False, # tell vector quantizer to return indices as bhw + ): + super().__init__() + self.image_key = image_key + self.encoder = Encoder(**ddconfig) + self.decoder = Decoder(**ddconfig) + self.quantize = VectorQuantizer( + n_embed, + embed_dim, + beta=0.25, + remap=remap, + sane_index_shape=sane_index_shape, + ) + self.quant_conv = torch.nn.Conv2d(ddconfig["z_channels"], embed_dim, 1) + self.post_quant_conv = torch.nn.Conv2d(embed_dim, ddconfig["z_channels"], 1) + if ckpt_path is not None: + self.init_from_ckpt(ckpt_path, ignore_keys=ignore_keys) + self.image_key = image_key + if colorize_nlabels is not None: + assert isinstance(colorize_nlabels, int) + self.register_buffer("colorize", torch.randn(3, colorize_nlabels, 1, 1)) + if monitor is not None: + self.monitor = monitor + self.scheduler_config = scheduler_config + self.lr_g_factor = lr_g_factor + + def init_from_ckpt(self, path, ignore_keys=list()): + sd = torch.load(path, map_location="cpu")["state_dict"] + keys = list(sd.keys()) + for k in keys: + for ik in ignore_keys: + if k.startswith(ik): + print("Deleting key {} from state_dict.".format(k)) + del sd[k] + self.load_state_dict(sd, strict=False) + print(f"VQModel loaded from {path}") + + def encode(self, x): + h = self.encoder(x) + h = self.quant_conv(h) + quant, emb_loss, info = self.quantize(h) + return quant, emb_loss, info + + def decode(self, quant): + quant = self.post_quant_conv(quant) + dec = self.decoder(quant) + return dec + + def decode_code(self, code_b): + quant_b = self.quantize.embed_code(code_b) + dec = self.decode(quant_b) + return dec + + def forward(self, input): + quant, diff, _ = self.encode(input) + dec = self.decode(quant) + return dec, diff + + def get_input(self, batch, k): + x = batch[k] + if len(x.shape) == 3: + x = x[..., None] + x = x.permute(0, 3, 1, 2).to(memory_format=torch.contiguous_format) + return x.float() + + def get_last_layer(self): + return self.decoder.conv_out.weight + + def log_images(self, batch, **kwargs): + log = dict() + x = self.get_input(batch, self.image_key) + x = x.to(self.device) + xrec, _ = self(x) + if x.shape[1] > 3: + # colorize with random projection + assert xrec.shape[1] > 3 + x = self.to_rgb(x) + xrec = self.to_rgb(xrec) + log["inputs"] = x + log["reconstructions"] = xrec + return log + + def to_rgb(self, x): + assert self.image_key == "segmentation" + if not hasattr(self, "colorize"): + self.register_buffer("colorize", torch.randn(3, x.shape[1], 1, 1).to(x)) + x = F.conv2d(x, weight=self.colorize) + x = 2.0 * (x - x.min()) / (x.max() - x.min()) - 1.0 + return x From 694b1d4e3860341b295d0e1b65daba0205b4e7e3 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:24:53 +0000 Subject: [PATCH 2/7] feat(chameleon): Orin SM87 INT8/QuaRot frontend Jetson AGX Orin (SM87) Chameleon-7B path aligned with the upstream rtx_sm87 naming: - flash_rt/models/chameleon/pipeline_rtx.py: one chameleon_forward serving prefill and decode on the SM80 CUTLASS INT8/INT4 rowwise GEMMs with QuaRot-Hadamard rotations (correctness requirement, not an optimization) and the ffn_down clamp on the last 4 layers (FP16 65504 overflow guard). - _chameleon_quant.py: INT8/INT4 weight quantization + Hadamard packing from the BF16 checkpoint. - _chameleon_spec.py: declarative weight spec with an inlined, Chameleon-specific _llm_block (no bias terms, no FP8 scales). - ChameleonTorchFrontendRtxSm87 (chameleon_rtx_sm87.py): set_prompt / prefill / decode_step / generate, FLASHRT_CHAMELEON_SM87_FORCE escape hatch. - hardware/rtx/attn_backend_chameleon.py: FA2 fwd_fp16_causal is mandatory for decode (bottom-right causal semantics); the backend raises rather than falling back to a top-left cuBLAS mask, which would be silently wrong. Runtime numbers (21.07 tok/s, 16/16 bit-identical greedy vs HF BF16) were measured on Orin hardware in the derivative repo and still need SM87 validation here. --- flash_rt/frontends/torch/_chameleon_quant.py | 251 ++++++ flash_rt/frontends/torch/_chameleon_spec.py | 85 +++ .../frontends/torch/chameleon_rtx_sm87.py | 716 ++++++++++++++++++ .../hardware/rtx/attn_backend_chameleon.py | 237 ++++++ flash_rt/models/chameleon/pipeline_rtx.py | 305 ++++++++ 5 files changed, 1594 insertions(+) create mode 100644 flash_rt/frontends/torch/_chameleon_quant.py create mode 100644 flash_rt/frontends/torch/_chameleon_spec.py create mode 100644 flash_rt/frontends/torch/chameleon_rtx_sm87.py create mode 100644 flash_rt/hardware/rtx/attn_backend_chameleon.py create mode 100644 flash_rt/models/chameleon/pipeline_rtx.py diff --git a/flash_rt/frontends/torch/_chameleon_quant.py b/flash_rt/frontends/torch/_chameleon_quant.py new file mode 100644 index 00000000..0a1a1229 --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_quant.py @@ -0,0 +1,251 @@ +"""Low-bit weight quantizers for the Chameleon-7B GEMMs (Orin SM87). + +Checkpoint-agnostic and frontend-agnostic: every function takes plain tensor +lists and returns plain tensors, so the same code serves any Chameleon +frontend. + +Layout contract — get this wrong and the GEMM silently returns garbage: + +* The declarative weight spec hands us per-projection FP16 weights in + ``[K, N]`` row-major (``Cat``/``FusedGateUp`` followed by ``T()``). +* The CUTLASS SM80 rowwise GEMMs consume the **B** operand as ``[N, K]`` + ColumnMajor with an ``[N]`` FP32 ``RowBroadcast`` scale, so we transpose to + ``[N, K]`` and quantize each of the N output rows symmetrically. +* The **A** operand is ``[M, K]`` RowMajor with ``M`` consecutive FP32 + ``ColBroadcast`` scales, produced at runtime by the fused norm/quant kernels. + +INT4 additionally applies the QuaRot rotation ``W_rot = (H_K @ W)/sqrt(K)`` +offline; the matching activation rotation happens online in the fused +RMSNorm+FHT kernels. Values are packed 2/byte with the even index in the low +nibble (``cutlass::int4b_t`` order). +""" + +from __future__ import annotations + +import logging +from typing import Dict, List, Tuple + +import torch + +logger = logging.getLogger(__name__) + +INT8_QUANT_MAX = 127.0 +INT8_QUANT_EPS = 1e-12 +INT4_QUANT_MAX = 7.0 +INT4_QUANT_EPS = 1e-10 + +#: The seven per-layer GEMMs, in the order the pipeline consumes them. +PROJECTIONS = ("q", "k", "v", "o", "gate", "up", "d") + +#: Projections whose K == D (a power of two) and so can take a full-width +#: Hadamard rotation. ``d`` has K == Dff == 11008 and needs block-H128. +INT4_POW2_PROJECTIONS = ("q", "k", "v", "o", "gate", "up") + + +def split_fused_projections(qkv_w: List[torch.Tensor], + gu_w: List[torch.Tensor], + o_w: List[torch.Tensor], + d_w: List[torch.Tensor], + *, D: int, Dff: int) -> Dict[str, List[torch.Tensor]]: + """Materialize the seven per-projection ``[K, N]`` weight lists. + + ``qkv_w[li]`` is ``[D, 3D]`` row-major and ``gu_w[li]`` is ``[D, 2*Dff]``: + after ``Cat(dim=0) -> T().contiguous()`` the fused row stride is ``3D`` + (resp. ``2*Dff``), so a *byte-offset* split would read the projections + column-interleaved. Only a ``fused[:, lo:hi].contiguous()`` slice recovers + the original ``q_proj`` / ``k_proj`` / ``v_proj`` blocks. + """ + out: Dict[str, List[torch.Tensor]] = {k: [] for k in PROJECTIONS} + for li, (qkv, gu) in enumerate(zip(qkv_w, gu_w)): + if tuple(qkv.shape) != (D, 3 * D): + raise RuntimeError( + f"layer {li}: expected fused qkv_w {(D, 3 * D)}, " + f"got {tuple(qkv.shape)}") + if tuple(gu.shape) != (D, 2 * Dff): + raise RuntimeError( + f"layer {li}: expected fused gu_w {(D, 2 * Dff)}, " + f"got {tuple(gu.shape)}") + out["q"].append(qkv[:, 0:D].contiguous()) + out["k"].append(qkv[:, D:2 * D].contiguous()) + out["v"].append(qkv[:, 2 * D:3 * D].contiguous()) + out["gate"].append(gu[:, 0:Dff].contiguous()) + out["up"].append(gu[:, Dff:2 * Dff].contiguous()) + out["o"] = list(o_w) + out["d"] = list(d_w) + return out + + +def quantize_per_row_int8(w_kn: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-output-row symmetric INT8 of a ``[K, N]`` FP16 weight. + + Returns ``(q [N, K] int8, scale [N] fp32)``, both contiguous on the + weight's device. + """ + w_f32 = w_kn.float().transpose(0, 1).contiguous() # [N, K] + scale = torch.clamp(w_f32.abs().amax(dim=1) / INT8_QUANT_MAX, + min=INT8_QUANT_EPS).float().contiguous() # [N] + q = torch.clamp(torch.round(w_f32 / scale[:, None]), + -127, 127).to(torch.int8).contiguous() # [N, K] + return q, scale + + +def hadamard_gpu(n: int, device="cuda") -> torch.Tensor: + """Unnormalised Sylvester Hadamard ``H_n`` (fp32). ``n`` must be a power of 2.""" + H = torch.ones(1, 1, dtype=torch.float32, device=device) + while H.shape[0] < n: + H = torch.cat([torch.cat([H, H], 1), torch.cat([H, -H], 1)], 0) + return H + + +def _pack_int4_rows(w_rot_nk: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-output-row symmetric INT4 of an already-rotated ``[N, K]`` fp32 weight.""" + scale = (w_rot_nk.abs().amax(1) / INT4_QUANT_MAX).clamp_min(INT4_QUANT_EPS) + q = torch.clamp(torch.round(w_rot_nk / scale[:, None]), -7, 7).to(torch.int8) + lo = (q[:, 0::2] & 0xF).to(torch.uint8) + hi = (q[:, 1::2] & 0xF).to(torch.uint8) + return (lo | (hi << 4)).contiguous(), scale.float().contiguous() + + +class QuantizedWeights: + """Owns the quantized tensors and exposes the pointer dicts the pipeline wants. + + ``ptr[proj][li]`` is the packed weight pointer and ``scale_ptr[proj][li]`` + the matching ``[N]`` FP32 scale pointer. The tensors are kept alive in + ``_store`` for the object's lifetime — dropping this object invalidates + every pointer, including any already baked into a captured CUDA graph. + """ + + def __init__(self) -> None: + self._store: List[torch.Tensor] = [] + self.ptr: Dict[str, List[int]] = {k: [] for k in PROJECTIONS} + self.scale_ptr: Dict[str, List[int]] = {k: [] for k in PROJECTIONS} + self.precision: Dict[str, str] = {} + + def bytes(self) -> int: + return sum(t.numel() * t.element_size() for t in self._store) + + def _set(self, proj: str, li: int, q: torch.Tensor, s: torch.Tensor) -> None: + self._store.append(q) + self._store.append(s) + while len(self.ptr[proj]) <= li: + self.ptr[proj].append(0) + self.scale_ptr[proj].append(0) + self.ptr[proj][li] = q.data_ptr() + self.scale_ptr[proj][li] = s.data_ptr() + + def _drop(self, ptr: int) -> None: + if ptr: + self._store = [t for t in self._store if t.data_ptr() != ptr] + + +def quantize_int8_all(proj: Dict[str, List[torch.Tensor]], + *, num_layers: int) -> QuantizedWeights: + """INT8-quantize all seven projections for all layers.""" + qw = QuantizedWeights() + for li in range(num_layers): + for key in PROJECTIONS: + w = proj[key][li] + if w.dtype != torch.float16: + w = w.to(torch.float16) + qw._set(key, li, *quantize_per_row_int8(w)) + qw.precision.update({k: "int8" for k in PROJECTIONS}) + logger.info("INT8 quantized %d LLM GEMM weights (%d layers x %d proj), %.2f GB", + num_layers * len(PROJECTIONS), num_layers, len(PROJECTIONS), + qw.bytes() / 2 ** 30) + return qw + + +def quantize_int4_quarot(qw: QuantizedWeights, + proj: Dict[str, List[torch.Tensor]], + *, num_layers: int, D: int, + include_down: bool = False) -> QuantizedWeights: + """Replace the INT8 tensors of the rotatable projections with QuaRot INT4. + + ``include_down`` additionally rotates the FFN down projection with a + **block-diagonal** ``H_128`` because its K (11008) is not a power of two. + Unrotated per-row W4A4 measures cos 0.9494 and fails the gate, so the + rotation is mandatory rather than an optimization. + """ + Hm = hadamard_gpu(D) / (float(D) ** 0.5) + keys = list(INT4_POW2_PROJECTIONS) + Hb = None + if include_down: + keys.append("d") + Hb = hadamard_gpu(128) / (128.0 ** 0.5) + + for li in range(num_layers): + for key in keys: + w = proj[key][li].to(torch.float32) # [K, N] + if key == "d": + Kd = w.shape[0] + w_rot = (Hb.t() @ w.reshape(Kd // 128, 128, -1) + ).reshape(Kd, -1).t().contiguous() # [N, Kd] + else: + if w.shape[0] != D: + raise RuntimeError( + f"{key} layer {li}: expected K={D}, got {tuple(w.shape)}") + w_rot = (Hm @ w).t().contiguous() # [N, D] + qw._drop(qw.ptr[key][li]) + qw._set(key, li, *_pack_int4_rows(w_rot)) + qw.precision[key] = "int4" + del w, w_rot + del Hm, Hb + torch.cuda.empty_cache() + logger.info("QuaRot INT4: rotated+packed %d weights (%d layers x %d proj)%s", + num_layers * len(keys), num_layers, len(keys), + "" if include_down else "; down stays INT8") + return qw + + +def quantize_int8_hadamard(qw: QuantizedWeights, + proj: Dict[str, List[torch.Tensor]], + *, num_layers: int, D: int) -> QuantizedWeights: + """Re-quantize the rotatable projections as **Hadamard-rotated INT8** (W8A8+QuaRot). + + This configuration sits between two tiers: + + * plain per-row INT8 — 8-bit resolution, but *unconditioned*, so a row whose + amax is set by a massive-activation channel loses its remaining ~4090 + channels to rounding; + * QuaRot INT4 — conditioned by the rotation, but only 15 levels. + + Rotating at 8 bits gets both. The rotation is free at inference time: the + weight side folds offline here (``W_rot = H·W/sqrt(K)``) and the activation + side is fused into the norm kernels (``rms_norm_fht_int8_fp16`` and friends). + + Crucially it keeps **plain per-row scales**, so the unmodified + ``cutlass_int8_rowwise_*`` GEMMs are reused — the alternative outlier fixes + (group-128 / block-scaled) would each need a bespoke GEMM, and the measured + ceiling for a hand-written block-scaled s4 kernel on 16-SM Orin was only + 41 TOPS. Principle #17: pick the rotation that keeps you on the fast path. + + The FFN down projection is left as plain INT8: its K (11008) is not a power + of two, and its input is the un-rotated BF16 SiLU output. + """ + Hm = hadamard_gpu(D) / (float(D) ** 0.5) + for li in range(num_layers): + for key in INT4_POW2_PROJECTIONS: + w = proj[key][li].to(torch.float32) # [K, N] + if w.shape[0] != D: + raise RuntimeError( + f"{key} layer {li}: expected K={D}, got {tuple(w.shape)}") + w_rot = (Hm @ w).half() # [K, N] + qw._drop(qw.ptr[key][li]) + qw._set(key, li, *quantize_per_row_int8(w_rot)) + qw.precision[key] = "int8+hadamard" + del w, w_rot + del Hm + torch.cuda.empty_cache() + logger.info("W8A8+Hadamard: rotated %d weights (%d layers x %d proj); " + "down stays plain INT8", + num_layers * len(INT4_POW2_PROJECTIONS), num_layers, + len(INT4_POW2_PROJECTIONS)) + return qw + + +__all__ = [ + "INT8_QUANT_MAX", "INT4_QUANT_MAX", "PROJECTIONS", "INT4_POW2_PROJECTIONS", + "QuantizedWeights", "split_fused_projections", "quantize_per_row_int8", + "hadamard_gpu", "quantize_int8_all", "quantize_int8_hadamard", + "quantize_int4_quarot", +] diff --git a/flash_rt/frontends/torch/_chameleon_spec.py b/flash_rt/frontends/torch/_chameleon_spec.py new file mode 100644 index 00000000..c5671f80 --- /dev/null +++ b/flash_rt/frontends/torch/_chameleon_spec.py @@ -0,0 +1,85 @@ +"""Declarative weight spec for upstream Chameleon-7B (HF layout). + +Chameleon-7B LLM — 32 layers, MHA (num_kv_heads=32, no interleave). +``attention_bias`` and ``mlp_bias`` are both ``false`` in the checkpoint +config, so no bias items are emitted; the per-head QK-Norm weight *and +bias* tensors are part of the layer block. + +Per-head QK Norm prevents norm_fuse; Cat is used for QKV fusion +(not FusedQKV) because the per-head norms are handled as separate +TensorList items. + +The QK-Norm tensors are loaded verbatim (``ToFp16()`` only), preserving their +``(1, 128)`` shape. That is deliberate: ``qk_norm_rope_fused_fp16`` reads the +weight as a flat ``[head_dim]`` vector shared across heads, which is exactly +what the checkpoint's ``model_parallel_size == 1`` layout means (upstream +expands it with ``repeat_interleave`` at forward time). No reshape is needed. +""" + +from __future__ import annotations + +from flash_rt.executors.weight_loader import Item, LayerBlock, ModelWeightSpec +from flash_rt.executors.torch_weights import Attr, Cat, FusedGateUp, T, TensorList, ToFp16 + + +def _llm_block() -> LayerBlock: + """Chameleon-7B LLM — 32 layers, FP16 backbone (no quantized spec).""" + lp = "model.layers.{i}" + items = [ + # ── Fused QKV (MHA: no interleave, no norm_fuse) ── + Item("qkv_w", + Cat([f"{lp}.self_attn.q_proj.weight", + f"{lp}.self_attn.k_proj.weight", + f"{lp}.self_attn.v_proj.weight"], dim=0), + [T()], + TensorList("_llm_qkv_w")), + # ── O projection ── + Item("o_w", f"{lp}.self_attn.o_proj.weight", + [ToFp16(), T()], + TensorList("_llm_o_w")), + # ── Fused GateUp (no norm_fuse — per-head QK Norm incompatible) ── + Item("gu_w", + FusedGateUp(gate=f"{lp}.mlp.gate_proj.weight", + up=f"{lp}.mlp.up_proj.weight"), + [T()], + TensorList("_llm_gu_w")), + # ── Down projection ── + Item("d_w", f"{lp}.mlp.down_proj.weight", + [ToFp16(), T()], + TensorList("_llm_d_w")), + # ── Layer norms ── + Item("input_ln_w", f"{lp}.input_layernorm.weight", + [ToFp16()], TensorList("_llm_input_ln_w")), + Item("post_ln_w", f"{lp}.post_attention_layernorm.weight", + [ToFp16()], TensorList("_llm_post_ln_w")), + # ── Per-head Q/K Norm ── + Item("q_norm_w", f"{lp}.self_attn.q_norm.weight", + [ToFp16()], TensorList("_llm_q_norm_w")), + Item("q_norm_b", f"{lp}.self_attn.q_norm.bias", + [ToFp16()], TensorList("_llm_q_norm_b")), + Item("k_norm_w", f"{lp}.self_attn.k_norm.weight", + [ToFp16()], TensorList("_llm_k_norm_w")), + Item("k_norm_b", f"{lp}.self_attn.k_norm.bias", + [ToFp16()], TensorList("_llm_k_norm_b")), + ] + return LayerBlock(prefix_fmt="", num_layers=32, items=items, name="llm") + + +def build_spec() -> ModelWeightSpec: + """Chameleon-7B: 32 decoder layers + embedding, final norm, lm_head.""" + return ModelWeightSpec( + framework="torch", + blocks=[_llm_block()], + singletons=[ + Item("embed_w", "model.embed_tokens.weight", + [ToFp16()], Attr("_llm_embed_w")), + Item("norm_w", "model.norm.weight", + [ToFp16()], Attr("_llm_norm_w")), + # Live output projection. + Item("lm_head_w", "lm_head.weight", + [ToFp16()], Attr("_llm_lm_head_w")), + ], + ) + + +__all__ = ["build_spec"] diff --git a/flash_rt/frontends/torch/chameleon_rtx_sm87.py b/flash_rt/frontends/torch/chameleon_rtx_sm87.py new file mode 100644 index 00000000..60a7c5fb --- /dev/null +++ b/flash_rt/frontends/torch/chameleon_rtx_sm87.py @@ -0,0 +1,716 @@ +"""FlashRT — upstream Chameleon-7B VLM frontend for Jetson AGX Orin (SM87). + +Image+text -> text. Constructed directly (**not** via ``flash_rt.load_model``): +this is a chat-style VLM exposing ``set_prompt()`` + ``generate()``, whereas +``VLAModel.predict`` unconditionally reads ``result['actions']``. Same precedent +as ``qwen3_vl`` — see the redirect in ``flash_rt/api.py``. + +Production precision policy (all defaults) +------------------------------------------ +* **Q/K/V/O, FFN gate/up** — INT8 W8A8 **+ Hadamard rotation** (QuaRot at 8 + bits): weights rotated offline, activations rotated inside the fused norm + kernel; per-row dynamic activation scales. +* **FFN down** — INT8 W8A8 per-row dynamic. Not rotated: K=11008 is not a power + of two and its input is the un-rotated BF16 SiLU output. +* **lm_head** (65536x4096) — INT8 W8A8. 268 MB/token = 3.7 % of the decode + budget; FP16 would double that for no measured argmax benefit. +* **residual / QK-LayerNorm / RoPE / attention / KV cache** — FP16, with + ``ffn_down_clamp`` applied on the last ``ffn_down_clamp_last_n`` layers. +* **attention** — FA2 fp16 causal; ``split_kv_bias=4`` on decode, because FA2's + own heuristic returns ``num_splits=1`` at Chameleon's 32 Q heads. +* **VQ-GAN encoder** — FP16 convs + **fp32** codebook distance/argmin. + +Measured (ISL=1032, OSL=16, warm): greedy output **bit-identical to the HF bf16 +reference for 16/16 tokens**, worst per-layer cosine 0.9986, last-row logit +cosine 0.999968, **21.07 tok/s** decode, 273.8 ms prefill, 7.6 GB resident. + +Two non-obvious *correctness* requirements +------------------------------------------ +1. **The Hadamard rotation is not an optimization.** Chameleon's + massive-activation channels (measured: L31 row-0 channel d632 at 2.4e4 against + a row median of 1e4) pin the per-row INT8 amax and round that row's other + ~4090 channels to zero. Plain per-row INT8 reproduces only 8/16 reference + tokens; rotating fixes it. It is free because it preserves per-row scales, so + the stock ``cutlass_int8_rowwise_*`` GEMMs are reused unchanged. +2. **``ffn_down_clamp`` is not a tuning knob.** L31's down output reaches 2.6e5, + past FP16's 65504; without the clamp the residual stores ``inf`` and the final + RMSNorm poisons that row. See ``models/chameleon/pipeline_thor.py``. + +Dim policy: backbone dims (32 layers / 4096 / 32 heads / 11008 / vocab 65536) and +the special token ids are **hard-asserted** against ``config.json``, so this +frontend is Chameleon-7B-specific by construction rather than by convention. +Note ``config.json``'s ``bos_token_id`` is stale (says 1, which is ````); +``tokenizer.json`` gives `` = 0`` and that is what the processor emits. + +Preprocessing specifics +----------------------- +Chameleon needs PIL LANCZOS / 512 / ``u8*0.0078-1.0`` -> ``[-1, +0.989]`` +normalization and a bare 1024-token raster per image (no grid/newline +layout); the quantizers live in ``_chameleon_quant``. + +See ``docs/chameleon7b_rtx_sm87.md`` for the roofline, the measured lever menu, +and the dead-ends. +""" + +from __future__ import annotations + +import json +import logging +import os +import time +from pathlib import Path +from typing import List, Optional, Sequence + +import torch + +import flash_rt.flash_rt_kernels as fvk +from flash_rt.frontends.torch import _chameleon_quant as cq +from flash_rt.frontends.torch._chameleon_spec import build_spec +from flash_rt.hardware.rtx.attn_backend_chameleon import ( + ChameleonAttnBackend, make_chameleon_attention_spec) +from flash_rt.models.chameleon.pipeline_rtx import chameleon_forward + +logger = logging.getLogger(__name__) + +_FP16 = torch.float16 +_BF16 = torch.bfloat16 + + +class ChameleonTorchFrontendRtxSm87: + """Chameleon-7B image+text → text on Orin SM87. + + Typical use:: + + f = ChameleonTorchFrontendRtxSm87("/path/to/Chameleon_7B_mGPT") + f.set_prompt("Describe this image.", images=[pil_img]) + print(f.generate(max_new_tokens=32)) + """ + + # Token ids for this checkpoint. Hardcoded so an accelerated deployment does + # not depend on the training code being importable, but asserted against + # config.json at load so a different checkpoint fails loudly (principle #10). + BOS_ID = 0 + EOS_ID = 2 + IMG_PLACEHOLDER_ID = 8711 # + BOI_ID = 8197 # + EOI_ID = 8196 # + SEP_ID = 8710 # + IMG_ID_OFFSET = 4 # token_id = codebook_index + 4 + N_IMG_CODES = 8192 + IMG_TOKENS_PER_VIEW = 1024 # 512/16 = 32 -> 32x32 raster + + # Chameleon-7B backbone dims — hard-asserted against config.json. + D = 4096 + L = 32 + H = 32 + HKV = 32 + HD = 128 + DFF = 11008 + VOCAB = 65536 + + def __init__(self, checkpoint_dir: str, *, + max_seq: int = 2048, + use_int4: bool = False, + use_int4_down: bool = False, + use_hadamard: bool = True, + split_kv_bias: int = 4, + ffn_down_clamp: Optional[float] = None, + vq_argmin_fp32: bool = True, + free_fp16_weights: bool = True, + probe_layers: Optional[Sequence[int]] = None, + **_ignored) -> None: + cc = torch.cuda.get_device_capability(0) + if cc != (8, 7) and os.environ.get("FLASHRT_CHAMELEON_SM87_FORCE") != "1": + raise RuntimeError( + f"ChameleonTorchFrontendRtxSm87 targets SM87 (Orin); found SM{cc[0]}{cc[1]}. " + "The INT8/INT4 path needs ENABLE_SM80_INT8_CUTLASS=ON. Set " + "FLASHRT_CHAMELEON_SM87_FORCE=1 to override.") + + self.checkpoint_dir = Path(checkpoint_dir) + self.max_seq = int(max_seq) + self.use_int4_down = bool(use_int4_down) + self.use_int4 = bool(use_int4) or self.use_int4_down + # W8A8+QuaRot: Hadamard-rotate the six K=4096 projections and quantize + # at 8 bits. Conditions the massive-activation channels (which plain + # per-row INT8 cannot handle) without INT4's noise, and keeps per-row + # scales so the stock CUTLASS GEMM is reused. Ignored on the INT4 tier, + # which already rotates. + self.use_hadamard = bool(use_hadamard) and not self.use_int4 + self.vq_argmin_fp32 = bool(vq_argmin_fp32) + # Required for correctness: L31's down output reaches ~2.6e5, far past + # FP16's 65504. See chameleon_forward's docstring for the measured + # per-layer table and docs/chameleon_acceleration_methodology.md §2.1. + self.ffn_down_clamp = float( + os.environ.get("FLASHRT_CHAMELEON_DOWN_CLAMP", "60000.0") + if ffn_down_clamp is None else ffn_down_clamp) + self.ffn_down_clamp_last_n = int( + os.environ.get("FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N", "4")) + self._probe_layers = list(probe_layers) if probe_layers else [] + + self._validate_config() + t0 = time.perf_counter() + self._load_weights(free_fp16_weights=free_fp16_weights) + self._build_rope_tables() + self._allocate_buffers() + self._build_attn_backend(split_kv_bias=split_kv_bias) + self._load_processor() + self._load_vqgan() + + # Prompt state. + self.input_ids: Optional[torch.Tensor] = None + self.S = 0 + self._prompt_ready = False + self._timing: dict = {} + + logger.info( + "ChameleonTorchFrontendRtxSm87 ready in %.1fs — tier=%s max_seq=%d " + "gpu_mem=%.2f GB", + time.perf_counter() - t0, self.precision_tier, self.max_seq, + torch.cuda.memory_allocated() / 2 ** 30) + + # ================================================================== + # Load + # ================================================================== + + @property + def precision_tier(self) -> str: + if self.use_int4_down: + return "int4+down" + if self.use_int4: + return "int4" + return "int8+hadamard" if self.use_hadamard else "int8" + + def _validate_config(self) -> None: + cfg = json.loads((self.checkpoint_dir / "config.json").read_text()) + expect = { + "hidden_size": self.D, "num_hidden_layers": self.L, + "num_attention_heads": self.H, "num_key_value_heads": self.HKV, + "intermediate_size": self.DFF, "vocab_size": self.VOCAB, + } + for k, want in expect.items(): + got = cfg.get(k) + if got != want: + raise RuntimeError( + f"config.json {k}={got}, this frontend hardcodes {want}. " + "Chameleon-7B only.") + if cfg.get("attention_bias") or cfg.get("mlp_bias"): + raise RuntimeError( + "This frontend assumes attention_bias=false and mlp_bias=false " + "(upstream Chameleon). A biased checkpoint needs the *_bias " + "GEMM entries wired back in.") + if cfg.get("model_parallel_size", 1) != 1: + raise RuntimeError( + f"model_parallel_size={cfg.get('model_parallel_size')}: the " + "QK-Norm params would differ per head group, but " + "qk_norm_rope_fused_fp16 broadcasts one [head_dim] vector " + "across all heads. Only the 7B (mp=1) layout is supported.") + if self.max_seq > cfg.get("max_position_embeddings", 4096): + raise ValueError( + f"max_seq={self.max_seq} exceeds max_position_embeddings=" + f"{cfg['max_position_embeddings']}") + + # Token-id contract. config.json's bos_token_id is stale (says 1, which + # is ); tokenizer.json is authoritative and gives = 0. + vm = cfg.get("vocabulary_map") or {} + for tok, want in (("", self.IMG_PLACEHOLDER_ID), + ("", self.BOI_ID), + ("", self.EOI_ID), + ("", self.SEP_ID)): + if vm.get(tok) != want: + raise RuntimeError( + f"token id mismatch: {tok} is {vm.get(tok)} in config.json, " + f"expected {want}") + img_ids = sorted(v for k, v in vm.items() if k.startswith("IMGIMG")) + if (len(img_ids) != self.N_IMG_CODES + or img_ids[0] != self.IMG_ID_OFFSET + or img_ids[-1] != self.IMG_ID_OFFSET + self.N_IMG_CODES - 1): + raise RuntimeError( + f"expected {self.N_IMG_CODES} contiguous IMGIMG ids starting at " + f"{self.IMG_ID_OFFSET}; got {len(img_ids)} spanning " + f"[{img_ids[0]}, {img_ids[-1]}]") + self._config = cfg + + def _load_weights(self, *, free_fp16_weights: bool) -> None: + from flash_rt.executors.torch_weights import ( + MultiSafetensorsSource, WeightLoader) + + shards = sorted(self.checkpoint_dir.glob("model-*-of-*.safetensors")) + if not shards: + shards = sorted(self.checkpoint_dir.glob("*.safetensors")) + if not shards: + raise FileNotFoundError(f"no safetensors in {self.checkpoint_dir}") + + src = MultiSafetensorsSource([str(p) for p in shards], device="cuda") + WeightLoader(source=src, target=self, spec=build_spec()).run() + del src + + proj = cq.split_fused_projections( + self._llm_qkv_w, self._llm_gu_w, self._llm_o_w, self._llm_d_w, + D=self.D, Dff=self.DFF) + # The fused tensors are dead once split; release before quantizing so + # the peak is (fp16 split + int8) rather than (fused + split + int8). + self._llm_qkv_w = [] + self._llm_gu_w = [] + torch.cuda.empty_cache() + + self.qw = cq.quantize_int8_all(proj, num_layers=self.L) + if self.use_int4: + cq.quantize_int4_quarot(self.qw, proj, num_layers=self.L, D=self.D, + include_down=self.use_int4_down) + elif self.use_hadamard: + cq.quantize_int8_hadamard(self.qw, proj, num_layers=self.L, D=self.D) + + # lm_head: INT8 in both tiers (see pipeline docstring for the ROI). + self._lm_head_q, self._lm_head_s = cq.quantize_per_row_int8( + self._llm_lm_head_w.t().contiguous()) # [V, D] -> [K=D, N=V] + + if free_fp16_weights: + for key in cq.PROJECTIONS: + proj[key] = [] + self._llm_o_w = [] + self._llm_d_w = [] + self._llm_lm_head_w = None + torch.cuda.empty_cache() + + logger.info("weights: %.2f GB quantized (%s) + %.2f GB lm_head int8 + " + "%.2f GB embed fp16", + self.qw.bytes() / 2 ** 30, self.precision_tier, + self._lm_head_q.numel() / 2 ** 30, + self._llm_embed_w.numel() * 2 / 2 ** 30) + + def _build_rope_tables(self) -> None: + """cos/sin as ``[max_seq, HD]`` fp16, ``cat([f, f], -1)`` tiled. + + The kernel reads ``cos[s*HD + d]`` for d in ``[0, HD)``, so the + half-frequencies must be duplicated across the two halves; a + ``[max_seq, HD/2]`` table indexes out of stride and corrupts RoPE. Being + row-major with stride exactly HD is also what lets a decode step select + position ``pos`` by pointer arithmetic alone. + """ + theta = float(self._config.get("rope_theta", 10000.0)) + inv = 1.0 / (theta ** (torch.arange(0, self.HD, 2, dtype=torch.float32, + device="cuda") / self.HD)) + pos = torch.arange(self.max_seq, device="cuda", dtype=torch.float32) + f = pos[:, None] * inv[None, :] + full = torch.cat([f, f], dim=-1) + self._rope_cos = torch.cos(full).to(_FP16).contiguous() + self._rope_sin = torch.sin(full).to(_FP16).contiguous() + + def _allocate_buffers(self) -> None: + MS, D, Dff, V = self.max_seq, self.D, self.DFF, self.VOCAB + dev = "cuda" + z = lambda *sh, dt: torch.zeros(*sh, dtype=dt, device=dev) # noqa: E731 + + self._x = z(MS, D, dt=_FP16) # residual stream + self._xn = z(MS, D, dt=_FP16) # final-norm output + self._o_proj_out = z(MS, D, dt=_FP16) + self._int8_act_d = z(MS, D, dt=torch.int8) + self._int8_act_ff = z(MS, Dff, dt=torch.int8) + self._bf16_gate_ff = z(MS, Dff, dt=_BF16) + self._bf16_xn_ff = z(MS, Dff, dt=_BF16) + self._int4_act_d = z(MS, D // 2, dt=torch.uint8) if self.use_int4 else None + self._int4_act_ff = (z(MS, Dff // 2, dt=torch.uint8) + if self.use_int4_down else None) + + # Dynamic per-row activation scales — one shared [MS] vector per quant + # site, reused by every layer (decode never uses static calibration: + # that was fitted at prefill M=Se and does not describe one decode row). + self._act_scale = {k: z(MS, dt=torch.float32) + for k in ("qkv", "o", "gu", "down")} + self._lm_act = z(MS, D, dt=torch.int8) + self._lm_act_scale = z(MS, dt=torch.float32) + self._logits = z(1, V, dt=_BF16) + self._logits_all: Optional[torch.Tensor] = None + self._bf16_min = torch.finfo(_BF16).min + + self._probe_bufs = [z(MS, D, dt=_FP16) for _ in self._probe_layers] + self._probe_final = z(MS, D, dt=_FP16) if self._probe_layers else None + + self._tok_dev = z(1, dt=torch.long) + + def _build_attn_backend(self, *, split_kv_bias: int) -> None: + spec = make_chameleon_attention_spec( + num_layers=self.L, num_q_heads=self.H, num_kv_heads=self.HKV, + head_dim=self.HD, max_seq=self.max_seq) + self.attn = ChameleonAttnBackend(spec, max_seq=self.max_seq, + split_kv_bias=split_kv_bias) + + def _load_processor(self) -> None: + """Use the HF ChameleonProcessor verbatim. + + Reimplementing it is a correctness liability: the pipeline is + blend-RGBA-on-white → PIL **LANCZOS** shortest-edge 512 → center-crop + 512 → ``float32(float64(u8) * 0.0078) - 1.0``, giving ``[-1, +0.989]`` + (note: *not* ``[-1, 1]``). ``ChameleonImageProcessorFast`` silently + substitutes BICUBIC for LANCZOS, so the slow/PIL path is required — it + is what ``preprocessor_config.json`` selects by default, and it runs + once per image outside any graph. + """ + from transformers import AutoProcessor + self.processor = AutoProcessor.from_pretrained(str(self.checkpoint_dir)) + if int(getattr(self.processor, "image_seq_length", 0)) != self.IMG_TOKENS_PER_VIEW: + raise RuntimeError( + f"processor image_seq_length=" + f"{self.processor.image_seq_length}, expected " + f"{self.IMG_TOKENS_PER_VIEW}") + + def _load_vqgan(self) -> None: + """Load the HF ``ChameleonVQVAE`` encoder from the checkpoint shards.""" + from safetensors.torch import load_file + from transformers import ChameleonConfig, ChameleonVQVAE + + index = self.checkpoint_dir / "model.safetensors.index.json" + prefix = "model.vqmodel." + sd = {} + if index.exists(): + wmap = json.loads(index.read_text())["weight_map"] + keys = [k for k in wmap if k.startswith(prefix)] + for shard in sorted({wmap[k] for k in keys}): + full = load_file(str(self.checkpoint_dir / shard)) + for k in keys: + if k in full: + sd[k[len(prefix):]] = full[k] + del full + if not sd: + raise FileNotFoundError( + f"no {prefix}* weights found under {self.checkpoint_dir}") + + cfg = ChameleonConfig.from_pretrained(str(self.checkpoint_dir)) + vq = ChameleonVQVAE._from_config(cfg.vq_config) + missing, unexpected = vq.load_state_dict(sd, strict=False) + if unexpected: + raise RuntimeError(f"unexpected vqmodel keys: {unexpected[:4]}") + if missing: + raise RuntimeError(f"missing vqmodel keys: {missing[:4]}") + self.vqgan = vq.eval().to(device="cuda", dtype=_FP16) + # fp32 codebook for the distance/argmin (see _vq_encode). + self._vq_codebook_f32 = ( + self.vqgan.quantize.embedding.weight.detach().float().contiguous()) + logger.info("VQ-GAN encoder loaded (%d tensors, fp16 convs, " + "argmin in %s)", len(sd), + "fp32" if self.vq_argmin_fp32 else "fp16") + + # ================================================================== + # Image tokenization + # ================================================================== + + @torch.no_grad() + def _vq_encode(self, pixel_values: torch.Tensor) -> torch.Tensor: + """``[N, 3, 512, 512]`` fp32 in ``[-1, 0.989]`` → ``[N, 1024]`` token ids. + + The convs run fp16 but the codebook distance and argmin run **fp32**: + ``|z|^2 + |e|^2 - 2 z.e`` is cancellation-prone, and in fp16 it flips + codebook indices (measured 98.14 % → 99.02 % index match vs an all-fp32 + reference for <0.1 ms; see docs §2.4). Tokens are a row-major raster of + the 32x32 latent grid, and the id is simply ``code + 4``. + """ + px = pixel_values.to(device="cuda", dtype=_FP16) + h = self.vqgan.quant_conv(self.vqgan.encoder(px)) + if not self.vq_argmin_fp32: + _, _, idx = self.vqgan.quantize(h) + return idx.view(px.shape[0], -1) + self.IMG_ID_OFFSET + e = self._vq_codebook_f32 # [n_emb, dim] + z = h.permute(0, 2, 3, 1).contiguous().view(-1, e.shape[1]).float() + d = z.pow(2).sum(1, keepdim=True) + e.pow(2).sum(1) - 2.0 * (z @ e.t()) + return d.argmin(1).view(px.shape[0], -1) + self.IMG_ID_OFFSET + + # ================================================================== + # Prompt + # ================================================================== + + def set_prompt(self, text: str, images=None, *, + input_ids: Optional[Sequence[int]] = None) -> None: + """Tokenize, VQ-encode the images, and seed the residual stream. + + Args: + text: prompt containing one ```` per image, e.g. + ``"Describe this image."``. The processor expands each + placeholder to ```` + 1024x ```` + + ````, prepends BOS and appends the sep token. + images: list of PIL images (or a single image). + input_ids: bypass tokenization + VQ entirely and use these ids + verbatim. Used by the precision harness to feed the reference's + exact ids, which isolates LLM error from VQ-GAN index drift. + """ + t0 = time.perf_counter() + if input_ids is not None: + ids = torch.as_tensor(list(input_ids), dtype=torch.long, device="cuda") + n_img = 0 + else: + if images is not None and not isinstance(images, (list, tuple)): + images = [images] + n_text_img = text.count("") + if images and n_text_img != len(images): + raise ValueError( + f"text has {n_text_img} '' placeholders but " + f"{len(images)} images were given") + enc = self.processor(text=text, images=images if images else None, + return_tensors="pt") + ids = enc["input_ids"][0].to("cuda") + n_img = 0 + if images: + img_ids = self._vq_encode(enc["pixel_values"]) # [N, 1024] + n_img = img_ids.shape[0] + # Upstream substitutes at the id level: masked_scatter over + # input_ids == . Order is row-major per image, images in + # order, which matches the placeholder order in the string. + slot = ids == self.IMG_PLACEHOLDER_ID + want = n_img * self.IMG_TOKENS_PER_VIEW + if int(slot.sum()) != want: + raise RuntimeError( + f"{int(slot.sum())} placeholders vs {want} " + "VQ tokens") + ids = ids.masked_scatter(slot, img_ids.reshape(-1).to(ids.dtype)) + + S = int(ids.numel()) + if S > self.max_seq: + raise ValueError( + f"prompt is {S} tokens but max_seq={self.max_seq}. Note " + f"S = 1 + n_img*1026 + n_text + 1.") + if S + 1 > self.max_seq: + raise ValueError( + f"prompt {S} tokens leaves no room to decode within " + f"max_seq={self.max_seq}") + + # NOTE: S is used exactly, never padded. Padding would put junk rows in + # the KV cache that decode would then attend to (a prefill path without a KV cache would pad Se to even). + self.input_ids = ids + self.S = S + self.attn.reset_cache() + # Seed the residual stream. No sqrt(D) scaling — Chameleon feeds raw + # embeddings, and image tokens are ordinary vocab rows (no projector). + torch.index_select(self._llm_embed_w, 0, ids, out=self._x[:S]) + self._prompt_ready = True + self._timing = {"prompt_ms": (time.perf_counter() - t0) * 1e3, + "S": S, "n_images": n_img} + + # ================================================================== + # Forward + # ================================================================== + + def _dims(self) -> dict: + return {"D": self.D, "Dff": self.DFF, "L": self.L, "H": self.H, + "Hd": self.HD, "vocab": self.VOCAB} + + def _bufs(self, *, logits_ptr: int) -> dict: + return { + "x": self._x.data_ptr(), + "xn": self._xn.data_ptr(), + "o_proj_out": self._o_proj_out.data_ptr(), + "int8_act_d": self._int8_act_d.data_ptr(), + "int8_act_ff": self._int8_act_ff.data_ptr(), + "int4_act_d": self._int4_act_d.data_ptr() if self.use_int4 else 0, + "int4_act_ff": (self._int4_act_ff.data_ptr() + if self.use_int4_down else 0), + "bf16_gate_ff": self._bf16_gate_ff.data_ptr(), + "bf16_xn_ff": self._bf16_xn_ff.data_ptr(), + "logits": logits_ptr, + "lm_act": self._lm_act.data_ptr(), + "lm_act_scale": self._lm_act_scale.data_ptr(), + } + + def _weights(self) -> dict: + w = { + "rope_cos": self._rope_cos.data_ptr(), + "rope_sin": self._rope_sin.data_ptr(), + "final_norm_w": self._llm_norm_w.data_ptr(), + "lm_head_w": self._lm_head_q.data_ptr(), + "lm_head_w_scale": self._lm_head_s.data_ptr(), + "input_ln_w": [t.data_ptr() for t in self._llm_input_ln_w], + "post_ln_w": [t.data_ptr() for t in self._llm_post_ln_w], + "q_norm_w": [t.data_ptr() for t in self._llm_q_norm_w], + "q_norm_b": [t.data_ptr() for t in self._llm_q_norm_b], + "k_norm_w": [t.data_ptr() for t in self._llm_k_norm_w], + "k_norm_b": [t.data_ptr() for t in self._llm_k_norm_b], + } + for key in cq.PROJECTIONS: + w[f"{key}_w"] = self.qw.ptr[key] + w[f"{key}_w_scale"] = self.qw.scale_ptr[key] + return w + + def _scales_dev(self) -> dict: + return {f"act_{k}": [v.data_ptr()] * self.L + for k, v in self._act_scale.items()} + + def _probe(self) -> Optional[dict]: + if not self._probe_layers: + return None + return {"layers": self._probe_layers, + "bufs": [b.data_ptr() for b in self._probe_bufs], + "final_buf": self._probe_final.data_ptr()} + + def _require_prompt(self) -> None: + if not self._prompt_ready: + raise RuntimeError("call set_prompt() before prefill()/generate()") + + @torch.no_grad() + def prefill(self, *, logits_all: bool = False) -> torch.Tensor: + """Run the prompt through the 32 layers, filling the KV cache. + + Returns the masked BF16 logits: ``[1, vocab]`` for the last position, or + ``[S, vocab]`` when ``logits_all`` (teacher-forced comparison). + """ + self._require_prompt() + t0 = time.perf_counter() + if logits_all: + if self._logits_all is None or self._logits_all.shape[0] < self.S: + self._logits_all = torch.zeros(self.max_seq, self.VOCAB, + dtype=_BF16, device="cuda") + out = self._logits_all[:self.S] + logits_ptr = self._logits_all.data_ptr() + else: + out = self._logits + logits_ptr = self._logits.data_ptr() + + chameleon_forward( + fvk, self._bufs(logits_ptr=logits_ptr), self._weights(), + self._dims(), self._scales_dev(), + attn=self.attn, S=self.S, pos=None, stream=0, + use_int4=self.use_int4, use_int4_down=self.use_int4_down, + use_hadamard=self.use_hadamard, + ffn_down_clamp_value=self.ffn_down_clamp, + ffn_down_clamp_last_n=self.ffn_down_clamp_last_n, + logits_all=logits_all, probe=self._probe()) + self._mask_image_logits(out) + torch.cuda.synchronize() + self._timing["prefill_ms"] = (time.perf_counter() - t0) * 1e3 + return out + + @torch.no_grad() + def decode_step(self, token_id, *, pos: int, stream: int = 0) -> torch.Tensor: + """One decode step: embed ``token_id``, attend keys ``[0, pos]``. + + ``pos`` is the absolute KV position this token occupies, i.e. ``S`` for + the first generated token. + + ``stream`` must be the capture stream when this body is being recorded + into a CUDA graph. Launching the kernels on the legacy default stream + while another stream is capturing leaves them **out** of the graph + without raising — the replay then only re-runs the torch ops and the + logits never change, which looks exactly like a frozen/stale graph. + """ + self._require_prompt() + if isinstance(token_id, torch.Tensor): + self._tok_dev.copy_(token_id.reshape(1).to(torch.long)) + else: + self._tok_dev.fill_(int(token_id)) + torch.index_select(self._llm_embed_w, 0, self._tok_dev, out=self._x[:1]) + chameleon_forward( + fvk, self._bufs(logits_ptr=self._logits.data_ptr()), self._weights(), + self._dims(), self._scales_dev(), + attn=self.attn, S=1, pos=int(pos), stream=int(stream), + use_int4=self.use_int4, use_int4_down=self.use_int4_down, + use_hadamard=self.use_hadamard, + ffn_down_clamp_value=self.ffn_down_clamp, + ffn_down_clamp_last_n=self.ffn_down_clamp_last_n) + self._mask_image_logits(self._logits) + return self._logits + + def _mask_image_logits(self, logits: torch.Tensor) -> None: + """Suppress the 8192 image-codebook ids, as upstream does every forward. + + Upstream applies this inside ``forward`` at *all* positions, so it is + not something a ``LogitsProcessor`` could be used for and it cannot be + disabled through the generation config. + """ + lo = self.IMG_ID_OFFSET + logits[:, lo:lo + self.N_IMG_CODES].fill_(self._bf16_min) + + # ================================================================== + # Generate + # ================================================================== + + @torch.no_grad() + def generate(self, text: Optional[str] = None, images=None, *, + max_new_tokens: int = 32, eos_token_id: Optional[int] = None, + return_ids: bool = False, skip_special_tokens: bool = True): + """Greedy decode. Returns the decoded string (or the raw id list). + + Greedy only: the checkpoint's generation config is already + ``do_sample=False``, and argmax over BF16 logits is order-preserving, so + it matches an fp32 argmax except on exact BF16 ties. + """ + if text is not None: + self.set_prompt(text, images) + self._require_prompt() + eos = self.EOS_ID if eos_token_id is None else int(eos_token_id) + budget = min(max_new_tokens, self.max_seq - self.S) + if budget < max_new_tokens: + logger.warning("max_new_tokens clipped %d -> %d by max_seq=%d", + max_new_tokens, budget, self.max_seq) + + logits = self.prefill() + tok = int(torch.argmax(logits[0]).item()) + out: List[int] = [tok] + + t0 = time.perf_counter() + steps = 0 + for i in range(budget - 1): + if tok == eos: + break + logits = self.decode_step(tok, pos=self.S + i) + tok = int(torch.argmax(logits[0]).item()) + out.append(tok) + steps += 1 + torch.cuda.synchronize() + dt = time.perf_counter() - t0 + self._timing["decode_steps"] = steps + self._timing["decode_ms_per_token"] = (dt / steps * 1e3) if steps else 0.0 + self._timing["decode_tok_s"] = (steps / dt) if dt > 0 else 0.0 + + if out and out[-1] == eos: + out = out[:-1] + if return_ids: + return out + return self.processor.tokenizer.decode( + out, skip_special_tokens=skip_special_tokens) + + # ================================================================== + # Introspection + # ================================================================== + + def snapshot_probe(self) -> dict: + """Per-layer post-residual hidden states captured during the last forward.""" + if not self._probe_layers: + return {} + S = self.S + out = {f"layer_{li}": b[:S].clone() + for li, b in zip(self._probe_layers, self._probe_bufs)} + out["final_norm"] = self._probe_final[:S].clone() + return out + + def reset(self) -> None: + self.attn.reset_cache() + self._prompt_ready = False + self.input_ids = None + self.S = 0 + + @property + def timing(self) -> dict: + return dict(self._timing) + + def precision_spec(self) -> dict: + return { + "tier": self.precision_tier, + "llm_gemms": dict(self.qw.precision), + "lm_head": "int8", + "residual": "fp16", + "attention": "fp16 FA2 causal (split_kv_bias=" + f"{self.attn.split_kv_bias})", + "kv_cache": "fp16", + "ffn_down_clamp": f"{self.ffn_down_clamp:.0f} on last {self.ffn_down_clamp_last_n} layers", + "vqgan": f"fp16 convs, argmin " + f"{'fp32' if self.vq_argmin_fp32 else 'fp16'}", + } + + def get_model_info(self) -> dict: + return { + "model": "chameleon-7b", "arch": "rtx_sm87", + "layers": self.L, "hidden": self.D, "ffn": self.DFF, + "heads": f"{self.H}Q/{self.HKV}KV", "head_dim": self.HD, + "vocab": self.VOCAB, "max_seq": self.max_seq, + "precision": self.precision_spec(), + } + + +__all__ = ["ChameleonTorchFrontendRtxSm87"] diff --git a/flash_rt/hardware/rtx/attn_backend_chameleon.py b/flash_rt/hardware/rtx/attn_backend_chameleon.py new file mode 100644 index 00000000..8aafba63 --- /dev/null +++ b/flash_rt/hardware/rtx/attn_backend_chameleon.py @@ -0,0 +1,237 @@ +"""FlashRT — Chameleon-7B VLM attention backend for Jetson Orin (SM87). + +This backend owns a **real per-layer KV cache** so the LLM can decode +autoregressively (a prefill-only backend could instead share one K/V scratch +across all 32 layers, ``layer_stride = 0``). + +Two load-bearing design points, both measured — see +``docs/chameleon7b_rtx_sm87.md`` §2.1 and §3.1: + +1. **The K/V GEMMs write straight into the cache.** CUTLASS hard-wires the + output row stride to ``N`` (``cutlass_sm80_int8_rowwise_fp16out.cu:169-171``), + and a per-layer slab of ``[max_seq, num_kv_heads*head_dim]`` has exactly that + row stride, so no staging buffer or copy is needed for either prefill + (``M = S`` at the slab base) or decode (``M = 1`` at ``+ pos*row_stride``). + ``qk_norm_rope_fused_fp16`` then does QK-LayerNorm + RoPE in place. + +2. **Split-KV must be forced on with a biased ``num_sms``.** FA2's heuristic + (``fa2_wrapper_causal.cu:41-43,152-158``) returns ``num_splits = 1`` whenever + ``batch*num_q_heads*num_m_blocks >= 0.8 * (num_sms*2)``. Chameleon decode is + ``1*32*1 = 32`` against ``0.8*32 = 25.6``, so at the true SM count split-KV + silently does nothing (measured: bit-identical output, 1.05x). ``num_sms`` is + a pure heuristic knob in this wrapper, so ``split_kv_bias`` multiplies it; + bias 4 measured 204.9 -> 141.8 us (**1.44x**) at fp16-rounding-level delta. + This is why the Qwen3-VL split-KV lever does not transfer as-is — + that model has 16 Q heads and lands under the threshold naturally. + +FA2 ``fwd_fp16_causal`` is **mandatory** for decode: its causal mask is +bottom-right aligned (``fa2_wrapper_causal.cu:126-138``) so ``q=1, kv=N`` +attends all N keys. The cuBLAS ``attention_mha_causal_fp16`` fallback is **top-left** aligned (``softmax.cu:182-191`` masks with +``q = row % S_q``), so at ``S_q = 1`` only column 0 survives — it is silently +wrong rather than merely slow. This backend therefore raises instead of +degrading to it. +""" + +from __future__ import annotations + +import logging + +import torch + +from flash_rt.hardware.backend import AttentionBackendBase, AttentionSpec + +logger = logging.getLogger(__name__) + +SITE = "llm" + +#: FA2's own cap on the split count (``fa2_wrapper_causal.cu`` passes 128). +_MAX_SPLITS = 128 + + +class ChameleonAttnBackend(AttentionBackendBase): + """Chameleon-7B self-attention with a per-layer FP16 KV cache. + + Owns every buffer it needs (KV cache, Q/O staging, softmax LSE, split-KV + accumulators). Pointers are stable for the object's lifetime, so they are + safe to bake into a captured CUDA graph — but the backend must be kept + alive for as long as any such graph exists. + """ + + def __init__(self, spec: AttentionSpec, *, max_seq: int, + split_kv_bias: int = 4) -> None: + super().__init__(spec) + if set(spec.sites.keys()) != {SITE}: + raise ValueError( + f"ChameleonAttnBackend expects exactly the {SITE!r} site, " + f"got {set(spec.sites.keys())}") + + s = spec.site(SITE) + self.num_layers = int(s.num_layers) + self.num_q_heads_ = int(s.num_q_heads) + self.num_kv_heads_ = int(s.num_kv_heads) + self.head_dim_ = int(s.head_dim) + self.max_seq = int(max_seq) + self.split_kv_bias = int(split_kv_bias) + self.scale = float(self.head_dim_) ** -0.5 + + if self.head_dim_ != 128: + raise ValueError( + f"head_dim must be 128 (FA2 fp16 causal aborts otherwise, " + f"fa2_wrapper_causal.cu:235-243); got {self.head_dim_}") + + try: + import flash_rt.flash_rt_fa2 as _fa2 + except ImportError as e: # pragma: no cover + raise RuntimeError( + "Chameleon decode requires flash_rt_fa2 (the cuBLAS MHA " + "fallback is top-left-aligned causal and therefore WRONG at " + "q_seq=1). Rebuild with -DFLASHRT_ENABLE_FA2=ON.") from e + if not hasattr(_fa2, "fwd_fp16_causal"): + raise RuntimeError( + "flash_rt_fa2 lacks fwd_fp16_causal; Chameleon decode cannot " + "fall back to cuBLAS MHA (top-left-aligned mask is wrong at " + "q_seq=1). Rebuild with -DFA2_DTYPES='fp16;bf16' " + "-DFA2_HDIMS='128;256'.") + self._fa2_fwd = _fa2.fwd_fp16_causal + + dev, fp16, fp32 = "cuda", torch.float16, torch.float32 + kv_row = self.num_kv_heads_ * self.head_dim_ # 4096 == GEMM N + q_row = self.num_q_heads_ * self.head_dim_ + + # Per-layer KV cache. The [max_seq, kv_row] slab per layer is exactly a + # legal CUTLASS destination (row stride == N), which is what lets the + # K/V GEMMs write into it directly. + self.K_cache = torch.zeros(self.num_layers, self.max_seq, kv_row, + dtype=fp16, device=dev) + self.V_cache = torch.zeros(self.num_layers, self.max_seq, kv_row, + dtype=fp16, device=dev) + # Q staging (prefill writes S rows, decode row 0). O aliases Q, matching + # the caller convention: the pipeline reads its result back in place. + self.Q_buf = torch.zeros(self.max_seq, q_row, dtype=fp16, device=dev) + + lse_rows = ((self.max_seq + 127) // 128) * 128 + self.lse_buf = torch.zeros(1, self.num_q_heads_, lse_rows, + dtype=fp32, device=dev) + # Split-KV accumulators, sized for the decode case (seqlen_q == 1). + # Empty splits are self-initialised by the kernel, so no pre-fill. + self.lse_accum = torch.zeros(_MAX_SPLITS, 1, self.num_q_heads_, 1, + dtype=fp32, device=dev) + self.o_accum = torch.zeros(_MAX_SPLITS, 1, self.num_q_heads_, 1, + self.head_dim_, dtype=fp32, device=dev) + + self._num_sms = torch.cuda.get_device_properties( + torch.cuda.current_device()).multi_processor_count + + kv_gb = 2 * self.K_cache.numel() * 2 / 2 ** 30 + logger.info( + "ChameleonAttnBackend: L=%d %dQ/%dKV hd=%d max_seq=%d " + "KV=%.2f GB split_kv_bias=%d (num_sms %d->%d)", + self.num_layers, self.num_q_heads_, self.num_kv_heads_, + self.head_dim_, self.max_seq, kv_gb, self.split_kv_bias, + self._num_sms, self._num_sms * self.split_kv_bias) + + # ------------------------------------------------------------------ + # Layout / pointers + # ------------------------------------------------------------------ + + @property + def kv_layer_stride_bytes(self) -> int: + return self.max_seq * self.num_kv_heads_ * self.head_dim_ * 2 + + @property + def kv_row_stride_bytes(self) -> int: + return self.num_kv_heads_ * self.head_dim_ * 2 + + def _check_layer(self, layer_idx: int) -> None: + if not 0 <= layer_idx < self.num_layers: + raise IndexError( + f"layer_idx {layer_idx} out of range [0, {self.num_layers})") + + def get_slot_ptrs(self, site: str, layer_idx: int) -> dict: + """Prefill slots: Q/O staging plus the base of this layer's KV slab.""" + if site != SITE: + raise KeyError(f"unknown site {site!r}") + self._check_layer(layer_idx) + off = layer_idx * self.kv_layer_stride_bytes + q = self.Q_buf.data_ptr() + return {"Q": q, "O": q, + "K": self.K_cache.data_ptr() + off, + "V": self.V_cache.data_ptr() + off} + + def kv_row_ptrs(self, layer_idx: int, pos: int) -> tuple: + """Decode slots: the single KV row for absolute position ``pos``.""" + self._check_layer(layer_idx) + if not 0 <= pos < self.max_seq: + raise IndexError( + f"pos {pos} out of range [0, max_seq={self.max_seq}); raise " + f"max_seq at construction time") + off = layer_idx * self.kv_layer_stride_bytes + pos * self.kv_row_stride_bytes + return (self.K_cache.data_ptr() + off, self.V_cache.data_ptr() + off) + + # ------------------------------------------------------------------ + # Attention + # ------------------------------------------------------------------ + + def _fa2(self, layer_idx: int, q_seq: int, kv_seq: int, *, + use_split_kv: bool, stream: int) -> int: + qr = self.num_q_heads_ * self.head_dim_ + kr = self.num_kv_heads_ * self.head_dim_ + off = layer_idx * self.kv_layer_stride_bytes + q = self.Q_buf.data_ptr() + # num_sms == 0 disables the split-KV heuristic entirely; a biased value + # is the only way to reach num_splits > 1 at 32 Q heads (see module doc). + num_sms = self._num_sms * self.split_kv_bias if use_split_kv else 0 + self._fa2_fwd( + q, self.K_cache.data_ptr() + off, self.V_cache.data_ptr() + off, q, + self.lse_buf.data_ptr(), + self.lse_accum.data_ptr() if use_split_kv else 0, + self.o_accum.data_ptr() if use_split_kv else 0, + batch=1, seqlen_q=q_seq, seqlen_k=kv_seq, + num_heads_q=self.num_q_heads_, num_heads_kv=self.num_kv_heads_, + head_dim=self.head_dim_, + q_strides=(q_seq * qr, qr, self.head_dim_), + k_strides=(kv_seq * kr, kr, self.head_dim_), + v_strides=(kv_seq * kr, kr, self.head_dim_), + o_strides=(q_seq * qr, qr, self.head_dim_), + softmax_scale=self.scale, num_sms=num_sms, stream=stream) + return q + + def run_prefill(self, layer_idx: int, seq_len: int, *, stream: int = 0) -> int: + """Square causal attention over ``seq_len`` tokens. Result lands in Q_buf.""" + self._check_layer(layer_idx) + if seq_len > self.max_seq: + raise ValueError(f"seq_len {seq_len} > max_seq {self.max_seq}") + # Prefill already fills every SM (num_m_blocks = ceil(S/64)), so + # splitting KV would only add a combine pass. + return self._fa2(layer_idx, seq_len, seq_len, + use_split_kv=False, stream=stream) + + def run_decode(self, layer_idx: int, pos: int, *, stream: int = 0) -> int: + """One query row attending keys ``[0, pos]``. Result lands in Q_buf row 0.""" + self._check_layer(layer_idx) + return self._fa2(layer_idx, 1, pos + 1, + use_split_kv=self.split_kv_bias > 1, stream=stream) + + # No generic ``run(site, layer_idx, q_seq, ...)`` on purpose. Prefill and + # decode differ in more than q_seq here (split-KV on/off, and the caller + # must supply the absolute KV position rather than a length), so a shim that + # inferred the mode from ``q_seq == 1`` would be an untested footgun. The + # inherited ``AttentionBackendBase.run`` raises NotImplementedError, which + # is the behaviour we want if a generic caller ever appears. + + def reset_cache(self) -> None: + self.K_cache.zero_() + self.V_cache.zero_() + + +def make_chameleon_attention_spec(*, num_layers: int, num_q_heads: int, + num_kv_heads: int, head_dim: int, + max_seq: int) -> AttentionSpec: + spec = AttentionSpec() + spec.add_site(SITE, num_layers=num_layers, num_q_heads=num_q_heads, + num_kv_heads=num_kv_heads, head_dim=head_dim, + max_q_seq=max_seq, max_kv_seq=max_seq, causal=True) + return spec + + +__all__ = ["ChameleonAttnBackend", "make_chameleon_attention_spec", "SITE"] diff --git a/flash_rt/models/chameleon/pipeline_rtx.py b/flash_rt/models/chameleon/pipeline_rtx.py new file mode 100644 index 00000000..6e11cf39 --- /dev/null +++ b/flash_rt/models/chameleon/pipeline_rtx.py @@ -0,0 +1,305 @@ +"""FlashRT — Chameleon-7B VLM forward for Jetson AGX Orin (SM87). + +Standalone text-generating forward on the Chameleon-7B INT8/QuaRot-INT4 +kernel set (SM80 CUTLASS rowwise GEMMs + Hadamard rotations). Key design +points: + +1. **No attention bias.** Upstream Chameleon has ``attention_bias = false`` and + ``mlp_bias = false``, so the no-bias GEMM entries are used unconditionally. +2. **A real KV cache.** The K and V GEMMs write straight into the attention + backend's per-layer slab — legal because CUTLASS hard-wires its output row + stride to ``N``, which equals the cache's row stride. No staging, no copy. +3. **One code path for prefill and decode.** ``pos is None`` means prefill + (``S`` rows at the slab base, RoPE from position 0); ``pos`` set means a + single decode row written at ``pos``, with the cos/sin pointers advanced by + ``pos`` rows. ``qk_norm_rope_fused_fp16`` needs no modification for this: + it derives position as ``row / num_heads``, which is 0 at ``S = 1``, so the + position lives entirely in the table pointer. +4. **lm_head tail instead of an action head.** Final RMSNorm, then per-row INT8 + quantization of the wanted row(s), then an INT8 GEMM to BF16 logits. The + ``mask_image_logits`` step and the argmax live in the frontend (they are + torch ops on the logits view, which are graph-safe — verified). + +lm_head stays INT8 in both precision tiers: it is 268 MB/token = 1.74 ms = +3.7 % of the INT8 decode budget, and dropping to 15 INT4 levels over a +65536-row output is not worth 0.8 ms. + +Raw-pointer interface only (int pointers + Python primitives) for CUDA-Graph +safety: no torch ops, no allocation, no sync inside the forward. +""" + +from __future__ import annotations + +import ctypes + +_CUDART = None + + +def _gpu_copy(dst_ptr: int, src_ptr: int, nbytes: int, stream: int) -> None: + """Async D2D copy — the in-graph layer-probe mechanism.""" + global _CUDART + if _CUDART is None: + _CUDART = ctypes.CDLL("libcudart.so") + _CUDART.cudaMemcpyAsync( + ctypes.c_void_p(dst_ptr), ctypes.c_void_p(src_ptr), + ctypes.c_size_t(nbytes), 3, ctypes.c_void_p(stream)) + + +def _check(status, name: str, shape) -> None: + if status != 0: + raise RuntimeError(f"{name} failed: status={status} shape={shape}") + + +def chameleon_forward( + fvk, bufs, weights, dims, scales_dev, + *, attn, S: int, pos=None, stream: int = 0, + use_int4: bool = False, use_int4_down: bool = False, + use_hadamard: bool = False, + ffn_down_clamp_value: float = 60000.0, + ffn_down_clamp_last_n: int = 4, + logits_all: bool = False, probe=None, +) -> None: + """Run the 32-layer Chameleon-7B decoder and the lm_head. + + Args: + fvk: ``flash_rt.flash_rt_kernels`` module. + bufs / weights / dims / scales_dev: int-pointer dicts from the frontend. + attn: ``ChameleonAttnBackend`` (owns the KV cache). + S: rows to process. ``1`` on the decode path. + pos: ``None`` for prefill; otherwise the absolute KV position of the + single decode row. RoPE and the KV write both key off this. + ffn_down_clamp_value: symmetric clamp applied to the down-projection + output before the residual add; ``<= 0`` disables it. + **This is required for correctness, not a tuning knob.** Measured on + this checkpoint (ISL=1032), the reference's magnitudes are tiny + through L30 and then explode at L31 only: + + layer L28 L29 L30 L31 + residual 1616 1720 2032 266240 + down_out 1120 1032 1880 264192 + + 2.6e5 is far beyond FP16's 65504, so the store becomes ``inf`` and + the final RMSNorm then poisons that row's logits. Because the + pre-L31 residual is only ~2032, clamping the down output at 60000 + keeps the residual at ~62000 < 65504 — so no BF16 residual stream is + needed. See ``docs/chameleon7b_rtx_sm87.md``. Unlike the + Thor path we do **not** clamp the down *input*: ours is BF16 + (``cutlass_int8_silu_gated_bf16out``), whose range absorbs the + 151552 without issue. + logits_all: compute logits for **all** S rows (teacher-forced + precision comparison). Requires an ``[S, vocab]`` logits buffer. + probe: ``{"layers": [...], "bufs": [...], "final_buf": ptr}`` to snapshot + the post-residual hidden state; ``None`` disables it at zero cost. + """ + decode = pos is not None + if decode and S != 1: + raise ValueError(f"decode path requires S == 1, got {S}") + + D = int(dims["D"]) + Dff = int(dims["Dff"]) + L = int(dims["L"]) + H = int(dims["H"]) + Hd = int(dims["Hd"]) + V = int(dims["vocab"]) + + x_ptr = int(bufs["x"]) + xn_ptr = int(bufs["xn"]) + int8_act_d_ptr = int(bufs["int8_act_d"]) + int8_act_ff_ptr = int(bufs["int8_act_ff"]) + int4_act_d_ptr = int(bufs.get("int4_act_d", 0)) + int4_act_ff_ptr = int(bufs.get("int4_act_ff", 0)) + bf16_gate_ptr = int(bufs["bf16_gate_ff"]) + bf16_xn_ff_ptr = int(bufs["bf16_xn_ff"]) + o_proj_out_ptr = int(bufs["o_proj_out"]) + logits_ptr = int(bufs["logits"]) + lm_act_ptr = int(bufs["lm_act"]) + lm_scale_ptr = int(bufs["lm_act_scale"]) + + # RoPE position enters purely through the table pointer (see module doc). + rope_off = (pos if decode else 0) * Hd * 2 + cos_ptr = int(weights["rope_cos"]) + rope_off + sin_ptr = int(weights["rope_sin"]) + rope_off + + probe_map = None + probe_final = None + if probe is not None: + layers = probe.get("layers") or [] + pbufs = probe.get("bufs") or [] + if len(layers) != len(pbufs): + raise ValueError("probe['layers'] and probe['bufs'] length mismatch") + probe_map = {int(li): int(p) for li, p in zip(layers, pbufs)} + probe_final = int(probe.get("final_buf") or 0) + + # ── entry: fused RMSNorm + quantize (layer 0's input_layernorm) ── + if use_int4: + fvk.rms_norm_fht_int4_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int4_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.rms_norm_fht_int8_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int8_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + else: + fvk.rms_norm_int8_rowwise_fp16( + x_ptr, int(weights["input_ln_w"][0]), + int8_act_d_ptr, int(scales_dev["act_qkv"][0]), + S, D, 1e-5, int(stream)) + + act_d_ptr = int4_act_d_ptr if use_int4 else int8_act_d_ptr + + for li in range(L): + if decode: + K_ptr, V_ptr = attn.kv_row_ptrs(li, pos) + Q_ptr = int(attn.get_slot_ptrs("llm", li)["Q"]) + else: + slots = attn.get_slot_ptrs("llm", li) + Q_ptr, K_ptr, V_ptr = int(slots["Q"]), int(slots["K"]), int(slots["V"]) + + a_qkv = int(scales_dev["act_qkv"][li]) + a_o = int(scales_dev["act_o"][li]) + a_gu = int(scales_dev["act_gu"][li]) + a_d = int(scales_dev["act_down"][li]) + + # ── Q/K/V: K and V land directly in the KV cache ── + qkv_gemm = (fvk.cutlass_int4_rowwise_fp16out if use_int4 + else fvk.cutlass_int8_rowwise_fp16out) + for name, out_ptr in (("q_w", Q_ptr), ("k_w", K_ptr), ("v_w", V_ptr)): + _check(qkv_gemm(act_d_ptr, int(weights[name][li]), a_qkv, + int(weights[name + "_scale"][li]), out_ptr, + S, D, D, int(stream)), + f"{'int4' if use_int4 else 'int8'} {name}", (S, D, D)) + + # ── fused per-head QK LayerNorm(+bias) + rotate-half RoPE, in place ── + fvk.qk_norm_rope_fused_fp16( + Q_ptr, K_ptr, + int(weights["q_norm_w"][li]), int(weights["q_norm_b"][li]), + int(weights["k_norm_w"][li]), int(weights["k_norm_b"][li]), + cos_ptr, sin_ptr, + S, H, Hd, 1e-5, int(stream)) + + # ── causal MHA (result written back into the Q slot) ── + if decode: + attn.run_decode(li, pos, stream=int(stream)) + else: + attn.run_prefill(li, S, stream=int(stream)) + + # ── O projection ── + if use_int4: + fvk.fht_int4_quant_fp16(Q_ptr, int4_act_d_ptr, a_o, S, D, int(stream)) + elif use_hadamard: + fvk.fht_int8_quant_fp16(Q_ptr, int8_act_d_ptr, a_o, S, D, int(stream)) + else: + fvk.quantize_int8_rowwise_fp16(Q_ptr, int8_act_d_ptr, a_o, + S, D, int(stream)) + o_gemm = (fvk.cutlass_int4_rowwise_fp16out if use_int4 + else fvk.cutlass_int8_rowwise_fp16out) + _check(o_gemm(act_d_ptr, int(weights["o_w"][li]), a_o, + int(weights["o_w_scale"][li]), o_proj_out_ptr, + S, D, D, int(stream)), "o_proj", (S, D, D)) + + # ── residual_1 + post-attention RMSNorm + quantize ── + if use_int4: + fvk.residual_add_rms_norm_fht_int4_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int4_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.residual_add_rms_norm_fht_int8_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int8_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + else: + fvk.residual_add_rms_norm_int8_rowwise_fp16( + x_ptr, o_proj_out_ptr, int(weights["post_ln_w"][li]), + int8_act_d_ptr, a_gu, S, D, 1e-5, int(stream)) + + # ── FFN: gate -> BF16, up with fused SiLU(gate)* -> BF16 ── + if use_int4: + _check(fvk.cutlass_int4_rowwise_bf16out( + int4_act_d_ptr, int(weights["gate_w"][li]), a_gu, + int(weights["gate_w_scale"][li]), bf16_gate_ptr, + S, Dff, D, int(stream)), "int4 gate", (S, Dff, D)) + _check(fvk.cutlass_int4_silu_gated_bf16out( + int4_act_d_ptr, int(weights["up_w"][li]), a_gu, + int(weights["up_w_scale"][li]), bf16_gate_ptr, bf16_xn_ff_ptr, + S, Dff, D, int(stream)), "int4 up+silu", (S, Dff, D)) + else: + _check(fvk.cutlass_int8_rowwise_bf16out( + int8_act_d_ptr, int(weights["gate_w"][li]), a_gu, + int(weights["gate_w_scale"][li]), bf16_gate_ptr, + S, Dff, D, int(stream)), "int8 gate", (S, Dff, D)) + _check(fvk.cutlass_int8_silu_gated_bf16out( + int8_act_d_ptr, int(weights["up_w"][li]), a_gu, + int(weights["up_w_scale"][li]), bf16_gate_ptr, bf16_xn_ff_ptr, + S, Dff, D, int(stream)), "int8 up+silu", (S, Dff, D)) + + # ── down projection ── + if use_int4_down: + fvk.fht128_int4_quant_bf16(bf16_xn_ff_ptr, int4_act_ff_ptr, a_d, + S, Dff, int(stream)) + _check(fvk.cutlass_int4_rowwise_fp16out( + int4_act_ff_ptr, int(weights["d_w"][li]), a_d, + int(weights["d_w_scale"][li]), o_proj_out_ptr, + S, D, Dff, int(stream)), "int4 down", (S, D, Dff)) + else: + fvk.quantize_int8_rowwise(bf16_xn_ff_ptr, int8_act_ff_ptr, a_d, + S, Dff, int(stream)) + _check(fvk.cutlass_int8_rowwise_fp16out( + int8_act_ff_ptr, int(weights["d_w"][li]), a_d, + int(weights["d_w_scale"][li]), o_proj_out_ptr, + S, D, Dff, int(stream)), "int8 down", (S, D, Dff)) + + # Guard the FP16 residual against L31's massive down output (see the + # measured table in the docstring). Restricted to the last + # ``ffn_down_clamp_last_n`` layers: the magnitude grows monotonically + # with depth and L28 measures 1616, i.e. 37x below the clamp, so the + # earlier layers cannot reach it. Clamping all 32 layers instead costs + # 6.2 ms of a 281 ms prefill (2.2 %) for no effect. The Gate-1 harness + # reports per-layer clamp saturation, so a checkpoint that violates the + # assumption is detectable — raise this to L if that ever happens. + if ffn_down_clamp_value > 0.0 and li >= L - ffn_down_clamp_last_n: + fvk.clamp_inplace_fp16(o_proj_out_ptr, float(ffn_down_clamp_value), + S * D, int(stream)) + + # ── residual_2 (+ next layer's input_layernorm + quantize) ── + if li < L - 1: + if use_int4: + fvk.residual_add_rms_norm_fht_int4_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int4_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + elif use_hadamard: + fvk.residual_add_rms_norm_fht_int8_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int8_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + else: + fvk.residual_add_rms_norm_int8_rowwise_fp16( + x_ptr, o_proj_out_ptr, int(weights["input_ln_w"][li + 1]), + int8_act_d_ptr, int(scales_dev["act_qkv"][li + 1]), + S, D, 1e-5, int(stream)) + else: + fvk.residual_add_fp16(x_ptr, o_proj_out_ptr, S * D, int(stream)) + + if probe_map is not None and li in probe_map: + _gpu_copy(probe_map[li], x_ptr, S * D * 2, stream) + + # ── final RMSNorm ── + fvk.rms_norm_fp16(x_ptr, int(weights["final_norm_w"]), xn_ptr, + S, D, 1e-5, int(stream)) + if probe_final: + _gpu_copy(probe_final, xn_ptr, S * D * 2, stream) + + # ── lm_head: INT8 W8A8 -> BF16 logits ── + # Next-token prediction reads the last row; logits_all is for the + # teacher-forced precision comparison. + rows, row_off = (S, 0) if logits_all else (1, S - 1) + fvk.quantize_int8_rowwise_fp16(xn_ptr + row_off * D * 2, lm_act_ptr, + lm_scale_ptr, rows, D, int(stream)) + _check(fvk.cutlass_int8_rowwise_bf16out( + lm_act_ptr, int(weights["lm_head_w"]), lm_scale_ptr, + int(weights["lm_head_w_scale"]), logits_ptr, + rows, V, D, int(stream)), "lm_head", (rows, V, D)) + + +__all__ = ["chameleon_forward"] From 75a8ef88bd75cdea9e720dc4695f9098c86a4d99 Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 04:34:51 +0000 Subject: [PATCH 3/7] feat(chameleon): registry wiring, tests, examples and docs - Register ("chameleon", "torch", "thor") and ("chameleon", "torch", "rtx_sm87") in _PIPELINE_MAP and allow the SM87 key in _SM87_ALLOWED. - api.load_model redirect for config="chameleon" (chat-style VLM, same pattern as qwen3_vl): raises NotImplementedError pointing at the two direct-instantiation frontends. - tests/test_chameleon_thor_vqgan_backend.py: eager-vs-TRT VQGAN backend contract test. - scripts/: bench_chameleon_thor.py, check_chameleon_thor_precision.py, profile_chameleon_thor.py, chameleon_orin_check.py (Gate-1 harness) and build_vqgan_trt.py (now driven by the vendored flash_rt.models.chameleon.vqgan package); HF BF16 reference rows use transformers' ChameleonForConditionalGeneration directly. - examples/thor/chameleon_quickstart.py + README entry, benchmarks/chameleon_thor_latency.py. - Docs: chameleon_usage.md, chameleon_thor_sm110.md and chameleon7b_rtx_sm87.md; Chameleon rows in USAGE.md, README.md and docs/benchmark_comparison.md. Thor numbers were measured on Jetson AGX Thor (sm_110). All SM87 runtime numbers in the Orin doc (21.07 tok/s, 16/16 bit-identical greedy vs HF BF16) come from Orin hardware in the derivative repo and still need SM87 validation in this tree. --- README.md | 1 + USAGE.md | 33 + benchmarks/chameleon_thor_latency.py | 180 +++++ docs/benchmark_comparison.md | 21 + docs/chameleon7b_rtx_sm87.md | 746 +++++++++++++++++++++ docs/chameleon_thor_sm110.md | 344 ++++++++++ docs/chameleon_usage.md | 143 ++++ examples/thor/README.md | 19 + examples/thor/chameleon_quickstart.py | 95 +++ flash_rt/api.py | 18 + flash_rt/hardware/__init__.py | 19 + scripts/bench_chameleon_thor.py | 358 ++++++++++ scripts/build_vqgan_trt.py | 324 +++++++++ scripts/chameleon_orin_check.py | 392 +++++++++++ scripts/check_chameleon_thor_precision.py | 203 ++++++ scripts/profile_chameleon_thor.py | 115 ++++ tests/test_chameleon_thor_vqgan_backend.py | 18 + 17 files changed, 3029 insertions(+) create mode 100644 benchmarks/chameleon_thor_latency.py create mode 100644 docs/chameleon7b_rtx_sm87.md create mode 100644 docs/chameleon_thor_sm110.md create mode 100644 docs/chameleon_usage.md create mode 100644 examples/thor/chameleon_quickstart.py create mode 100644 scripts/bench_chameleon_thor.py create mode 100644 scripts/build_vqgan_trt.py create mode 100644 scripts/chameleon_orin_check.py create mode 100644 scripts/check_chameleon_thor_precision.py create mode 100644 scripts/profile_chameleon_thor.py create mode 100644 tests/test_chameleon_thor_vqgan_backend.py diff --git a/README.md b/README.md index 937bf5b5..5c5da722 100644 --- a/README.md +++ b/README.md @@ -1102,6 +1102,7 @@ examples/ - **Cosmos3-Nano text-to-video** (`config="cosmos3_video"`) — RTX 5090 BF16/FP8 denoise and complete benchmark workflow; [usage and performance](docs/cosmos3_video_usage.md) - **Cosmos3-Edge AV inverse dynamics and Reasoner** (`config="cosmos3_edge"`) — Jetson AGX Thor official baseline, 6.60x no-cache AV denoise, and NVFP4 multimodal chat decode; [complete usage and performance](docs/cosmos3_edge_thor.md) - **Qwen3-VL-8B** — RTX 5090 NVFP4/FP8 multimodal path, RTX 4090 official-FP8 path, and Jetson BF16 paths for Thor and Orin; [RTX 5090 usage](docs/qwen3_vl_nvfp4.md), [RTX 4090 usage](docs/qwen3_vl_fp8_sm89.md), [Jetson Thor usage](docs/qwen3_vl_thor.md), [Jetson Orin usage](docs/qwen3_vl_rtx_bf16.md) +- **Chameleon-7B** — Jetson Thor dynamic-FP8 prefill (~120 ms E2E, ~30 tok/s decode) and Jetson Orin INT8/QuaRot-INT4 path; [usage](docs/chameleon_usage.md), [Thor notes](docs/chameleon_thor_sm110.md), [Orin SM87 notes](docs/chameleon7b_rtx_sm87.md) - **MiniMax-Remover** — FP8 transformer + NVFP4 VAE video inpainting; [usage and performance](docs/minimax_remover_usage.md) - **MelBandRoformer** — kernelized FP8 audio source separation; [usage and performance](docs/melband_roformer_usage.md) - **OmniVoice TTS** — BF16/FP4 acceleration and HTTP serving; [serving quickstart](serving/omnivoice_agent/README.md) diff --git a/USAGE.md b/USAGE.md index 47e2414b..32b94bfd 100644 --- a/USAGE.md +++ b/USAGE.md @@ -681,6 +681,39 @@ once and reused, so per-chunk cost is just the infer row. This path does not change CMake targets, C++ bindings, or existing Pi0/Pi0.5/GROOT N1.6 runtime dispatch. +### Chameleon-7B (Thor) + +Standalone Chameleon-7B (text + image) is a direct-instantiation Thor +frontend — it is registered in `_PIPELINE_MAP` but is **not** dispatched by +`flash_rt.load_model` (same pattern as Qwen3-VL). The VQGAN image +tokenizer defaults to the generic eager Chameleon path; if compatible +TensorRT engines exist in the deployment, it is recommended to opt in +explicitly (`use_trt_vqgan=True`). A Jetson Orin (SM87) INT8/QuaRot +frontend is also available; see [`docs/chameleon_usage.md`](docs/chameleon_usage.md), +[`docs/chameleon_thor_sm110.md`](docs/chameleon_thor_sm110.md) and +[`docs/chameleon7b_rtx_sm87.md`](docs/chameleon7b_rtx_sm87.md). + +```python +from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + +fe = ChameleonTorchFrontendThor( + "/path/to/Chameleon_7B_mGPT", + use_fp8=True, # dynamic per-tensor FP8 (recommended default) + use_cuda_graph=True, + use_trt_vqgan=False, # generic default = eager VQGAN; set True when engines exist + target_size=512, +) + +out = fe.prefill("Describe the image.", [pil_image]) +# out["logits"]: (65536,) fp32 with mask_image_logits applied +``` + +FA4 attention is an explicit opt-in fast path +(`use_fa4_attn=True` / `FLASHRT_CHAMELEON_FA4_ATTN=1`, needs the +`thor-fa4` pip extra); measured transformer-prefill-only ≈ **104 ms** at +Se≈1056 vs ≈111 ms with CUTLASS FMHA. E2E (TRT VQGAN + FA4) ≈ **120 ms** +vs HF BF16 ≈ 403 ms transformer-only (~3.4×). + ### Wan2.2 TI2V-5B Wan2.2 TI2V-5B is exposed as an RTX SM120 official-pipeline baseline: diff --git a/benchmarks/chameleon_thor_latency.py b/benchmarks/chameleon_thor_latency.py new file mode 100644 index 00000000..04b95506 --- /dev/null +++ b/benchmarks/chameleon_thor_latency.py @@ -0,0 +1,180 @@ +"""Chameleon-7B (Thor sm_110) latency benchmark. + +Measures standalone Chameleon-7B prefill latency on real images with clean +stage separation: + +- ``--reuse-input-ids``: build real-image input ids once, then time only + embed + backbone + lm_head (transformer-prefill-only, HF-comparable). +- Default: full ``prefill()`` E2E including VQGAN tokenization. +- ``--use-trt-vqgan``: explicit TensorRT VQGAN opt-in (recommended when + compatible engines exist; the generic default stays eager). +- FA4 attention: enable with ``FLASHRT_CHAMELEON_FA4_ATTN=1`` (needs the + ``thor-fa4`` pip extra; prints whether it is active). + +Latency is wall-clock P50 (per CONTRIBUTING.md: quickstart --benchmark +style; CUDA-graph replayed latency is what the pipeline measures inside the +graph). Every result row records device, VQGAN backend, FA4 state, Se, +fp8/fp16 and graph settings for reproducible reporting. + +Usage: + + python benchmarks/chameleon_thor_latency.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image-dir /path/to/images \ + --iters 20 --warmup 5 +""" +from __future__ import annotations + +import argparse +import json +import pathlib +import time +from typing import Dict, List + +import torch + + +def _stats(xs: List[float]) -> Dict[str, float]: + a = sorted(xs) + n = len(a) + return {"mean": sum(a) / n, "p50": a[n // 2], "min": a[0], "max": a[-1]} + + +def _load_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + + paths = sorted( + p for p in image_dir.iterdir() + if p.suffix.lower() in (".jpg", ".jpeg", ".png", ".bmp")) + if not paths: + raise FileNotFoundError(f"No real images under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths], [str(p) for p in paths] + + +def _pad_ids(input_ids: List[int], pad_id: int = 1): + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _prefill_once(fe, prompt, images, cached_ids, use_graph: bool) -> Dict[str, float]: + """One timed prefill; returns stage latencies in ms.""" + times: Dict[str, float] = {} + t0 = time.perf_counter() + ids = fe.encode_prompt(prompt, images) if cached_ids is None else cached_ids + torch.cuda.synchronize() + times["encode_ms"] = (time.perf_counter() - t0) * 1000.0 + + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + torch.cuda.synchronize() + t1 = time.perf_counter() + + fe._embed_ids(padded) + torch.cuda.synchronize() + t2 = time.perf_counter() + + if use_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + torch.cuda.synchronize() + t3 = time.perf_counter() + + fe._project_last() + torch.cuda.synchronize() + t4 = time.perf_counter() + + times["prepare_ms"] = (t1 - t0) * 1000.0 - times["encode_ms"] + times["embed_ms"] = (t2 - t1) * 1000.0 + times["backbone_ms"] = (t3 - t2) * 1000.0 + times["lm_head_ms"] = (t4 - t3) * 1000.0 + times["transformer_ms"] = times["embed_ms"] + times["backbone_ms"] + times["lm_head_ms"] + times["total_ms"] = (t4 - t0) * 1000.0 + return times + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when engines are available; default is eager VQGAN)") + ap.add_argument("--reuse-input-ids", action="store_true", + help="time transformer prefill only (VQGAN excluded)") + ap.add_argument("--no-graph", action="store_true") + ap.add_argument("--use-fp16", action="store_true", + help="FP16 reference path instead of dynamic FP8") + ap.add_argument("--iters", type=int, default=20) + ap.add_argument("--warmup", type=int, default=5) + ap.add_argument("--output", default=None, help="JSON output path") + args = ap.parse_args() + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + from flash_rt.hardware.thor import fa4_backend + + images, image_paths = _load_images(pathlib.Path(args.image_dir), args.max_images) + use_graph = not args.no_graph + fp8 = not args.use_fp16 + + fe = ChameleonTorchFrontendThor( + args.checkpoint, use_fp8=fp8, use_cuda_graph=use_graph, + target_size=args.target_size, use_trt_vqgan=args.use_trt_vqgan) + + cached_ids = None + if args.reuse_input_ids: + cached_ids = fe.encode_prompt(args.prompt, images) + torch.cuda.synchronize() + + for _ in range(args.warmup): + _prefill_once(fe, args.prompt, images, cached_ids, use_graph) + + rows: Dict[str, List[float]] = {} + for _ in range(args.iters): + t = _prefill_once(fe, args.prompt, images, cached_ids, use_graph) + for k, v in t.items(): + rows.setdefault(k, []).append(v) + + result = { + "model": "chameleon-7b", + "device": torch.cuda.get_device_name(0), + "sm_count": torch.cuda.get_device_properties(0).multi_processor_count, + "checkpoint": args.checkpoint, + "image_paths": image_paths, + "prompt": args.prompt, + "target_size": args.target_size, + "fp8": fp8, + "cuda_graph": use_graph, + "vqgan_backend": fe.vqgan_backend, + "fa4_attn": fe.fa4_attn_active, + "fa4_status": fa4_backend.status(), + "reuse_input_ids": bool(args.reuse_input_ids), + "Se": int(fe.Se), + "latency_ms": {k: _stats(v) for k, v in rows.items()}, + } + for k, v in result["latency_ms"].items(): + print(f"[chameleon] {k:14s} p50={v['p50']:8.1f} ms mean={v['mean']:8.1f}") + print(f"[chameleon] device={result['device']} vqgan={fe.vqgan_backend} " + f"fa4={fe.fa4_attn_active} fp8={fp8} graph={use_graph} Se={fe.Se}") + + if args.output: + pathlib.Path(args.output).write_text(json.dumps(result, indent=2)) + print(f"[chameleon] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/docs/benchmark_comparison.md b/docs/benchmark_comparison.md index 50fad3f5..3294eead 100644 --- a/docs/benchmark_comparison.md +++ b/docs/benchmark_comparison.md @@ -255,6 +255,27 @@ TRT-aligned FP4 loop comparison, using the same quantization scheme. | 25 | ~304 ms | **97.5 ms** | **~3.1x** | | 50 | ~608 ms | **155.8 ms** | **~3.9x** | +## Chameleon-7B + +Baseline: HF `transformers` BF16 eager, transformer-only (input ids +built by FlashRT, model forward timed). FlashRT rows are the same harness +(real image `hand_1.jpg`, prompt "Describe the image.", target_size=512, +Se≈1053-1072, wall-clock P50). VQGAN backend and FA4 state are recorded per +row — generic default is eager VQGAN + CUTLASS FMHA; TRT VQGAN and FA4 are +explicit opt-ins. + +| FlashRT FP8 | VQGAN | FA4 | Latency | Speedup vs HF | +|---|---:|---:|---:|---:| +| transformer-only | (n/a) | on | **104.2 ms** | **3.9×** | +| transformer-only | (n/a) | off | 111.2 ms | 3.6× | +| E2E | TRT opt-in | on | **120.2 ms** | **3.4×** | +| E2E | eager | on | 177.3 ms | 2.3× | +| E2E | eager | off | ~190 ms | ~2.1× | + +| Baseline (HF BF16) | Latency | +|---|---:| +| transformer-only | 402.9 ms | + ## Qwen3-8B LLM rows list the baseline and FlashRT measurements without speedup. diff --git a/docs/chameleon7b_rtx_sm87.md b/docs/chameleon7b_rtx_sm87.md new file mode 100644 index 00000000..2c369f16 --- /dev/null +++ b/docs/chameleon7b_rtx_sm87.md @@ -0,0 +1,746 @@ +# Chameleon-7B VLM on Jetson AGX Orin (SM87) — FlashRT adaptation + +> **Status: Phase 1 complete — Gate 1 PASSES.** Production tier is +> **INT8 W8A8 + Hadamard (QuaRot at 8 bits)**: greedy output is **bit-identical +> to the HF bf16 reference for 16/16 tokens**, worst layer cosine 0.9986, +> last-row logit cosine 0.999968, at **21.07 tok/s decode and 273.8 ms warm +> prefill** (ISL=1032) — both at their measured ceilings (§4.2). Both defects found during bring-up were fixed by +> **quantization-method changes, not precision fallbacks**: a basis rotation for +> the massive-activation outliers (§4.5) and a ported clamp for the L31 FP16 +> overflow (§4.6). +> +> This is the authoritative document for **upstream Chameleon-7B as an image+text +> → text VLM** on Orin SM87. It is a knowledge migration from the RynnVLA-001 +> port in the derivative repo this was migrated from (same Chameleon-7B +> backbone, but *prefill-only* + action head); the cross-model mechanism +> analysis and the decode-regime methodology on the same platform that it built +> on are documented in that repo. + +--- + +## 0. Conclusion first + +**Shipped configuration** — `ChameleonTorchFrontendRtxSm87`, all defaults: + +| | | +|---|---| +| LLM GEMMs (Q/K/V/O, gate, up) | **INT8 W8A8 + Hadamard rotation** (QuaRot at 8 bits), per-row dynamic activation scales | +| FFN down | INT8 W8A8, per-row dynamic (K=11008 is not a power of two) | +| lm_head | INT8 W8A8 | +| residual / QK-LayerNorm / RoPE / attention / KV cache | FP16, with `ffn_down_clamp=60000` on the last 4 layers | +| attention | FA2 fp16 causal, `split_kv_bias=4` on decode | +| VQ-GAN encoder | FP16 convs, **fp32** codebook distance/argmin | + +**Result** (ISL=1032 = one 512² image + prompt, OSL=16, warm): + +| metric | value | +|---|---| +| greedy output vs HF bf16 reference | **bit-identical, 16/16 tokens** | +| worst per-layer residual cosine (L0..L31) | **0.9986** | +| last-row logit cosine | **0.999968** | +| decode | **21.07 tok/s** (47.5 ms/token) | +| prefill (LLM) | **273.8 ms** — GEMM at 91 % of the achievable CUTLASS ceiling | +| steady GPU memory | **7.6 GB** | + +**Why it works, in one line:** the Chameleon backbone's massive-activation +channels are fixed by a **basis rotation** rather than by more bits, finer +granularity, smoothing, or a per-layer precision fallback — and because the +rotation preserves per-row scales it reuses the stock CUTLASS INT8 GEMMs, so it +costs nothing (§4.5). + +### What this port had to add over RynnVLA-001 + +RynnVLA-001 shares this exact backbone but is **prefill-only** +(the RynnVLA pipeline in the derivative repo sets `layer_stride_llm = 0`, so all +32 layers share one K/V scratch — no KV cache, lm_head loaded but idle, no +sampler). So: + +* **Prefill / TTFT** — RynnVLA's whole Orin tuning ladder (L2 swizzle Id4, + stages-5, t256x128, vec8 elementwise) lives in the kernels and transferred for + free; measured 91 % of ceiling with zero new tuning. +* **Decode (M=1)** — a different regime (weight-bandwidth-bound, not + FLOPs-bound) with different levers. This was the new work, and it still needed + **zero new CUDA kernels** for the QK-Norm / RoPE / KV-write path (§3.1): the + K/V GEMMs write straight into the cache slab and the RoPE position rides in + the cos/sin table pointer. + +The one kernel this port *did* add is the INT8 twin of the existing FHT pack +(§4.5) — a device-side function inside `csrc/kernels/fht_int4.cu`, no new GEMM. + +## 1. Platform + +| Field | Value | +|---|---| +| Device | NVIDIA Jetson AGX Orin 64 GB (`torch.cuda.get_device_properties`: `Orin`, 61.4 GB) | +| GPU | SM **8.7** (Ampere), **16 SMs**, L2 4 MB | +| Memory | LPDDR5X unified, 204.8 GB/s spec | +| FP8 / FP4 | **not native** (Ada sm89+/Hopper/Blackwell only) → INT8/INT4 is the only low-bit route | +| CUDA / torch | 12.2 / 2.3.0 | +| Build | `cmake -B build -S . -DGPU_ARCH=87 -DFA2_ARCH_NATIVE_ONLY=ON -DFA2_HDIMS='128;256' -DFA2_DTYPES='fp16;bf16'` | + +**Measured bandwidth (this is the load-bearing calibration).** Three different numbers, and +picking the wrong one produces >100 % "efficiency" nonsense: + +| probe | result | use for | +|---|---|---| +| D2D copy, 512 MB, read+write | **124.5–126.0 GB/s** | copy-bound ops | +| single-stream vectorized reduce (`roofline.py --measure-bw`) | 99 GB/s | nothing — undersaturated | +| **best achieved by a real weight-streaming kernel** (int8 gate GEMM @ M=1) | **173.3 GB/s** = 85 % of spec | **the decode roofline denominator** | + +The RynnVLA-001 documentation in the derivative repo quotes 166 GB/s for a D2D +copy; we measure 124.5 in this container. Treat 173 GB/s (kernel-achieved, +read-dominated) as the decode ceiling. + +> ⚠️ `/sys/devices/gpu.0/devfreq/*/cur_freq` is **not readable in this container**, so clocks +> cannot be locked or even observed. Every number below is warm (≥30 warmup iters) and a median +> over ≥50 iters. Cross-config *ratios* are trustworthy; absolute values carry DVFS uncertainty +> (idle 306 MHz vs 1300.5 MHz loaded — principle #16). + +## 2. Phase 0 — probe verdicts + +### 2.1 R1 — FA2 split-KV is a silent no-op at 32 Q heads ⚠️ **and the fix is one Python argument** + +`csrc/attention/fa2_wrapper_causal.cu:41-43,152-158`: + +``` +num_splits = fa2_num_splits_heuristic_causal(batch*num_heads_q*num_m_blocks, num_sms*2, ...) + → if (batch_nheads_mblocks >= 0.8f * num_SMs) return 1; +``` + +Chameleon decode: `1*32*1 = 32` vs `0.8 * (16*2) = 25.6` → **`num_splits = 1`**. Passing the +accumulators does nothing. RynnBrain's +12 % split-KV win (documented in the +derivative repo) worked *only* because Qwen3-VL-2B has **16** Q heads. + +`num_sms` is a pure heuristic knob in this wrapper, so biasing it selects the split count. +Measured (q=1, kv=1040, 32 Q heads, head_dim 128, fp16): + +| `num_sms` passed | latency | speedup | max abs diff vs no-split | +|---|---|---|---| +| 0 (baseline, no accum) | 204.9 µs | 1.00× | — | +| 16 (**real SM count**) | 195.2 µs | 1.05× | **0.000e+00** ← proves `num_splits=1` | +| 32 | 149.8 µs | 1.37× | 1.221e-04 | +| **64** | **141.8 µs** | **1.44×** | 1.221e-04 | +| 128 | 154.8 µs | 1.32× | 1.221e-04 (over-split) | + +**Verdict: ship a `split_kv_bias` backend parameter, default 4× (`num_sms=64`).** 1.44× on +decode attention, pure Python, graph-safe. The 1.221e-04 delta is fp16 accumulation-order noise +(fp16 eps ≈ 9.8e-4 at magnitude 1), not an error. + +### 2.2 R2/R3 — M=1 GEMM: INT8 needs no GEMV; INT4 needs a small-M tile + +All shapes at M=1, achieved GB/s = weight bytes / time, normalized to the **173.3 GB/s** +kernel-achieved read ceiling. + +| shape | INT8 variant | µs | GB/s | %ceil | INT4 variant | µs | GB/s | %ceil | +|---|---|---|---|---|---|---|---|---| +| Q/K/V/O 4096×4096 | fp16out | 144.2 | 116.4 | 67 % | fp16out | 121.4 | 69.1 | **40 %** ⚠️ | +| gate 11008×4096 | bf16out | 260.2 | 173.3 | **100 %** | bf16out | 178.5 | 126.3 | 73 % | +| up+silu 11008×4096 | silu_gated | 273.3 | 165.0 | 95 % | silu_gated | 165.7 | 136.1 | 79 % | +| down 4096×11008 | fp16out | 267.3 | 168.7 | 97 % | fp16out | 170.7 | 132.1 | 76 % | +| lm_head 65536×4096 | bf16out | 1743.8 | 153.9 | 89 % | *(stays int8)* | — | — | — | +| **per token (GEMM only)** | | **47.70 ms** | | | | **33.45 ms** | | | +| **→ tok/s (GEMM only)** | | **21.0** | | | | **29.9** | | | + +**Verdict 1 — INT8: ship CUTLASS as-is, do not write a GEMV.** Four of five shapes are at +89–100 % of the achieved read ceiling. Measured 47.70 ms vs the predicted 44 ms weight floor — +**measured ≈ predicted, so the bottleneck model is correct** (principle #15). + +**Verdict 2 — INT4 delivers only 1.43×, not 2×.** Root cause confirmed in source: +`csrc/gemm/cutlass_sm80_int4_rowwise.cu:61` defines exactly one tile (`GemmShape<128,128,128>`) +with **no `M<=64` dispatcher**, whereas INT8 dispatches `M<=64 → 64×128` +(`cutlass_sm80_int8_rowwise_fp16out.cu:330-333`). At M=1 INT4 therefore wastes a 128-row tile — +visible as Q/K/V/O at **40 %** of ceiling (69.1 GB/s) versus INT8's 67 % (116.4 GB/s) on the +same shape with half the bytes. → **`cutlass_sm80_int4_rowwise_t64x128.cu` is justified**; +predicted recovery 33.45 → ~27 ms/token (~37 tok/s GEMM-only). + +**Verdict 3 — the planned "M=1 up-projection split" lever is DEAD. Dropped before writing any +code.** The hypothesis was that `cutlass_int8_silu_gated_bf16out` (128×128 tile only, +`cutlass_sm80_int8_silu_gated.cu:54`) would lose ~93 µs/layer at M=1 versus a +`bf16out` (64-tile) + `silu_mul_qwen36_bf16` split. Measured: **273.3 µs vs 260.3 µs = 13 µs**, +i.e. 0.4 ms/token ≈ 0.9 % — and the split adds a `silu_mul` launch plus an 11008-element bf16 +round trip that roughly cancels it. The prediction was **7× too optimistic**; both GEMMs are +already bandwidth-bound. (Principle #13: microbenchmark before writing the kernel.) + +### 2.3 R4 — HF reference: **PASS** + +Stock transformers 4.57.1 has `ChameleonForConditionalGeneration`, but its `ChameleonLayerNorm` +builds weights of shape `(num_heads, head_dim) = (32,128)` +(`transformers/models/chameleon/modeling_chameleon.py:187-202, 281-282`) while this Lumina-mGPT +checkpoint stores `(1,128)` — so it **cannot be loaded directly**. + +Working recipe (also: 4.57 *rejects* `state_dict=` together with a checkpoint path, so a +naked-constructor pattern is required — exactly the one used by the HF reference builder in +`scripts/chameleon_orin_check.py`): + +1. read both shards, `repeat_interleave(32, dim=0)` the **128** `self_attn.{q,k}_norm.{weight,bias}` tensors; +2. `torch.set_default_dtype(torch.bfloat16)`; `ChameleonForConditionalGeneration(cfg)` with `cfg._attn_implementation = "eager"`; +3. `load_state_dict(sd, strict=False)` → **0 missing, 0 unexpected** (548 tensors); `.eval().cuda()` → 13.1 GB. + +Verified in the same run: `mask_image_logits` is live — logits over ids **4..8195** come back at +`-3.390e+38` = `finfo(bf16).min`, while the text range max is `-13.375`. Since +`model_parallel_size == 1`, the expansion is a pure broadcast, so this reference is numerically +equivalent to official `facebook/chameleon-7b`. + +### 2.4 R5 — VQ-GAN codebook argmin precision + +`ChameleonVQVAE._from_config` + the 129 `model.vqmodel.*` tensors load with **0 missing / 0 +unexpected** (confirming the checkpoint is encoder-only and so is the HF module). Codebook index +match on a deterministic 512×512 input, versus a full-fp32 reference: + +| configuration | index match | +|---|---| +| fp32 convs + fp32 argmin | 100.00 % | +| **fp16 convs + fp16 argmin** | 98.14 % | +| **fp16 convs + fp32 distance/argmin** | 99.02 % | + +The fp32-argmin fix helps (`z²+e²−2ez` is cancellation-prone in fp16; +`modeling_chameleon.py:850-861`) and costs <0.1 ms. This probe used random noise; the +~92 % divergence recorded by the derivative repo's layerwise precision tests was on **real +images**, so re-measure on real content at Phase 5 before declaring the fix sufficient. + +### 2.5 R8 — decode-graph primitives are graph-safe: **PASS** + +Captured `index_select(emb, 0, tok, out=x)` → lm_head stand-in → `mask_view.fill_(bf16_min)` → +`argmax(out=)` → `tok.copy_(out_tok)`: capture succeeded (no `code=13`), the **stale-value test +passed** (changing the seed token changed the embedding output, maxdiff 4.85 — i.e. the graph +re-reads `tok` rather than baking it), and the logit mask survived replay. So no fp16 +embedding-lookup kernel is needed. + +### 2.6 R6 — deferred + +Original `original_tokenizers/vqgan.{yaml,ckpt}` vs HF `model.vqmodel.*` equivalence only gates +the **TRT engine** track (a TRT engine built from `vqgan.ckpt` must produce the same tokens as +the safetensors weights). Deferred to Phase 5. + +## 3. Design + +### 3.1 Zero new CUDA kernels for the QK-Norm / RoPE / KV path + +Three source facts combine to make the existing prefill kernel cover decode too: + +1. CUTLASS int8/int4 GEMM output row stride is hard-wired to `N` + (`cutlass_sm80_int8_rowwise_fp16out.cu:169-171`), so a `[32, max_seq, 32, 128]` fp16 KV cache — + whose per-layer slab is a contiguous `[max_seq, 4096]` with row stride exactly 4096 == N — is + a **legal GEMM destination**. `AlignmentC=8` (16 B) is satisfied by both the layer and row offsets. +2. `qk_norm_rope_fused_fp16` is in-place with implicit row stride `dim=128` and derives position + as `seq_pos = row / num_heads` (`qk_norm_rope_fused.cu:56-57, 65-66`) — so at `seq_len=1` every + row maps to row 0 of whatever cos/sin pointer it is handed. +3. The RoPE tables are C-contiguous `[max_seq, 128]` fp16 (`ChameleonTorchFrontendRtxSm87` + builds them that way), so position `pos` is `data_ptr() + pos*128*2` bytes. + +Therefore: + +* **prefill** — point the K and V GEMMs at `Kcache + li*layer_stride` / `Vcache + li*layer_stride`, + then call `qk_norm_rope_fused_fp16` unchanged (V needs no transform); +* **decode** — point them at `+ pos*4096*2` and call the same kernel with `seq_len=1` and cos/sin + pre-offset by `pos*128*2`. + +Also required: **`Se` must not be even-padded** (the RynnVLA pipeline in the derivative repo +pads `Se` for FP8 GEMM alignment) — with a real KV cache the pad row is junk that decode *will* +attend to, and CUTLASS constrains only `K`. + +Attention correctness: FA2 causal is **bottom-right aligned** +(`fa2_wrapper_causal.cu:126-138`), so `q=1, kv=N` attends all N keys. The cuBLAS fallback +`attention_mha_causal_fp16` is **top-left aligned** (`softmax.cu:182-191` masks with +`q = row % S_q`) and is therefore *silently wrong* at q=1 — the Chameleon backend must **raise** +rather than degrade to it. + +### 3.2 Precision policy + +| Component | Default (lossless tier) | Opt-in tier | +|---|---|---| +| Q/K/V/O, gate/up, down | INT8 W8A8, per-output-row weight scale, dynamic per-row act | QuaRot W4A4 (`use_int4`), down via block-diagonal `H_128` (`use_int4_down`) | +| lm_head (65536×4096) | INT8 (268 MB/token = 1.74 ms = 3.7 % of budget) | stays INT8 — never int4 | +| residual / RMSNorm / QK-LayerNorm / RoPE / attention / KV cache | FP16 | unchanged | +| VQ-GAN encoder | FP16 convs + **fp32 distance/argmin** | TRT FP16 (Phase 5) | + +Decode always uses **dynamic per-row** activation quant — never the prefill static calibration, +which was fitted at M=Se and does not describe a single decode row. + +### 3.3 Token contract (upstream Chameleon — do NOT reuse RynnVLA's ids) + +`[BOS 0] + n_img × ([8197 ] + [8711 ]×1024 + [8196 ]) + text + [8710 sep]`, +so `S = 1 + n_img*1026 + n_text + 1`. Image token id = **VQ codebook index + 4**, exactly, for +all 8192 codes; the 1024 tokens are a raster scan of the 32×32 latent grid. + +> ⚠️ **Trap:** the RynnVLA port in the derivative repo hardcodes +> `: 8710, : 8720`. Both are **wrong** for upstream Chameleon (8710 is the +> `sep_token`). Its `_init_special_token_ids` also sets ids 65536-65539, which are out of range +> for `vocab_size=65536`. This is one of three reasons the Chameleon frontend is standalone +> rather than a subclass of the RynnVLA frontend — see §4. +> +> ⚠️ `config.json` says `bos_token_id: 1`, which is **stale** (``); `tokenizer.json` gives +> ` = 0` and that is what the processor emits. + +### 3.4 Why the frontend is standalone (not a subclass) + +Three of the most attractive inheritable helpers from the derivative repo's RynnVLA port are +*actively wrong* for upstream Chameleon: its `_preprocess_image` is bicubic/384/`x*2-1` where +Chameleon needs PIL **LANCZOS**/512/`u8*0.0078-1.0` → `[-1, +0.989]`; its `_vqgan_encode` emits a +grid+newline token layout instead of a bare 1024 raster; its `_load_tokenizer` / +`_init_special_token_ids` produce the wrong ids above. Genuinely reusable are the quantizers and +`_split_fused_llm_weights` — extracted to `flash_rt/frontends/torch/_chameleon_quant.py`. + +## 4. Phase 1 — implementation and Gate 1 + +### 4.1 Shipped + +| file | role | +|---|---| +| `flash_rt/frontends/torch/_chameleon_quant.py` | checkpoint-agnostic INT8 / QuaRot-INT4 weight quantizers + the fused-projection split | +| `flash_rt/frontends/torch/_chameleon_spec.py` | weight spec — an inlined, Chameleon-specific `_llm_block()` (no dependency on any other model's spec) + `embed`/`norm`/`lm_head` singletons | +| `flash_rt/hardware/rtx/attn_backend_chameleon.py` | `ChameleonAttnBackend` — real per-layer FP16 KV cache, prefill + decode, `split_kv_bias` | +| `flash_rt/models/chameleon/pipeline_rtx.py` | one `chameleon_forward` serving both prefill (`pos=None`) and decode (`S=1`, `pos` set) | +| `flash_rt/frontends/torch/chameleon_rtx_sm87.py` | `ChameleonTorchFrontendRtxSm87` — `set_prompt` / `prefill` / `decode_step` / `generate` | +| `scripts/chameleon_orin_check.py` | Gate-1 harness (HF reference + graph safety + overflow guard) | +| `flash_rt/hardware/__init__.py`, `flash_rt/api.py` | dispatch entry + `_SM87_ALLOWED` + chat-VLM redirect | + +**Zero new CUDA kernels**, as predicted in §3.1. + +### 4.2 Measured performance (warm p50 over 10 iters after 2 discarded) + +| quantity | measured | predicted (§5) | verdict | +|---|---|---|---| +| **decode, ISL=1032** | **47.5 ms/token = 21.07 tok/s** | 53.7 ms = 18.6 tok/s | **beats prediction** -> bottleneck model correct | +| **prefill LLM, ISL=1032** | **273.8 ms** (min 270.6) | ~255 ms | **within 7 %** | +| prefill, plain INT8 (no rotation) | 282.2 ms | — | rotation is free at prefill too | +| prefill, **first call** | 460 ms | — | **1.68x cold-start penalty** — CUTLASS workspace `cudaMalloc` + JIT | +| load / steady memory | 29 s / **7.6 GB** | — | fp16 originals freed after quantization | + +> WARNING: an earlier revision of this doc claimed "prefill 496 ms, 1.9x worse +> than predicted — unexplained". That was **our measurement error**: the Gate-1 +> harness calls `prefill()` exactly once, so the number included first-call +> CUTLASS workspace allocation and JIT. **There is no prefill gap.** Principle +> #16 exists for exactly this reason — warm before every measurement. + +### 4.2.1 Per-kernel breakdown (torch.profiler, S=1032, post-warmup) + +Total GPU 281.1 ms (measured before the clamp restriction in §4.2.3): + +| kernel | ms | % | calls | +|---|---|---|---| +| CUTLASS INT8 GEMM — Q/K/V/O `(1032,4096,4096)` | 72.71 | 25.9 % | 128 | +| CUTLASS INT8 GEMM — up + fused SiLU-gate | 57.81 | 20.6 % | 32 | +| CUTLASS INT8 GEMM — gate (bf16 out) | 49.36 | 17.6 % | 32 | +| CUTLASS INT8 GEMM — down (t256x128) | 47.61 | 16.9 % | 32 | +| FA2 fp16 causal | 14.07 | 5.0 % | 32 | +| `residual_add_rms_norm_fht` (rotation fused into the norm) | 11.52 | 4.1 % | 63 | +| `qk_norm_rope_fused_fp16` | 9.36 | 3.3 % | 32 | +| `quantize_int8_rowwise_vec8` (bf16, pre-down) | 6.38 | 2.3 % | 32 | +| `clamp_inplace_fp16` | 6.23 | 2.2 % | 32 | +| `fht_int8_quant` (pre-O) | 3.90 | 1.4 % | 32 | +| lm_head + tail | 2.17 | 0.8 % | 4 | + +=> **GEMM 229.2 ms (81.5 %)**, elementwise tail 37.8 ms (13.5 %), attention 14.07 ms (5.0 %). + +### 4.2.2 The real GEMM ceiling is 64.4 TOPS, not 84.8 — GEMM tuning is spent + +The 84.8 TOPS quoted by the RynnVLA-001 documentation in the derivative repo is +the **raw `mma.s8` issue rate** from a register-only probe. What CUTLASS actually +achieves on its best-case shape is lower: big-square probes measure **58.7 TOPS +at 4096^3 and 64.4 TOPS at 8192^3**. Against that realistic ceiling: + +| shape | ms/call | TOPS | vs 64.4 ceiling | +|---|---|---|---| +| Q/K/V/O `(1032,4096,4096)` | 0.567 | 61.1 | **95 %** | +| gate / up `(1032,11008,4096)` | 1.694 | 54.9 | 85 % | +| down `(1032,4096,11008)` | 1.540 | 60.4 | 94 % | +| **whole LLM** | **229.2** | **58.7** | **91 %** | + +The prefill GEMMs are at **91 % of the achievable CUTLASS ceiling**, and the +isolated probe reproduces the in-pipeline time to within 0.5 % (0.567 vs +0.568 ms on Q/K/V/O) — so there is no pipeline overhead left to recover. This +independently confirms RynnVLA-001's conclusion that the GEMM ladder (swizzle +Id4 / stages-5 / t256x128) is spent. Only `gate/up` at 85 % shows slack, and +RynnVLA already swept tiles there and measured 256x128 as "only ~2 % better, not +worth a 4th instantiation". + +> WARNING: **the roofline probe itself had a DVFS bug**, found here. Whichever +> shape was measured *first* was penalised by clock ramp: Q/K/V/O reported +> **28.3 TOPS** measured first versus **61.1** for the identical shape after +> adding a 3-second saturating pre-ramp, and the per-shape TOPS ascended purely +> in measurement order (28.3 -> 53.8 -> 60.5). `_ramp_clocks()` now runs before +> any timing in the roofline script (in the derivative repo, +> `scripts/bench/orin_int8_roofline.py`). Any earlier per-shape number from that +> script is suspect. + +### 4.2.3 Clamp restricted to the last 4 layers: -6.3 ms + +`clamp_inplace_fp16` cost 6.23 ms (2.2 %) across all 32 layers, but the measured +magnitudes (§4.6) grow monotonically with depth and L28 is 1616 — **37x below the +60000 clamp** — so early layers can never reach it. Restricting it to the last +`ffn_down_clamp_last_n` layers (default 4): + +| | before | after | +|---|---|---| +| prefill warm p50 | 280.1 ms | **273.8 ms** | +| decode | 20.96 tok/s | **21.07 tok/s** | +| L31 / final / logit cosine, greedy text | 0.999722 / 0.999447 / 0.999968 / 16-of-16 | **bit-identical** | + +The Gate-1 harness reports per-layer clamp saturation, so a checkpoint that +violates the monotonicity assumption is detectable; set +`FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N=32` if that ever happens. + + +### 4.3 Gate 1 — PASS + +Real image (`FlashRT.png`), `"Describe this image."`, ISL=1032, OSL=16, +tier **INT8+Hadamard**: + +| check | result | gate | verdict | +|---|---|---|---| +| **greedy text identical to HF** | **16/16 tokens** | 16/16 | **PASS** | +| worst layer cosine (L0..L31) | **0.9986** | ≥0.97 | PASS | +| L31 / final-norm cosine | **0.999722 / 0.999447** | — | PASS | +| last-row logit cosine | **0.999968** | ≥0.999 | PASS | +| graph safety (capture + stale-value) | cos 0.9884 between two seed tokens | not frozen | PASS | +| FP16 residual finiteness | no inf/nan | finite | PASS | +| argmax, image positions (1026) | 89.6 % exact / **95.0 % tie-adjusted** | — | informational (§4.7) | +| argmax, text positions (6) | 5/6 | — | **not binding** — n=6 is too small; one near-tie flip moves it 17 points | + +Both engines produce `"The image is a logo for the company Flexsteel. The logo is a"`. + +### 4.4 Root cause: the row-0 massive activation + +Probing every layer 20-31 against the reference localizes the failure precisely. +It is **not** spread across the tensor — it is **row 0, the BOS/attention-sink +token**, in the last four layers: + +| layer | cosine (all rows) | worst-row cosine | worst row | FlashRT max\|x\| | ref max\|x\| | +|---|---|---|---|---|---| +| L24 | 0.9966 | 0.982 | row4 | 1993 | 2512 | +| L27 | 0.9922 | 0.994 | row4 | 1990 | 2512 | +| **L28** | 0.850 | **0.688** | **row0** | 240 | 1632 | +| **L31** | 0.691 | **−0.384** | **row0** | 3502 | 23936 | + +(ISL=20 text-only.) The reference's L31 row-0 norm is **42971 vs a median of +10456**, concentrated in a few channels — **d632 = 23936**, then d808, d1282, +d2669. FlashRT has 1225 in d632. + +Mechanism: per-row INT8 activation quantization sets `scale = amax/127` from that +outlier, so the other ~4090 channels of row 0 round to zero and the row's +direction is destroyed (cosine goes *negative*). This is the documented +Chameleon massive-activation zone (the skill's backbone profile names d671/d579 +for the L15→L19 band; here it is d632 in the L28→L31 band), and per principle +#3/#17 the fix is **basis rotation, not smoothing**. + +### 4.5 The tier ladder: rotation × bit-width (W8A8+Hadamard wins) + +Two independent error sources act here, and each shipped tier only fixed one: + +* **outlier conditioning** — a row whose amax is set by a massive-activation + channel loses its other ~4090 channels to rounding. Fixed by a **basis + rotation**, not by finer granularity or smoothing (principle #17). +* **quantization noise** — the resolution left for the other 1031 ordinary rows. + Fixed by **more bits**. + +Measured on the same prompt, all three tiers: + +| ISL | tier | L24 | L28 | L31 | final | last-row logit cos | greedy prefix | +|---|---|---|---|---|---|---|---| +| 7 | INT8 (per-row) | 0.9987 | 0.699 | 0.508 | 0.794 | 0.998850 | 8/12 | +| 7 | INT4 (QuaRot) | 0.9985 | **0.981** | **0.960** | 0.955 | 0.999484 | 8/12 | +| 1032 | INT8 (per-row) | 0.9952 | 0.9953 | 0.9983 | 0.9969 | 0.999916 | 8/16 | +| 1032 | INT4+down | 0.9334 | 0.9313 | — | — | 0.998881 | **0/16** | +| **1032** | **INT8+Hadamard** | **0.9989** | **0.9989** | **0.99972** | **0.99945** | **0.999968** | **16/16** | + +> This **contradicts the RynnVLA-001 Orin conclusion** ("both INT4 tiers beat +> INT8 at every layer probe on every frame"). That doc is not wrong — it +> measured a *prefill-only VLA at fixed Se*; the verdict is ISL-dependent, and +> a VLM's production ISL sits in the opposite regime. + +So the INT8-vs-INT4 verdict *inverts with sequence length* — at short ISL the +sink row is 1/7 of the tensor and rotation dominates; at long ISL it is 1/1032 +and 4-bit noise dominates. That inversion is the tell that the two tiers were +each solving half the problem. **Rotating at 8 bits solves both and strictly +dominates**, which is why it is the default. + +Cost: **one new device-side pack function** (`quant_int8`) inside the existing +`csrc/kernels/fht_int4.cu`, plus templating its three kernels on the output +width — the norm and the radix-16 register FHT are shared verbatim with the INT4 +path. **No new GEMM**: because the rotation keeps plain per-row scales, the +unmodified `cutlass_int8_rowwise_*` kernels consume the rotated activations +directly. Measured decode **20.96 tok/s vs 19.9** for plain INT8, i.e. no +throughput cost (the FHT rides inside an already-optimized fused norm kernel; +the difference is within the ±3 % process-to-process variance this platform +shows). + +The weight side folds offline (`W_rot = H·W/√K`, `quantize_int8_hadamard`); the +activation side is fused into the norm (`rms_norm_fht_int8_fp16`, +`residual_add_rms_norm_fht_int8_fp16`, `fht_int8_quant_fp16`). The FFN **down** +projection stays plain INT8: K=11008 is not a power of two and its input is the +un-rotated BF16 SiLU output. + +**Why not the alternatives** (principle #17's measured ladder on this backbone): +SmoothQuant reached only 0.641 and outlier-splitting 0.970 on the A4 variant of +this problem, while group-128 / block-scaled schemes need a bespoke GEMM whose +hand-written ceiling on 16-SM Orin measured just 41 TOPS. A per-layer FP16 +fallback would also have worked, but it is checkpoint-specific tuning that +permanently costs throughput — the rotation is free and generalizes. + +### 4.6 SOLVED — the FP16 overflow, via RynnVLA-002's `ffn_down_clamp` + +With a real image the reference's L31 residual reaches **max|x| = 89088**, above +FP16's 65504, so FlashRT stored `inf` and the final RMSNorm turned that row's +logits into `nan`. It affects **both** precision tiers — it is a property of the +residual *dtype*, not of the quantization. + +The first instinct (a BF16 residual stream, ~1 new kernel) was **wrong** — the +answer already existed in the lineage. The derivative repo's Chameleon +acceleration-methodology documentation records this exact failure for the +Chameleon backbone, and its FP8 optimization playbook had already flagged the +missing clamp as a *latent, unverified* risk for RynnVLA-001: + +> **可复用到 001**：001 当前**无 clamp**，是潜在的 inf 风险点（尤其长序列）。直接移植 +> `clamp_inplace_fp16` 即可。 + +This port empirically confirmed that prediction. + +**Why a clamp is sufficient** — measured per-layer magnitudes in the bf16 +reference (ISL=1032). The explosion is confined to **exactly one layer**: + +| quantity | L28 | L29 | L30 | **L31** | +|---|---|---|---|---| +| residual | 1616 | 1720 | 2032 | **266240** | +| o_proj output | 76 | 78 | 80 | 1056 | +| down **input** (gu) | 1128 | 1528 | 6080 | **151552** | +| down **output** | 1120 | 1032 | 1880 | **264192** | + +Because the pre-L31 residual is only ~2032, clamping the down **output** at +60000 leaves the residual at ~62000 < 65504. So one `clamp_inplace_fp16` +(already in `flash_rt_kernels`, CUDA-Graph safe) removes the overflow with +**zero new kernels and no dtype change**. + +Unlike RynnVLA-002's Thor path we do **not** need to clamp the down *input*: +ours is BF16 (`cutlass_int8_silu_gated_bf16out`), whose range absorbs 151552 +without issue. The clamp is applied on every layer because L0-L30 are three +orders of magnitude below it and therefore untouched; cost is 32 extra +elementwise launches (<0.3 % of the decode budget, ~0.7 % of prefill). + +**Result** (ISL=1032, INT8, real image): + +| | before | after | +|---|---|---| +| L31 cosine | `nan` (inf) | **0.998266** | +| final-norm cosine | `nan` | **0.996923** | +| L31 max\|x\| | `inf` | 60160 (saturating at the clamp, as intended) | +| last-row logit cosine | 0.999916 | 0.999916 | +| greedy prefix vs HF | 8/16 | 8/16 | + +Exposed as `ffn_down_clamp` (default 60000, env `FLASHRT_CHAMELEON_DOWN_CLAMP`), +named after the corresponding RynnVLA-002 env var in the derivative repo. + +⚠️ **The clamp did not change the text divergence** (still 8/16). That confirms +the overflow was confined to the sink row's post-L31 residual, which feeds only +that row's final norm — so it was never the cause of the divergence. The +remaining gap is ordinary INT8 error at high-confidence text decisions and is +still open; see §6 for the ranked options inherited from the RynnVLA lineage. + +Also note the **gate itself was wrong** at first: an absolute +"max|x| < 30000" threshold fails by construction on a backbone whose reference +legitimately runs at 2.6e5. The correct gate is **finiteness**, with clamp +saturation reported as information. + +### 4.7 Measurement-hygiene finding: argmax-over-all-positions is meaningless here + +1024 of the 1032 teacher-forced positions are **image** positions. At those the +model predicts a next token while all 8192 image ids are masked out of the +logits (§3.3), so the winner is an arbitrary low-confidence text token — median +reference top1−top2 gap **0.250** on a logit scale of ~20, versus **0.895** at +the 6 text positions. An unsplit "argmax match = 87.21 %" therefore says almost +nothing about generation quality. The harness now reports image and text +positions separately and gates only on text positions, and additionally +classifies a mismatch as a **BF16 tie** when the reference's top-2 gap is within +one BF16 ULP (58 of the 132 mismatches were ties). + +### 4.8 Two harness traps worth remembering + +* **A forward hook that returns a value replaces the module output.** Using + `dict.setdefault(...)` inside a `register_forward_hook` lambda returns the + stored tensor, which silently substituted a *CPU* tensor for + `model.model.norm`'s output and crashed `lm_head` with a device mismatch. + Always `return None`. +* **Launching on stream 0 while another stream is capturing silently drops the + kernels from the graph.** The first graph-safety run reported + `stale-value: FAIL (frozen)` with cos exactly 1.0000 — not because anything + was baked, but because only the torch ops got captured and none of the `fvk` + kernels did. `decode_step` now takes an explicit `stream` argument. + +## 5. Roofline ladder — predicted vs finally measured (principle #15) + +Decode reads 6.745 G params/token (32 layers 6.476 G + lm_head 0.268 G). **MHA +with 32 KV heads makes the KV cache 4x heavier than a GQA model** — 0.524 MB per +token of context, so 0.55 GB at S=1040 and **2.15 GB at S=4096, where KV would +dominate an int4 tier.** + +| tier | weight floor @173 GB/s | + KV @S~1040 | GEMM-only µbench | predicted total | **finally measured** | +|---|---|---|---|---|---| +| **int8 (+Hadamard, shipped)** | 39.0 ms | +3.2 ms | 47.70 ms | 53.7 ms = 18.6 tok/s | **47.5 ms = 21.07 tok/s** | +| int4 (as built) | 19.5 ms | +3.2 ms | 33.45 ms | 39.5 ms = 25.3 tok/s | not shipped (loses on precision, §4.5) | + +Decode came in **13 % better than predicted** — the prediction charged full price +for attention and the elementwise tail, but `split_kv_bias` (§2.1) and the fused +FHT norm absorb part of it. A prediction that is close *and* slightly pessimistic +is the sign the bottleneck model is right (a large gap in either direction would +mean the model of the bottleneck is wrong, not that there is tuning left). + +**Prefill**: predicted ~255 ms by scaling RynnVLA's measured Se=1214 numbers to +Se=1032; **measured 273.8 ms warm** (within 7 %), of which GEMM is 229.2 ms at +**91 % of the achievable CUTLASS ceiling** (§4.2.2). Image tokenize adds ~53 ms +(PyTorch VQ-GAN; ~27 ms with a 512x512 TRT engine, not built). + +> Superseded numbers, kept so they don't re-mislead: an earlier revision of this +> section predicted **18.6 tok/s** decode and this doc once reported **496 ms** +> prefill and an **84.8 TOPS** GEMM target. Current values: **21.07 tok/s**, +> **273.8 ms**, and a **64.4 TOPS** achievable ceiling. See §4.2 for why the +> 496 ms was a cold-start artifact and §4.2.2 for the ceiling correction. + +## 6. Ranked lever menu + +### Precision — status: closed + +| # | lever | outcome | +|---|---|---| +| 1 | **W8A8 + Hadamard (QuaRot at 8 bits)** | **DONE — this closed it.** greedy 8/16 → **16/16**, worst layer 0.9946 → 0.9986, last-row logit 0.999916 → 0.999968, at no throughput cost. Default tier. | +| 2 | `ffn_down_clamp` (ported from RynnVLA-002) | **DONE** — removed the L31 FP16 `inf` (§4.6) | +| ~~3~~ | ~~Tier-3 FP16 fallback for L31~~ | **not needed** — the rotation fixed the same layer at 8 bits. A per-layer precision fallback is checkpoint-specific tuning and costs throughput permanently; prefer the quantization method. | +| ~~4~~ | ~~AWQ / SmoothQuant per-K smoothing~~ | **not needed** — and principle #17's measured ladder on this backbone puts smoothing (0.641) far below rotation (0.9914). Kept only as a fallback if a future checkpoint defeats rotation. | +| 5 | ISL-adaptive tier selection | **obsolete** — W8A8+Hadamard wins at both short and long ISL, so there is nothing to switch between | +| ~~6~~ | ~~BF16 residual stream~~ | **superseded by the clamp** — would have cost a new `qk_norm_rope_fused_bf16` kernel to fix what one existing elementwise kernel already fixes | + +**Decode (M=1, weight-bandwidth-bound)** + +| # | lever | predicted | effort | status | +|---|---|---|---|---| +| 1 | `use_int4` / `use_int4_down` | **1.43×** (measured, not 2× — §2.2) | trivial, already built | Phase 2 | +| 2 | `cutlass_sm80_int4_rowwise_t64x128.cu` + `M<=64` dispatcher | int4 33.45 → ~27 ms = **+20 %** | ~150 lines + 1 CMake line | Phase 3 — **justified by §2.2** | +| 3 | `split_kv_bias = 4` (`num_sms=64`) | attention **1.44×** = +1.7 ms/token (+3–5 %), more at long S | 1 Python arg | Phase 1 — **measured (§2.1)** | +| 4 | Per-position decode CUDA graph | 0–15 % throughput, −6 s startup | medium | Phase 4 | +| 5 | INT8 Q/K/V/O at 67 % of ceiling (small-N tail: 4096/128 = 32 tiles on 16 SMs) | up to +12 % if it reached 100 % | high (hand GEMV) | open | +| 6 | Devpos kernel + fp16 seqused-splitkv FA2 → *one* decode graph | 0 % throughput; removes capture cost + `max_new_tokens` cap | high (1 `.cu` + FA2 rebuild) | deferred | +| 7 | Reduced lm_head (drop rows 4..8195) | +0.5–1 % | low | deferred | +| 8 | INT8 KV cache | +3.8 % @1040, **+12 % @4096** | high (needs a 32Q/32KV variant; RynnBrain measured break-even) | S≥4096 only | +| ~~9~~ | ~~M=1 up-projection split~~ | ~~+7 %~~ → **measured 0.9 %, net ≈0** | — | **DEAD (§2.2)** | +| 10 | Speculative decode | 1.5–2× | N/A — no draft model | — | + +**Prefill / TTFT (FLOPs-bound — levers do not transfer)** + +| # | lever | predicted | status | +|---|---|---|---| +| 1 | `use_int4` / `use_int4_down` | LLM ~265 → ~165 ms (**−38 % TTFT**) | already built | +| 2 | 512×512 TRT VQ-GAN engine | 53 → 27 ms (−9 % TTFT) | `scripts/build_vqgan_trt.py`, inputs present in `original_tokenizers/` | +| 3 | GEMM-util / elementwise fusion | **≈0** — measured spent at 74 % of the 84.8 TOPS mma peak | closed | +| 4 | Per-Se prefill CUDA graph | 0–5 %, and a full capture per new prompt length; `Se` cannot be bucketed (padding poisons the KV cache) | **reject for a VLM** | + +## 7. Dead-ends (measured — do not re-walk) + +| direction | result | one-line reason | +|---|---|---| +| FA2 split-KV with the real `num_sms=16` | **bit-identical, 1.05×** | `32 >= 0.8*32` → `num_splits=1`; the heuristic disables itself at 32 Q heads | +| M=1 up-projection split (`bf16out` + `silu_mul`) | 13 µs/layer ≈ 0.9 %, net ≈0 | both GEMMs already bandwidth-bound; the extra launch + bf16 round trip cancels it | +| INT4 at M=1 expecting 2× | **1.43×** | single 128×128 tile, no `M<=64` dispatcher | +| **`use_int4` / `use_int4_down` as the VLM default** | **L24/L28 cosine 0.933 vs INT8's 0.995; greedy prefix 0/16 vs 8/16** | at production ISL the sink row is 1/1032 of the sequence, so INT8's per-row damage is diluted and 4-bit noise on the other 1031 rows dominates (§4.5). INT4 still wins at short ISL — the verdict is ISL-dependent | +| **unsplit argmax match as a precision metric** | "87.21 %" says nothing | 1024/1032 teacher-forced positions are image positions whose logits are fully masked → arbitrary low-confidence winners (§4.7) | +| forward hook capturing via `dict.setdefault` | CPU/CUDA device crash in `lm_head` | a hook returning non-None **replaces** the module output | +| launching pipeline kernels on `stream=0` during graph capture | `stale-value: FAIL (frozen)`, cos exactly 1.0000 | kernels on the default stream are silently *not* recorded; only the torch ops were captured | +| `attention_mha_causal_fp16` for decode | silently wrong | top-left-aligned causal mask; at `S_q=1` only column 0 survives | +| stock `from_pretrained(..., state_dict=...)` | `ValueError` in 4.57 | use naked ctor + `load_state_dict` | +| stock transformers loading this ckpt unmodified | shape mismatch on 128 tensors | `ChameleonLayerNorm` wants `(32,128)`, ckpt has `(1,128)` | +| **measuring prefill on the first call** | 460-496 ms vs 273.8 ms warm | 1.68x cold-start penalty from CUTLASS workspace `cudaMalloc` + JIT; produced a phantom "1.9x prefill gap" | +| **roofline probe without a clock pre-ramp** | first shape measured reads 28.3 TOPS vs 61.1 warm | Orin DVFS ramps 306 -> 1300 MHz; per-shape TOPS ascend in measurement order | +| clamping the down output on all 32 layers | 6.23 ms (2.2 %) for no effect on 28 of them | magnitudes grow monotonically with depth; L28 is 37x below the clamp | +| trusting 84.8 TOPS as the GEMM target | it is the **raw mma issue rate**, not achievable | CUTLASS peaks at 64.4 TOPS big-square on this part; we are at 91 % of *that* | +| single-stream reduce as a bandwidth probe | 99 GB/s | undersaturated; use a real weight-streaming kernel (173 GB/s) | +| int8 `sum(dtype=int64)` as a read-BW probe | 9.3 GB/s | ALU-bound reduction, not a bandwidth measurement | + +## 8. Reproduction + +```bash +# Gate 1: correctness vs the HF bf16 reference + graph safety + fp16 health +PYTHONPATH=. python3 scripts/chameleon_orin_check.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image FlashRT.png --prompt "Describe this image." --steps 16 +# ... --text-only fast smoke, no image +# ... --int4 / --int4-down alternative tiers +# ... --vq-fp16-argmin measure VQ index drift instead of avoiding it + +# M=1 decode roofline (no checkpoint) — the "do we need a GEMV?" gate +# (the roofline probe script lives in the derivative repo this was migrated +# from: scripts/bench/orin_int8_roofline.py) +python3 scripts/bench/orin_int8_roofline.py --decode + +# large-M prefill roofline at the production shape +python3 scripts/bench/orin_int8_roofline.py --M 1032 +``` + +Minimal use: + +```python +from flash_rt.frontends.torch.chameleon_rtx_sm87 import ChameleonTorchFrontendRtxSm87 +f = ChameleonTorchFrontendRtxSm87("/path/to/Chameleon_7B_mGPT") # int8+hadamard +f.set_prompt("Describe this image.", images=[pil_img]) +print(f.generate(max_new_tokens=32)) +``` + +`load_model(config="chameleon")` deliberately raises with this snippet: it is a +chat VLM (`set_prompt` + `generate`), not the VLA `predict()` surface. + +**Build** (the INT8 FHT kernels ride in the existing `ENABLE_SM80_INT8_CUTLASS` +source list, so no CMake change is needed): + +```bash +cmake -B build -S . -DGPU_ARCH=87 -DFA2_ARCH_NATIVE_ONLY=ON \ + -DFA2_HDIMS='128;256' -DFA2_DTYPES='fp16;bf16' +cmake --build build --target flash_rt_kernels -j6 +``` + +## 9. Constructor parameters (classified: required / recommended / experimental / informational) + +`ChameleonTorchFrontendRtxSm87(checkpoint_dir, **kwargs)`. There is no CLI/server +entry yet (out of scope this round), so these are the frontend kwargs. + +**Required** + +| param | notes | +|---|---| +| `checkpoint_dir` | path to the Chameleon-7B checkpoint. Backbone dims and the special token ids are **hard-asserted** against its `config.json` | + +**Recommended (safe defaults; tune for deployment)** + +| param | default | notes | +|---|---|---| +| `max_seq` | 2048 | sizes the KV cache (`2 × 32 × max_seq × 4096 × 2 B` = 2.15 GB at 2048) and the RoPE tables. Hard-checked against `max_position_embeddings=4096`. `S = 1 + n_img*1026 + n_text + 1` | +| `use_hadamard` | **`True`** | the W8A8+QuaRot tier. **Do not disable in production** — plain per-row INT8 reproduces only 8/16 reference tokens (§4.5). Kept switchable for A/B only | +| `split_kv_bias` | `4` | multiplies the `num_sms` passed to FA2 so split-KV actually engages at 32 Q heads (§2.1). `1` disables | +| `ffn_down_clamp` | `60000` | **correctness requirement**, not a knob (§4.6). Env: `FLASHRT_CHAMELEON_DOWN_CLAMP` | +| `ffn_down_clamp_last_n` | `4` | layers to clamp, counted from the end. `32` = all layers (safe but costs 2.2 % of prefill). Env: `FLASHRT_CHAMELEON_DOWN_CLAMP_LAST_N` | +| `vq_argmin_fp32` | `True` | fp32 codebook distance/argmin; costs <0.1 ms and lifts index match 98.1 % → 99.0 % (§2.4) | +| `free_fp16_weights` | `True` | drop the 13 GB fp16 originals after quantization | + +**Experimental (off by default)** + +| param | default | notes | +|---|---|---| +| `use_int4` | `False` | QuaRot W4A4 on the six K=4096 projections. Wins only at very short ISL; **below `use_hadamard` at production ISL** (§4.5) | +| `use_int4_down` | `False` | additionally int4 the FFN down via block-H128. **Not recommended** — worst tier measured at production ISL (0/16 greedy) | +| `probe_layers` | `None` | list of layer indices to snapshot post-residual hidden states for `snapshot_probe()`; zero cost when `None` | + +**Informational** + +| param | notes | +|---|---| +| `precision_tier` / `precision_spec()` / `get_model_info()` | report the resolved configuration; `timing` carries `prompt_ms` / `prefill_ms` / `decode_tok_s` | + +**Not accepted** (unlike the RynnVLA frontends): `use_fp8`, `use_fp4`, +`use_fp8_attn`, `use_awq_v_proj`, `num_views`, `action_dim`, +`action_chunk_size`, `state_dim` — SM87 has no FP8/FP4 tensor cores and this is +not a VLA. Unknown kwargs are swallowed by `**_ignored`. + +## 10. Tier status (all four verified to run and produce finite logits) + +| tier | flag | verdict | +|---|---|---| +| **int8+hadamard** | default | **production** — Gate 1 PASS, 16/16 greedy match, 21.07 tok/s | +| int8 plain | `use_hadamard=False` | works; loses the outlier conditioning (8/16 greedy) — kept for A/B | +| int4 (QuaRot) | `use_int4=True` | works; best at very short ISL, below int8+hadamard at production ISL | +| int4+down | `use_int4_down=True` | works but **not recommended** — worst at production ISL (§4.5) | diff --git a/docs/chameleon_thor_sm110.md b/docs/chameleon_thor_sm110.md new file mode 100644 index 00000000..af665b6b --- /dev/null +++ b/docs/chameleon_thor_sm110.md @@ -0,0 +1,344 @@ +# 标准 Chameleon-7B @ Thor SM110 — 权威文档 + +**平台**: Jetson AGX Thor (SM110, aarch64) · CUDA 13.0 · transformers 4.43+ +**模型**: 标准/独立 Chameleon-7B(纯 LLM 主干 + VQGAN 图像 tokenizer,**无 ActionHead / ActionVAE**) +**生产方案**: 全 32 层**运行时动态 per-tensor FP8**(实现位于 Chameleon 专用的 `flash_rt/models/chameleon/pipeline_thor.py::chameleon_forward`,迁移自 derivative repo 的 RynnVLA-002 Thor 移植)+ 通用 eager Chameleon VQGAN 默认路径 + cuBLASLt 逐 shape autotune + L31 selective clamp;TensorRT VQGAN 仅显式 opt-in +**版本**: v1.4 (2026-08,新增 KV-cache 增量解码 `generate_greedy`:30.4 tok/s,逐 token 与全前缀重算 oracle 一致) + +> 与 RynnVLA-002 Thor 的关系:本文档描述的是**独立/标准 Chameleon-7B**(纯文本+图像对话骨干,直接 frontend `encode_prompt`/`prefill`/`generate_greedy`,不是 VLA `predict()` 接口),没有 ActionHead/ActionVAE。其动态 FP8 Chameleon 主干实现与 Thor attention backend 迁移自 derivative repo 中 RynnVLA-002 的移植。 + +--- + +## 0. 结论先行 + +- **资产路径**:`/path/to/Chameleon_7B_mGPT`(注意实际目录名是 `mGPT` 不是 `mGP`)。包含权重 shards、tokenizer、`original_tokenizers/vqgan.{yaml,ckpt}`。 +- **HF 直接加载会失败**:当前 `transformers` 的 `ChameleonForConditionalGeneration.from_pretrained` 在该 checkpoint 上因 `q_norm`/`k_norm` 形状为旧式 `[1,128]`(而不是新式 `[32,128]`)而报错/静默错载。**变通方案**:生产路径直接走 FlashRT 自己的声明式 `WeightLoader`(完全绕开 HF `from_pretrained`);或如 `scripts/check_chameleon_thor_precision.py` 那样,用裸 `ChameleonForConditionalGeneration` 构造器 + `load_state_dict(strict=False)` 加载 HF 参考模型(上游没有 vendor 目录,脚本直接从 `transformers` 导入 `ChameleonForConditionalGeneration`,可用 `--skip-hf` 跳过 HF 参考比对)。 +- **精度已验证(真实图片,非合成 token id)**: + - FlashRT FP16 vs HF BF16(last-token logits cosine,mask_image_logits 后):**0.9999997**,greedy next-token 完全一致。 + - FlashRT 动态 FP8 vs FlashRT FP16:**0.99999999**,greedy next-token 完全一致,top-10 overlap 1.0。 +- **VQGAN backend policy(框架定位)**:FlashRT 面向**通用/标准 Chameleon**时保持框架通用性——VQGAN **默认走 eager** Chameleon tokenization(`use_trt_vqgan=False`),不默认依赖 RynnVLA/TensorRT engine,保证框架自身能力可独立运行。**若部署环境存在可用 TRT engine,建议显式开启**(`use_trt_vqgan=True` 或脚本 `--use-trt-vqgan`;实测 VQGAN 74.9→17.3ms,TRT E2E ~121ms vs eager ~190ms)。这与此前 RynnVLA 等**专用模型**的策略不同:专用模型以降低模型耗时为核心目标,可默认/直接使用 TRT 等加速路径;通用模型必须保持 FlashRT 框架能力为默认,TRT 只是显式 opt-in 的**建议加速项**。输出 JSON 记录实际 backend(`eager`/`trt`)。 +- **最新端到端性能(真实图片 `hand_1.jpg`,prompt "Describe the image.",target_size=512,stage-aware benchmark,含 §4.11 融合 kernel + §4.12 FA4 之后)**: + + | 口径 | VQGAN backend | FlashRT FP8 p50/mean | 说明 | + |---|---|--:|---| + | 默认 E2E | eager | **~190 ms** | VQGAN 74.9ms 主导;eager 无 TRT 时瓶颈在 VQGAN | + | 显式 opt-in E2E | TRT | **121.1 / 121.2 ms** | TRT VQGAN 17.5ms + transformer 103.5ms(含 FA4) | + | transformer-prefill-only(FA4) | eager ids reused | **101.9 / 102.0 ms** | HF-comparable,不含 VQGAN,50 iter | + + > **2026-08-05 复测(单热窗口,20 iter,`benchmarks/chameleon_thor_latency.py`)**: + > transformer-only FA4 off **111.2 ms** / FA4 on **104.2 ms**(−7.0);E2E eager+FA4 **177.3 ms**; + > E2E TRT+FA4 **120.2 ms**。与下表历史值差异在热噪声(±5%)内;PR 面文档 + > (`docs/chameleon_usage.md`、`docs/benchmark_comparison.md`、USAGE.md)统一用复测值。 + + Roofline 结论(详见 §4.10-4.12):Se=1056/1072 时理论工作量约 **14.3-14.5 TFLOP**;按 240 TFLOP/s 计,乐观 compute floor 约 **59-60 ms**。per-shape GEMM 微测(§4.11)证实 GEMM tactic 已接近 Thor 实测天花板(32 层 GEMM-only ≈61.9ms),因此与 floor 的差距主要来自非 GEMM 工作。§4.11 融合 RMSNorm/SwiGLU+amax(117.5→110.9ms)、§4.12 引入 FA4 attention(110.9→**101.9ms**,58.3% of 240TFLOP/s,**1.71× floor**)。剩余空间集中在 O-projection 量化(无自然融合点)与 KV-cache 增量解码(后者已在 §4.13 落地)。 +- **FA4 attention(显式 opt-in)**:参考上游 PR [`flashrt-project/FlashRT#163`](https://github.com/flashrt-project/FlashRT/pull/163)(GROOT N1.7 Thor NVFP4+FA4,单图 51.6→29.9ms,1.70×)。Chameleon 形状(Se=1056,32 head,HD=128,causal)实测 FA4 比仓库 CUTLASS causal FMHA **快 2.75×**(450.5→163.9 µs/层),输出 cos=0.99999994;集成后 transformer-only FP8 **-8.4ms**。依赖 `pip install .[thor-fa4]`(nvidia-cutlass-dsl==4.5.1 + quack-kernels==0.4.1),通过 `FLASHRT_CHAMELEON_FA4_ATTN=1` 或构造参数 `use_fa4_attn=True` 开启,backend 不可用时自动回退 CUTLASS FMHA。 +- **KV-cache 增量解码(2026-08 新增,详见 §4.13)**:`generate_greedy` 现为一次 prefill + M=1 增量 decode(`chameleon_decode_step`),稳态 **30.4 tok/s**(32.9 ms/token),墙钟约 **2.8×** 于全前缀重算;逐 token 与 eager 全前缀重算 oracle 完全一致(32-token 生成 38/38)。新增 bottom-right 对齐 causal FMHA 符号 `fmha_fp16_causal_br`(decode 时 SQ=1`=2, override via `eos_token_id`) or `max_seq`. Requires + the dynamic-FP8 path (`use_fp8=True`); the eager full-recompute path is + retained as `_generate_greedy_recompute` for oracle comparisons. + Benchmark via `scripts/bench_chameleon_thor.py --generate-greedy N`. +- The TRT VQGAN path uses a square `target_size×target_size` bicubic + resize while eager uses aspect-preserving `var_center_crop` — token + counts can differ slightly between backends (expected behavior + difference, not a bug). diff --git a/examples/thor/README.md b/examples/thor/README.md index 5e1e2124..1a87b54a 100644 --- a/examples/thor/README.md +++ b/examples/thor/README.md @@ -91,6 +91,25 @@ Quality measurements, the `max_pixels` resolution knob, and the per-projection `wq_overrides` sweep surface are in [`docs/qwen3_vl_thor.md`](../../docs/qwen3_vl_thor.md). +## Chameleon-7B (multimodal chat) + +`chameleon_quickstart.py` runs standalone Chameleon-7B (image + text) with +the dynamic per-tensor FP8 backbone and CUDA-graph prefill: + +```bash +python examples/thor/chameleon_quickstart.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image /path/to/image.jpg \ + --prompt "Describe the image." \ + --benchmark +``` + +Add `--use-trt-vqgan` when compatible TensorRT VQ-GAN engines exist +(`scripts/build_vqgan_trt.py`), and `FLASHRT_CHAMELEON_FA4_ATTN=1` for the +optional FA4 attention fast path. Measured ~190 ms E2E prefill (eager +VQGAN), ~120 ms with TRT VQGAN + FA4, and ~30 tok/s incremental decode. +Full details in [`docs/chameleon_usage.md`](../../docs/chameleon_usage.md). + ## Thor VLA performance ### Precision (Pi0.5, 2-view LIBERO) diff --git a/examples/thor/chameleon_quickstart.py b/examples/thor/chameleon_quickstart.py new file mode 100644 index 00000000..c071f28d --- /dev/null +++ b/examples/thor/chameleon_quickstart.py @@ -0,0 +1,95 @@ +#!/usr/bin/env python +"""Chameleon-7B (Thor sm_110) quickstart. + +Runs the standalone Chameleon-7B image+text frontend on a real image and +reports prefill latency. The frontend is a direct-instantiation class +(``ChameleonTorchFrontendThor``, registered in ``_PIPELINE_MAP`` but not +dispatched by ``flash_rt.load_model`` — same pattern as Qwen3-VL / Nex-N2 +/ LingBot). + +Build first (one shared module — the Chameleon kernels live inside +flash_rt_kernels; FA4 is optional): + + cmake -B build -S . -DGPU_ARCH=110 + cmake --build build -j + pip install -e ".[torch]" # add ,thor-fa4 for the FA4 fast path + +Run: + + python examples/thor/chameleon_quickstart.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image /path/to/hand_1.jpg \ + --prompt "Describe the image." + +Expected on Thor (dynamic FP8, CUDA graph, target_size=512, eager VQGAN): +~190 ms/prefill E2E; with --use-trt-vqgan (engines present) ~120 ms; +with FA4 enabled (FLASHRT_CHAMELEON_FA4_ATTN=1) the transformer part drops +to ~104 ms. The script prints the actual VQGAN backend and FA4 status. +""" +import argparse +import time + +import torch + +from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor +from flash_rt.hardware.thor import fa4_backend + + +def main() -> None: + ap = argparse.ArgumentParser(description="Chameleon-7B Thor quickstart") + ap.add_argument("--checkpoint", required=True, + help="Chameleon-7B dir (model-*-of-*.safetensors + config.json)") + ap.add_argument("--image", required=True, help="real image path (jpg/png)") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when engines are available; default is eager VQGAN)") + ap.add_argument("--use-fp16", action="store_true", + help="Use the FP16 reference path instead of dynamic FP8") + ap.add_argument("--no-graph", action="store_true", help="Disable CUDA Graph") + ap.add_argument("--iters", type=int, default=10, help="timed replays") + ap.add_argument("--benchmark", action="store_true", + help="report wall-clock prefill latency (P50)") + args = ap.parse_args() + + from PIL import Image + + image = Image.open(args.image).convert("RGB") + fa4 = fa4_backend.is_available() + print(f"[chameleon] FA4 attention available: {fa4} ({fa4_backend.status()})" + f"{'' if fa4 else ' <-- CUTLASS FMHA will be used; pip install .[thor-fa4]'}") + + fe = ChameleonTorchFrontendThor( + args.checkpoint, + use_fp8=not args.use_fp16, + use_cuda_graph=not args.no_graph, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + ) + print(f"[chameleon] VQGAN backend: {fe.vqgan_backend} " + f"FA4 active: {fe.fa4_attn_active}") + + out = fe.prefill(args.prompt, [image]) + print(f"[chameleon] Se={out['Se']} (real_len={len([i for i in out['input_ids'] if i != 1])}) " + f"logits={tuple(out['logits'].shape)}") + top = int(torch.argmax(out["logits"]).item()) + print(f"[chameleon] greedy next-token id: {top}") + + if args.benchmark: + ts = [] + for _ in range(args.iters): + torch.cuda.synchronize() + t0 = time.perf_counter() + fe.prefill(args.prompt, [image]) + torch.cuda.synchronize() + ts.append((time.perf_counter() - t0) * 1000.0) + ts.sort() + p50 = ts[len(ts) // 2] + print(f"[chameleon] prefill P50: {p50:.1f} ms over {args.iters} iters " + f"(wall-clock, includes VQGAN; fp8={not args.use_fp16}, " + f"graph={not args.no_graph}, vqgan={fe.vqgan_backend})") + + +if __name__ == "__main__": + main() diff --git a/flash_rt/api.py b/flash_rt/api.py index 13d4e20e..c3ebe5db 100644 --- a/flash_rt/api.py +++ b/flash_rt/api.py @@ -510,6 +510,24 @@ def load_model(checkpoint, framework="torch", num_views=2, autotune=3, "use_fp4_decoder/use_fa4 are unsupported with framework=" "'jetson_pi'; use the Thor torch FP4/FA4 frontend instead") + # Chameleon-7B is a chat-style VLM, not a VLA: its frontends expose + # set_prompt(...) + generate() rather than predict(images, ...), so + # VLAModel's result['actions'] contract does not apply. Registered in + # _PIPELINE_MAP for discoverability only. + if config == "chameleon": + raise NotImplementedError( + "config='chameleon' is a chat-style VLM and is not served through " + "load_model's VLA wrapper. Construct it directly:\n" + " from flash_rt.frontends.torch.chameleon_rtx_sm87 import " + "ChameleonTorchFrontendRtxSm87 # Jetson Orin SM87\n" + " from flash_rt.frontends.torch.chameleon_thor import " + "ChameleonTorchFrontendThor # Jetson Thor SM110\n" + " f = ChameleonTorchFrontendRtxSm87('/path/to/Chameleon_7B_mGPT')\n" + " f.set_prompt('Describe this image.', images=[img])\n" + " print(f.generate(max_new_tokens=32))\n" + "See docs/chameleon7b_rtx_sm87.md, docs/chameleon_thor_sm110.md " + "and docs/chameleon_usage.md.") + if framework == "jetson_pi": if config not in ("pi0", "pi05", "llm", "mllm"): raise ValueError( diff --git a/flash_rt/hardware/__init__.py b/flash_rt/hardware/__init__.py index 5d231ead..c3bd1372 100644 --- a/flash_rt/hardware/__init__.py +++ b/flash_rt/hardware/__init__.py @@ -159,6 +159,24 @@ def detect_arch() -> str: ("flash_rt.frontends.torch.qwen3_vl_rtx_bf16", "Qwen3VlTorchFrontendRtxBF16"), + # ── Chameleon-7B ── + # Direct frontend (set_prompt/generate), not the VLA predict() surface. + # Orin SM87: INT8/INT4+QuaRot-Hadamard path; compute in + # flash_rt/models/chameleon/pipeline_rtx.py. Registered for resolver / + # direct-construction discovery only; load_model(config="chameleon") + # raises a redirect because this exposes set_prompt() + generate(), + # not the VLA predict() surface. See docs/chameleon7b_rtx_sm87.md. + ("chameleon", "torch", "rtx_sm87"): + ("flash_rt.frontends.torch.chameleon_rtx_sm87", + "ChameleonTorchFrontendRtxSm87"), + # Thor SM110: dynamic-FP8 backbone (optional NVFP4 FFN), attention via + # the dedicated Chameleon Thor backend (FA4 -> CUTLASS causal FMHA -> + # cuBLAS fallback). Compute in flash_rt/models/chameleon/pipeline_thor.py. + # See docs/chameleon_usage.md. + ("chameleon", "torch", "thor"): + ("flash_rt.frontends.torch.chameleon_thor", + "ChameleonTorchFrontendThor"), + # Cosmos3-Edge official Thor baseline. ("cosmos3_edge", "torch", "thor"): ("flash_rt.frontends.torch.cosmos3_edge_thor", "Cosmos3EdgeTorchFrontendThor"), @@ -195,6 +213,7 @@ def detect_arch() -> str: # resolving here and crashing later at the first kernel launch. _SM87_ALLOWED = { ("pi05", "torch", "rtx_sm87"), + ("chameleon", "torch", "rtx_sm87"), ("qwen3_vl", "torch", "rtx_sm87"), } diff --git a/scripts/bench_chameleon_thor.py b/scripts/bench_chameleon_thor.py new file mode 100644 index 00000000..d05f119f --- /dev/null +++ b/scripts/bench_chameleon_thor.py @@ -0,0 +1,358 @@ +"""Real-image latency benchmark for standalone Chameleon-7B on Thor. + +Measures HF BF16, FlashRT FP16, and FlashRT dynamic FP8 prefill +latency on the same real-image prompt. Inputs are always real images +(from a user-supplied directory), never synthetic token ids. + +Usage +----- + PYTHONPATH=. python scripts/bench_chameleon_thor.py \\ + --checkpoint /path/to/Chameleon_7B_mGPT \\ + --image-dir /path/to/images \\ + --prompt "Describe the image." \\ + --iters 10 --warmup 2 \\ + --output /tmp/chameleon_thor_bench.json +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import time +from typing import Dict, List + +import numpy as np + + +def _stats(xs: List[float]) -> Dict[str, float]: + a = np.asarray(xs, dtype=np.float64) + return { + "mean": float(a.mean()), + "p50": float(np.percentile(a, 50)), + "min": float(a.min()), + "max": float(a.max()), + } + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths], [str(p) for p in paths] + + +def _pad_ids(input_ids: List[int], pad_id: int = 1) -> tuple[List[int], int]: + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _estimate_prefill_tflops(Se: int) -> float: + D, Dff, L, vocab = 4096, 11008, 32, 65536 + + def gemm_flops(M: int, N: int, K: int) -> int: + return 2 * M * N * K + + per_layer_gemm = ( + gemm_flops(Se, 3 * D, D) + + gemm_flops(Se, D, D) + + gemm_flops(Se, 2 * Dff, D) + + gemm_flops(Se, D, Dff) + ) + per_layer_attn = 4 * Se * Se * D + lm_head = gemm_flops(1, vocab, D) + return (L * (per_layer_gemm + per_layer_attn) + lm_head) / 1e12 + + +def _roofline(Se: int, prefill_ms: float, peak_tflops: float) -> Dict[str, float]: + tflops = _estimate_prefill_tflops(Se) + achieved = tflops / (prefill_ms / 1000.0) if prefill_ms > 0 else 0.0 + floor_ms = tflops / peak_tflops * 1000.0 if peak_tflops > 0 else 0.0 + return { + "estimated_tflops": float(tflops), + "assumed_peak_tflops": float(peak_tflops), + "achieved_tflops": float(achieved), + "efficiency_vs_peak": float(achieved / peak_tflops) if peak_tflops > 0 else 0.0, + "optimistic_compute_floor_ms": float(floor_ms), + "measured_over_floor": float(prefill_ms / floor_ms) if floor_ms > 0 else 0.0, + } + + +def _run_flashrt_prefill_once(fe, prompt: str, images, cached_ids, + *, use_cuda_graph: bool): + import torch + + times = {} + t0 = time.perf_counter() + if cached_ids is None: + ids = fe.encode_prompt(prompt, images) + else: + ids = cached_ids + torch.cuda.synchronize() + t1 = time.perf_counter() + + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + torch.cuda.synchronize() + t2 = time.perf_counter() + + fe._embed_ids(padded) + torch.cuda.synchronize() + t3 = time.perf_counter() + + if use_cuda_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + torch.cuda.synchronize() + t4 = time.perf_counter() + + fe._project_last() + torch.cuda.synchronize() + t5 = time.perf_counter() + + times["encode_ms"] = (t1 - t0) * 1000.0 + times["prepare_ms"] = (t2 - t1) * 1000.0 + times["embed_ms"] = (t3 - t2) * 1000.0 + times["backbone_ms"] = (t4 - t3) * 1000.0 + times["lm_head_ms"] = (t5 - t4) * 1000.0 + times["transformer_prefill_ms"] = times["embed_ms"] + times["backbone_ms"] + times["lm_head_ms"] + times["total_ms"] = (t5 - t0) * 1000.0 + return times, fe.Se, real_len, fe.vqgan_backend + + +def _bench_flashrt(checkpoint_dir: pathlib.Path, prompt: str, images, + *, use_fp8: bool, use_cuda_graph: bool, + target_size: int, use_trt_vqgan: bool, + trt_vqgan_engine_dir: str | None, + iters: int, warmup: int, + reuse_input_ids: bool, + generate_greedy: int, + peak_tflops: float) -> Dict: + import torch + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + fe = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=use_fp8, use_cuda_graph=use_cuda_graph, + target_size=target_size, use_trt_vqgan=use_trt_vqgan, + trt_vqgan_engine_dir=trt_vqgan_engine_dir) + + cached_ids = None + input_build_ms = None + if reuse_input_ids: + t0 = time.perf_counter() + cached_ids = fe.encode_prompt(prompt, images) + torch.cuda.synchronize() + input_build_ms = (time.perf_counter() - t0) * 1000.0 + + for _ in range(warmup): + _run_flashrt_prefill_once( + fe, prompt, images, cached_ids, use_cuda_graph=use_cuda_graph) + + stage_values: Dict[str, List[float]] = {} + Se = real_len = None + backend = fe.vqgan_backend + for _ in range(iters): + times, Se, real_len, backend = _run_flashrt_prefill_once( + fe, prompt, images, cached_ids, use_cuda_graph=use_cuda_graph) + for k, v in times.items(): + stage_values.setdefault(k, []).append(v) + + stage_stats = {k: _stats(v) for k, v in stage_values.items()} + prefill_ms = stage_stats["transformer_prefill_ms"]["p50"] + result = { + "Se": int(Se), + "real_len": int(real_len), + "vqgan_backend": backend, + "fa4_attn": fe.fa4_attn_active, + "reuse_input_ids": bool(reuse_input_ids), + "one_time_input_build_ms": input_build_ms, + "latency_ms": stage_stats["total_ms"], + "stage_breakdown_ms": stage_stats, + "roofline": _roofline(int(Se), prefill_ms, peak_tflops), + } + + if generate_greedy > 0: + for _ in range(max(1, min(warmup, 2))): + fe.generate_greedy(prompt, images, max_new_tokens=generate_greedy) + gen_lat = [] + for _ in range(iters): + t0 = time.perf_counter() + out = fe.generate_greedy(prompt, images, max_new_tokens=generate_greedy) + torch.cuda.synchronize() + gen_lat.append((time.perf_counter() - t0) * 1000.0) + result["generate_greedy"] = { + "max_new_tokens": int(generate_greedy), + "latency_ms": _stats(gen_lat), + "ms_per_token": _stats([x / generate_greedy for x in gen_lat]), + "output_token_count": len(out["input_ids"]), + } + + del fe + torch.cuda.empty_cache() + return result + + +def _bench_hf(checkpoint_dir: pathlib.Path, prompt: str, images, + *, target_size: int, use_trt_vqgan: bool, + trt_vqgan_engine_dir: str | None, + iters: int, warmup: int) -> Dict: + import torch + from transformers import AutoConfig + from safetensors.torch import load_file + + try: + from transformers import ChameleonForConditionalGeneration as _Cls + except (ImportError, ModuleNotFoundError) as e: + print(f"[bench] HF BF16 reference unavailable ({e}); skipping") + return None + + cfg = AutoConfig.from_pretrained(str(checkpoint_dir)) + cfg.rope_scaling = None + if not hasattr(cfg, "rope_theta") or cfg.rope_theta is None: + cfg.rope_theta = 10000.0 + + model = _Cls(cfg) + sd = {} + for shard in sorted(checkpoint_dir.glob("model-*-of-*.safetensors")): + sd.update(load_file(str(shard))) + model.load_state_dict(sd, strict=False, assign=False) + model = model.to(torch.bfloat16).cuda().eval() + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + fe = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=False, use_cuda_graph=False, + target_size=target_size, use_trt_vqgan=use_trt_vqgan, + trt_vqgan_engine_dir=trt_vqgan_engine_dir) + ids = fe.encode_prompt(prompt, images) + backend = fe.vqgan_backend + del fe + torch.cuda.empty_cache() + + ids_t = torch.tensor([ids], dtype=torch.long, device="cuda") + + def _fwd(): + with torch.no_grad(): + model(input_ids=ids_t, use_cache=False) + torch.cuda.synchronize() + + for _ in range(warmup): + _fwd() + lat = [] + for _ in range(iters): + t0 = time.perf_counter() + _fwd() + lat.append((time.perf_counter() - t0) * 1000.0) + del model + torch.cuda.empty_cache() + return {"Se": len(ids), "vqgan_backend": backend, "latency_ms": _stats(lat)} + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when available; default is eager VQGAN)") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--iters", type=int, default=10) + ap.add_argument("--warmup", type=int, default=2) + ap.add_argument("--no-graph", action="store_true", + help="Disable CUDA Graph capture for FlashRT paths") + ap.add_argument("--reuse-input-ids", action="store_true", + help="Build real-image input ids once and benchmark transformer prefill only") + ap.add_argument("--stage-breakdown", action="store_true", + help="Include per-stage timing in JSON output (currently always collected)") + ap.add_argument("--generate-greedy", type=int, default=0, + help="Also benchmark full-prefix greedy generation for N new tokens") + ap.add_argument("--peak-tflops", type=float, default=240.0, + help="Measured Thor FP8 GEMM plateau used for roofline efficiency") + ap.add_argument("--skip-hf", action="store_true") + ap.add_argument("--output", default="/tmp/chameleon_thor_bench.json") + args = ap.parse_args() + + import torch + + checkpoint_dir = pathlib.Path(args.checkpoint) + image_dir = pathlib.Path(args.image_dir) + images, image_paths = _load_real_images(image_dir, args.max_images) + device_name = torch.cuda.get_device_name(0) + + result: Dict = { + "checkpoint": str(checkpoint_dir), + "image_dir": str(image_dir), + "image_paths": image_paths, + "prompt": args.prompt, + "num_images": len(images), + "target_size": args.target_size, + "use_trt_vqgan": bool(args.use_trt_vqgan), + "trt_vqgan_engine_dir": args.trt_vqgan_engine_dir, + "vqgan_backend_requested": "trt" if args.use_trt_vqgan else "eager", + "device": device_name, + "iters": args.iters, + "warmup": args.warmup, + "graph": not args.no_graph, + "reuse_input_ids": bool(args.reuse_input_ids), + "stage_breakdown": bool(args.stage_breakdown), + "generate_greedy": int(args.generate_greedy), + "peak_tflops": float(args.peak_tflops), + } + + print("[bench] FlashRT FP16...") + result["flashrt_fp16"] = _bench_flashrt( + checkpoint_dir, args.prompt, images, use_fp8=False, + use_cuda_graph=not args.no_graph, target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup, + reuse_input_ids=args.reuse_input_ids, + generate_greedy=args.generate_greedy, + peak_tflops=args.peak_tflops) + print(f"[bench] FlashRT FP16: {result['flashrt_fp16']}") + + print("[bench] FlashRT dynamic FP8...") + result["flashrt_fp8"] = _bench_flashrt( + checkpoint_dir, args.prompt, images, use_fp8=True, + use_cuda_graph=not args.no_graph, target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup, + reuse_input_ids=args.reuse_input_ids, + generate_greedy=args.generate_greedy, + peak_tflops=args.peak_tflops) + print(f"[bench] FlashRT FP8: {result['flashrt_fp8']}") + + if not args.skip_hf: + print("[bench] HF BF16 (eager)...") + result["hf_bf16"] = _bench_hf( + checkpoint_dir, args.prompt, images, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + iters=args.iters, warmup=args.warmup) + print(f"[bench] HF BF16: {result['hf_bf16']}") + + with open(args.output, "w") as f: + json.dump(result, f, indent=2) + print(f"[bench] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/build_vqgan_trt.py b/scripts/build_vqgan_trt.py new file mode 100644 index 00000000..405eb14f --- /dev/null +++ b/scripts/build_vqgan_trt.py @@ -0,0 +1,324 @@ +#!/usr/bin/env python3 +"""Build TensorRT FP16 engines for the Chameleon VQ-GAN encoder. + +Exports fixed-shape ONNX per resolution, then compiles TRT engines. +Engines are cached at ~/.flash_rt/trt_engines/vqgan/ (or --output_dir). + +Must be run on the TARGET hardware (engines are not portable across GPUs). + +Usage +----- +python scripts/build_vqgan_trt.py \ + --cfg_path /path/to/chameleon/tokenizer/vqgan.yaml \ + --ckpt_path /path/to/chameleon/tokenizer/vqgan.ckpt \ + --resolutions 384x512 384x384 384x672 512x512 \ + --verify +""" + +import argparse +import hashlib +import json +import os +import platform +import sys +import time +from pathlib import Path + +import numpy as np +import torch +import torch.nn as nn +import tensorrt as trt + +from flash_rt.models.chameleon.vqgan import ImageTokenizer + + +class VQGANEncoderWrapper(nn.Module): + """Image tensor -> codebook indices. + + Input : x of shape (B, 3, H, W), float32 in [-1, 1] + Output : indices of shape (B, H/16, W/16), int64 + """ + + def __init__(self, vq_model: nn.Module): + super().__init__() + self.encoder = vq_model.encoder + self.quant_conv = vq_model.quant_conv + # Re-host the codebook as a self-contained nn.Embedding so this + # wrapper is a fully standard nn.Module (no external parameter sharing). + n_e, e_dim = vq_model.quantize.embedding.weight.shape + self.codebook = nn.Embedding(n_e, e_dim) + with torch.no_grad(): + self.codebook.weight.copy_(vq_model.quantize.embedding.weight) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.encoder(x) # (B, 2*z_channels, H/16, W/16) when double_z=True + h = self.quant_conv(h) # (B, e_dim, H/16, W/16) + + b, c, hh, ww = h.shape + # (B, e_dim, H/16, W/16) -> (B, H/16, W/16, e_dim) -> (B*N, e_dim) + z_flat = h.permute(0, 2, 3, 1).contiguous().view(-1, c) + e = self.codebook.weight # (n_e, e_dim) + + # ||z - e||^2 = ||z||^2 + ||e||^2 - 2 z·e + d = ( + (z_flat * z_flat).sum(dim=1, keepdim=True) + + (e * e).sum(dim=1) + - 2.0 * torch.matmul(z_flat, e.t()) + ) # (B*N, n_e) + idx = torch.argmin(d, dim=1) # (B*N,) + idx = idx.view(b, hh, ww).to(torch.int64) + return idx + + +def build_vqmodel(cfg_path: str, ckpt_path: str, device: torch.device) -> nn.Module: + tokenizer = ImageTokenizer(cfg_path=cfg_path, ckpt_path=ckpt_path, device=device) + vq_model = tokenizer._vq_model.eval() + for p in vq_model.parameters(): + p.requires_grad_(False) + return vq_model + + +def parse_resolution(s: str) -> tuple: + h, w = s.lower().split("x") + return int(h), int(w) + + +def export_onnx(vq_model, height, width, batch, opset, output_path, device): + wrapper = VQGANEncoderWrapper(vq_model).to(device).eval() + dummy = torch.randn(batch, 3, height, width, device=device, dtype=torch.float32) + export_kwargs = dict( + input_names=["image"], + output_names=["indices"], + dynamic_axes=None, + opset_version=opset, + do_constant_folding=True, + ) + # PyTorch 2.5+ defaults to dynamo=True which emits TRT-incompatible + # IR-10/opset-18 nodes. Force legacy exporter on those versions. + # On older PyTorch (< 2.5) the kwarg doesn't exist and isn't needed. + _torch_ver = tuple(int(x) for x in torch.__version__.split(".")[:2]) + if _torch_ver >= (2, 5): + export_kwargs["dynamo"] = False + torch.onnx.export(wrapper, dummy, output_path, **export_kwargs) + print(f" [onnx] exported {output_path} (shape=[{batch},3,{height},{width}])") + return wrapper + + +def build_engine(onnx_path, engine_path, workspace_gb, opt_level, fp16=True): + logger = trt.Logger(trt.Logger.WARNING) + builder = trt.Builder(logger) + + # Choose network creation flags. TRT 11+ supports STRONGLY_TYPED which + # makes the engine honor the ONNX's native dtypes verbatim (so FP16 + # weights stay FP16). Older TRT uses EXPLICIT_BATCH + BuilderFlag.FP16 + # which silently demotes to FP32 for some ops on Ada (we observed this + # on TRT 11.1 — the engine ran ~2× slower in FP32 by default). + use_strong = (fp16 and + hasattr(trt.NetworkDefinitionCreationFlag, "STRONGLY_TYPED")) + + if use_strong: + # Strongly typed mode requires the ONNX itself to be FP16. + # Auto-convert from the FP32 ONNX (idempotent: skips if already FP16). + from onnxconverter_common import float16 + import onnx as _onnx_mod + + onnx_path_obj = Path(onnx_path) + fp16_onnx = onnx_path_obj.with_suffix(".fp16.onnx") + if not fp16_onnx.exists(): + print(f" [onnx-fp16] converting {onnx_path_obj.name} → " + f"{fp16_onnx.name}") + mdl = _onnx_mod.load(str(onnx_path_obj)) + mdl16 = float16.convert_float_to_float16(mdl, keep_io_types=True) + _onnx_mod.save(mdl16, str(fp16_onnx)) + else: + print(f" [onnx-fp16] reusing existing {fp16_onnx.name}") + parse_path = str(fp16_onnx) + network = builder.create_network( + 1 << int(trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED)) + elif hasattr(trt.NetworkDefinitionCreationFlag, "EXPLICIT_BATCH"): + # Legacy path (TRT < 10). + parse_path = str(onnx_path) + network = builder.create_network( + 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) + else: + # Very old TRT — default network. + parse_path = str(onnx_path) + network = builder.create_network(0) + + parser = trt.OnnxParser(network, logger) + with open(parse_path, "rb") as f: + if not parser.parse(f.read()): + for i in range(parser.num_errors): + print(f" [trt] parse error: {parser.get_error(i)}") + raise RuntimeError(f"Failed to parse ONNX: {parse_path}") + + config = builder.create_builder_config() + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, + int(workspace_gb * (1 << 30))) + if fp16 and not use_strong and hasattr(trt.BuilderFlag, "FP16"): + # Old TRT path that needs the explicit FP16 flag. + config.set_flag(trt.BuilderFlag.FP16) + config.builder_optimization_level = opt_level + + t0 = time.time() + serialized = builder.build_serialized_network(network, config) + elapsed = time.time() - t0 + + if serialized is None: + raise RuntimeError(f"TRT engine build failed for {onnx_path}") + + with open(engine_path, "wb") as f: + f.write(serialized) + size_mb = os.path.getsize(engine_path) / (1024 * 1024) + print(f" [trt] built {engine_path} ({size_mb:.1f} MB, {elapsed:.1f}s)") + return engine_path + + +def verify_engine(engine_path, torch_wrapper, height, width, batch, device): + logger = trt.Logger(trt.Logger.WARNING) + runtime = trt.Runtime(logger) + with open(engine_path, "rb") as f: + engine = runtime.deserialize_cuda_engine(f.read()) + if engine is None: + print(" [verify] FAILED: could not deserialize engine") + return False + + context = engine.create_execution_context() + stream = torch.cuda.current_stream() + + h_lat, w_lat = height // 16, width // 16 + inp = torch.randn(batch, 3, height, width, device=device, dtype=torch.float32) + out = torch.zeros(batch, h_lat, w_lat, device=device, dtype=torch.int64) + + # TRT may use int32 for output; detect binding dtype + out_name = "indices" + out_dtype_trt = engine.get_tensor_dtype(out_name) + if out_dtype_trt == trt.DataType.INT32: + out_buf = torch.zeros(batch, h_lat, w_lat, device=device, dtype=torch.int32) + else: + out_buf = out + + context.set_tensor_address("image", inp.data_ptr()) + context.set_tensor_address("indices", out_buf.data_ptr()) + context.execute_async_v3(stream_handle=stream.cuda_stream) + stream.synchronize() + + if out_dtype_trt == trt.DataType.INT32: + trt_indices = out_buf.to(torch.int64) + else: + trt_indices = out_buf + + with torch.no_grad(): + pt_indices = torch_wrapper(inp) + + match = (trt_indices == pt_indices).sum().item() + total = trt_indices.numel() + mismatch_pct = 100.0 * (1.0 - match / total) + ok = mismatch_pct < 0.5 + print(f" [verify] match={match}/{total} ({100*match/total:.2f}%), " + f"mismatch={mismatch_pct:.3f}% {'OK' if ok else 'WARN'}") + return ok + + +def compute_ckpt_hash(ckpt_path: str) -> str: + h = hashlib.sha256() + with open(ckpt_path, "rb") as f: + h.update(f.read(65536)) + return h.hexdigest()[:16] + + +def main(): + parser = argparse.ArgumentParser(description="Build TRT engines for VQ-GAN encoder") + parser.add_argument("--cfg_path", type=str, required=True, + help="Path to the Chameleon VQ-GAN vqgan.yaml config") + parser.add_argument("--ckpt_path", type=str, required=True, + help="Path to the Chameleon VQ-GAN vqgan.ckpt checkpoint") + parser.add_argument("--resolutions", nargs="+", default=["384x512", "384x384", "384x672", "512x512"], + help="HxW resolutions to build engines for") + parser.add_argument("--batch", type=int, default=1) + parser.add_argument("--output_dir", type=str, + default=str(Path.home() / ".flash_rt" / "trt_engines" / "vqgan")) + parser.add_argument("--workspace_gb", type=float, default=2.0) + parser.add_argument("--opt_level", type=int, default=5, + help="TRT builder optimization level (0-5)") + parser.add_argument("--opset", type=int, default=17) + parser.add_argument("--keep_onnx", action="store_true", + help="Keep intermediate ONNX files in output_dir") + parser.add_argument("--verify", action="store_true", + help="Run parity check TRT vs PyTorch after build") + args = parser.parse_args() + + device = torch.device("cuda") + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + print(f"TensorRT {trt.__version__}") + print(f"Platform: {platform.machine()}") + print(f"Output dir: {output_dir}") + print(f"Resolutions: {args.resolutions}") + print() + + vq_model = build_vqmodel(args.cfg_path, args.ckpt_path, device) + ckpt_hash = compute_ckpt_hash(args.ckpt_path) + + manifest_path = output_dir / "manifest.json" + if manifest_path.exists(): + with open(manifest_path) as f: + manifest = json.load(f) + manifest["build_date"] = time.strftime("%Y-%m-%dT%H:%M:%S") + else: + manifest = { + "trt_version": trt.__version__, + "platform": platform.machine(), + "ckpt_hash": ckpt_hash, + "build_date": time.strftime("%Y-%m-%dT%H:%M:%S"), + "batch": args.batch, + "precision": "fp16", + "engines": {}, + } + + for res_str in args.resolutions: + height, width = parse_resolution(res_str) + assert height % 16 == 0 and width % 16 == 0, f"H/W must be multiples of 16, got {res_str}" + h_lat, w_lat = height // 16, width // 16 + + print(f"── {res_str} ({height}×{width} → {h_lat}×{w_lat} latent) ──") + + engine_name = f"vqgan_encoder_b{args.batch}_{height}x{width}_fp16.engine" + engine_path = output_dir / engine_name + + # Export ONNX + onnx_path = output_dir / f"vqgan_encoder_{height}x{width}.onnx" + wrapper = export_onnx(vq_model, height, width, args.batch, args.opset, + str(onnx_path), device) + + # Build TRT engine + build_engine(str(onnx_path), str(engine_path), args.workspace_gb, args.opt_level) + + # Verify + if args.verify: + verify_engine(str(engine_path), wrapper.to(device), height, width, args.batch, device) + + # Clean ONNX + if not args.keep_onnx: + onnx_path.unlink(missing_ok=True) + + manifest["engines"][res_str] = { + "file": engine_name, + "height": height, + "width": width, + "input_shape": [args.batch, 3, height, width], + "output_shape": [args.batch, h_lat, w_lat], + } + print() + + # Write manifest + manifest_path = output_dir / "manifest.json" + with open(manifest_path, "w") as f: + json.dump(manifest, f, indent=2) + print(f"Manifest written: {manifest_path}") + print("Done.") + + +if __name__ == "__main__": + main() diff --git a/scripts/chameleon_orin_check.py b/scripts/chameleon_orin_check.py new file mode 100644 index 00000000..e8886ff5 --- /dev/null +++ b/scripts/chameleon_orin_check.py @@ -0,0 +1,392 @@ +"""Gate-1 correctness harness for Chameleon-7B on Orin SM87. + +Compares the FlashRT INT8/INT4 frontend against a **stock transformers 4.57.1** +``ChameleonForConditionalGeneration`` reference (bf16, eager attention) on the +*same* token ids, and runs the CUDA-Graph safety gate on the decode body. + +Two non-obvious things this handles: + +* **The checkpoint does not load into stock transformers as-is.** Its + ``ChameleonLayerNorm`` builds ``(num_heads, head_dim) = (32,128)`` weights + while this Lumina-mGPT export stores ``(1,128)`` — the shard is + ``model_parallel_size`` x ``head_dim`` and upstream expands it with + ``repeat_interleave`` at forward time. We expand it at load instead, which is + exactly equivalent for the 7B (mp=1) layout. transformers 4.57 also rejects + ``from_pretrained(..., state_dict=...)``, so the model is built with a naked + constructor + ``load_state_dict``. + +* **VQ-GAN index drift would poison every number.** FlashRT runs the encoder + convs in fp16, so codebook indices can differ from an all-fp32 reference. The + FlashRT side therefore *exports* the ids it computed and the reference is fed + those verbatim, isolating LLM error from tokenizer error. Run with + ``--vq-fp16-argmin`` to measure the drift itself instead. + +Usage: + PYTHONPATH=. python scripts/chameleon_orin_check.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image FlashRT.png --prompt "Describe this image." --steps 16 + ... --int4 # QuaRot W4A4 tier + ... --text-only # skip the image (fast smoke) +""" + +from __future__ import annotations + +import argparse +import sys + +import torch + +PROBE_LAYERS = [0, 4, 8, 12, 16, 20, 24, 28, 31] + +# Chameleon suppresses the 8192 image-codebook ids at every forward, so they +# carry no information and must be excluded from any similarity metric — +# including them makes cosine NaN (finfo(bf16).min squared overflows fp32). +IMG_LO, IMG_HI = 4, 8196 + + +def cosine(a: torch.Tensor, b: torch.Tensor) -> float: + a = a.detach().float().flatten() + b = b.detach().float().flatten() + return float(a @ b / (a.norm() * b.norm() + 1e-30)) + + +def text_slice(logits: torch.Tensor) -> torch.Tensor: + """Drop the masked image-id band before comparing logits.""" + return torch.cat([logits[..., :IMG_LO], logits[..., IMG_HI:]], dim=-1) + + +# ====================================================================== +# Reference +# ====================================================================== + +def load_reference(ckpt: str): + """Stock transformers Chameleon, bf16, eager, with qk_norm expanded.""" + import json + from pathlib import Path + from safetensors.torch import load_file + from transformers import ChameleonConfig, ChameleonForConditionalGeneration + + ckpt_p = Path(ckpt) + index = json.loads((ckpt_p / "model.safetensors.index.json").read_text()) + sd, n_exp = {}, 0 + for shard in sorted(set(index["weight_map"].values())): + full = load_file(str(ckpt_p / shard)) + for k, t in full.items(): + if ((".q_norm." in k or ".k_norm." in k) + and t.dim() == 2 and t.shape[0] == 1): + t = t.repeat_interleave(32, dim=0) + n_exp += 1 + sd[k] = t + del full + if n_exp != 128: + print(f" [warn] expanded {n_exp} qk_norm tensors, expected 128") + + cfg = ChameleonConfig.from_pretrained(ckpt) + cfg._attn_implementation = "eager" + torch.set_default_dtype(torch.bfloat16) + model = ChameleonForConditionalGeneration(cfg) + torch.set_default_dtype(torch.float32) + missing, unexpected = model.load_state_dict(sd, strict=False) + missing = [m for m in missing if "inv_freq" not in m] + if missing or unexpected: + raise RuntimeError(f"ref load: missing={missing[:4]} unexpected={unexpected[:4]}") + del sd + return model.eval().cuda() + + +@torch.no_grad() +def reference_forward(model, ids: list): + """Teacher-forced forward. Returns (per-layer hidden states, logits). + + Hidden states come from forward hooks on the decoder layers, NOT from + ``output_hidden_states=True``: HF's ``all_hidden_states`` replaces the last + entry with the *post-final-norm* tensor, so comparing it against a + pre-norm probe shows a false cosine collapse. + """ + caught = {} + handles = [] + for li in PROBE_LAYERS: + def hook(_m, _inp, out, li=li): + caught[li] = (out[0] if isinstance(out, tuple) else out)[0].float().cpu() + return None # a non-None hook return REPLACES the output + handles.append(model.model.layers[li].register_forward_hook(hook)) + + def norm_hook(_m, _inp, out): + caught["final_norm"] = out[0].float().cpu() + return None + handles.append(model.model.norm.register_forward_hook(norm_hook)) + try: + t = torch.tensor([ids], device="cuda") + logits = model(input_ids=t).logits[0].float().cpu() + finally: + for h in handles: + h.remove() + return caught, logits + + +# ====================================================================== +# Gates +# ====================================================================== + +def gate_graph_safety(front) -> bool: + """Capture the decode body and prove it is not frozen (stale-value test).""" + print("\n── graph-safety gate (decode body) ──") + pos = front.S + tok_a, tok_b = 16853, 40000 + try: + front.decode_step(tok_a, pos=pos) # warm every M=1 shape first + front.decode_step(tok_a, pos=pos) + torch.cuda.synchronize() + g = torch.cuda.CUDAGraph() + s = torch.cuda.Stream() + with torch.cuda.stream(s): + with torch.cuda.graph(g, stream=s): + # The kernels MUST be launched on the capture stream; on stream 0 + # they are silently not recorded and the replay looks frozen. + front.decode_step(front._tok_dev, pos=pos, + stream=int(s.cuda_stream)) + print(" capture : OK (no code=13)") + except Exception as e: # pragma: no cover + print(f" capture : FAIL — {type(e).__name__}: {e}") + return False + + front._tok_dev.fill_(tok_a); g.replay(); torch.cuda.synchronize() + la = front._logits.clone() + front._tok_dev.fill_(tok_b); g.replay(); torch.cuda.synchronize() + lb = front._logits.clone() + c = cosine(text_slice(la), text_slice(lb)) + frozen = torch.equal(la, lb) + print(f" stale-value : {'FAIL (frozen)' if frozen else 'PASS'} " + f"(cos between two seed tokens = {c:.4f})") + return not frozen + + +def gate_overflow(front) -> bool: + """FP16 residual health. + + The gate is **finiteness**, not an absolute magnitude. Chameleon's L31 + residual legitimately reaches ~2.6e5 in the bf16 reference (the massive + activation), so any absolute threshold below that would fail by + construction. What must not happen is inf/nan, which the + ``ffn_down_clamp`` prevents by capping the down output just under FP16's + 65504. Saturation at the clamp is therefore *expected* at L31 and is + reported for information only. + """ + print("\n── fp16 residual health (finite + clamp saturation) ──") + snaps = front.snapshot_probe() + if not snaps: + print(" (frontend built without probe_layers — skipped)") + return True + clamp = float(getattr(front, "ffn_down_clamp", 0.0) or 0.0) + bad, sat = [], [] + worst, worst_k = 0.0, "" + for k, v in snaps.items(): + if not bool(torch.isfinite(v).all()): + bad.append(k) + m = float(v[torch.isfinite(v)].abs().max()) if v.numel() else 0.0 + if clamp and m >= clamp * 0.98: + sat.append(k) + if m > worst: + worst, worst_k = m, k + print(f" max finite |x| : {worst:.0f} at {worst_k} " + f"(fp16 max 65504, clamp {clamp:.0f})") + print(f" saturating at clamp: {sat if sat else 'none'} " + f"{'(expected at L31)' if sat else ''}") + ok = not bad + print(f" inf/nan : {bad if bad else 'none'} -> " + f"{'PASS' if ok else 'FAIL'}") + return ok + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image", default=None) + ap.add_argument("--prompt", default="Describe this image.") + ap.add_argument("--steps", type=int, default=16, help="greedy tokens to compare") + ap.add_argument("--max-seq", type=int, default=1280) + ap.add_argument("--int4", action="store_true") + ap.add_argument("--int4-down", action="store_true") + ap.add_argument("--text-only", action="store_true") + ap.add_argument("--vq-fp16-argmin", action="store_true", + help="measure VQ index drift instead of avoiding it") + ap.add_argument("--split-kv-bias", type=int, default=4) + ap.add_argument("--skip-ref", action="store_true", + help="run only the FlashRT-side gates") + args = ap.parse_args() + + from flash_rt.frontends.torch.chameleon_rtx_sm87 import ChameleonTorchFrontendRtxSm87 + + print("=" * 68) + print("Chameleon-7B Orin SM87 — Gate 1") + print("=" * 68) + + front = ChameleonTorchFrontendRtxSm87( + args.checkpoint, max_seq=args.max_seq, + use_int4=args.int4, use_int4_down=args.int4_down, + split_kv_bias=args.split_kv_bias, + vq_argmin_fp32=not args.vq_fp16_argmin, + probe_layers=PROBE_LAYERS) + print(f"tier={front.precision_tier} spec={front.precision_spec()}") + + # ---- prompt ---- + images = None + text = args.prompt + if not args.text_only: + if args.image: + from PIL import Image + images = [Image.open(args.image).convert("RGB")] + else: + import numpy as np + from PIL import Image + images = [Image.fromarray(np.random.RandomState(0).randint( + 0, 256, (480, 640, 3), dtype=np.uint8))] + print(" [note] no --image given; using deterministic noise") + text = "" + args.prompt + front.set_prompt(text, images=images) + ids = front.input_ids.tolist() + print(f"\nISL={len(ids)} (1 BOS + n_img*1026 + text + 1 sep) " + f"images={front.timing['n_images']} " + f"prompt_ms={front.timing['prompt_ms']:.0f}") + + # ---- FlashRT teacher-forced logits + probes ---- + lg_frt = front.prefill(logits_all=True).float().cpu() + probes = front.snapshot_probe() + print(f"prefill_ms={front.timing['prefill_ms']:.0f}") + + ok_overflow = gate_overflow(front) + ok_graph = gate_graph_safety(front) + + # ---- greedy text ---- + front.set_prompt(text, images=images) + frt_ids = front.generate(max_new_tokens=args.steps, return_ids=True) + frt_txt = front.processor.tokenizer.decode(frt_ids, skip_special_tokens=True) + tm = front.timing + print(f"\n── generation ──\n FlashRT ids : {frt_ids}" + f"\n FlashRT text: {frt_txt!r}" + f"\n decode : {tm['decode_ms_per_token']:.2f} ms/token " + f"= {tm['decode_tok_s']:.2f} tok/s (ISL={len(ids)}, OSL={args.steps})") + + if args.skip_ref: + print("\n(--skip-ref: reference comparison not run)") + return 0 if (ok_graph and ok_overflow) else 1 + + # ---- reference ---- + print(f"\n── HF reference (bf16 eager) ──") + ref = load_reference(args.checkpoint) + ref_h, ref_lg = reference_forward(ref, ids) + + print("\n layer cosine norm-ratio FlashRT|max| ref|max|") + worst = 1.0 + for li in PROBE_LAYERS: + a = probes[f"layer_{li}"].float().cpu() + b = ref_h[li] + c = cosine(a, b) + worst = min(worst, c) + print(f" L{li:<5d} {c:.6f} {float(a.norm()/b.norm()):.4f} " + f"{float(a.abs().max()):9.1f} {float(b.abs().max()):9.1f}") + c_fn = cosine(probes["final_norm"].float().cpu(), ref_h["final_norm"]) + print(f" final {c_fn:.6f}") + + # logits + argmax over the whole teacher-forced sequence + a_lg = text_slice(lg_frt) + b_lg = text_slice(ref_lg) + c_last = cosine(a_lg[-1], b_lg[-1]) + am_frt = a_lg.argmax(-1) + am_ref = b_lg.argmax(-1) + exact = am_frt == am_ref + + # A BF16 tie is not a precision failure: when the reference's top-1 and + # top-2 are within one BF16 ULP the winner is numerically arbitrary, and + # any engine may legitimately pick either. Classify those separately + # instead of scoring them as errors. + top2 = b_lg.topk(2, dim=-1).values + gap = (top2[:, 0] - top2[:, 1]).abs() + ulp = top2[:, 0].abs() * 2 ** -8 # BF16 has 8 mantissa bits + tied = gap <= ulp + real_bad = (~exact) & (~tied) + n = len(am_ref) + match_exact = float(exact.float().mean()) + match_adj = float((exact | tied).float().mean()) + print(f"\n last-row logit cosine : {c_last:.6f}") + print(f" argmax exact match : {match_exact*100:.2f}% " + f"({int(exact.sum())}/{n})") + print(f" of {int((~exact).sum())} mismatches: {int(((~exact) & tied).sum())} " + f"are BF16 ties (gap <= 1 ulp), {int(real_bad.sum())} are real") + print(f" tie-adjusted match : {match_adj*100:.2f}%") + if int(real_bad.sum()): + g = gap[real_bad] + print(f" real-mismatch ref gap : median={float(g.median()):.4f} " + f"max={float(g.max()):.4f} (logit scale " + f"~{float(top2[:, 0].abs().median()):.1f})") + match = match_adj + + # Split by position class. At an *image* position the model predicts the + # next token while all 8192 image ids are masked out of the logits, so the + # winner is an arbitrary low-confidence text token — averaging over the 1024 + # image positions swamps the handful that actually drive generation. Gate on + # the text positions only. + ids_t = torch.tensor(ids) + is_img = ((ids_t >= IMG_LO) & (ids_t < IMG_HI)) | (ids_t == 8197) | (ids_t == 8196) + for label, sel in (("image", is_img), ("text ", ~is_img)): + k = int(sel.sum()) + if not k: + continue + e = float(exact[sel].float().mean()) + a = float((exact | tied)[sel].float().mean()) + print(f" {label} positions ({k:4d}) : exact {e*100:6.2f}% " + f"tie-adjusted {a*100:6.2f}% " + f"median ref gap {float(gap[sel].median()):.3f}") + text_sel = ~is_img + if int(text_sel.sum()): + match = float((exact | tied)[text_sel].float().mean()) + + # greedy text identity + with torch.no_grad(): + gen = ref.generate(input_ids=torch.tensor([ids], device="cuda"), + max_new_tokens=args.steps, do_sample=False, + num_beams=1) + ref_new = gen[0, len(ids):].tolist() + ref_txt = front.processor.tokenizer.decode(ref_new, skip_special_tokens=True) + n_pref = 0 + for x, y in zip(frt_ids, ref_new): + if x != y: + break + n_pref += 1 + print(f"\n reference ids : {ref_new}") + print(f" reference text: {ref_txt!r}") + print(f" identical prefix: {n_pref}/{min(len(frt_ids), len(ref_new))} tokens") + + # ---- verdict ---- + cos_gate = 0.99 if front.use_int4 else 0.97 + n_text = int(text_sel.sum()) + checks = [ + ("worst layer cosine", worst >= cos_gate, f"{worst:.4f} >= {cos_gate}"), + ("last-row logit cosine", c_last >= 0.999, f"{c_last:.6f} >= 0.999"), + ("greedy text identical", n_pref == len(ref_new), + f"{n_pref}/{len(ref_new)}"), + ("graph safety", ok_graph, ""), + ("fp16 residual finite", ok_overflow, ""), + ] + # Only binding with a meaningful sample: an image-heavy prompt leaves a + # handful of text positions, where one near-tie flip swings the rate by + # >15 points. Greedy text identity above is the metric that actually + # tracks generation quality. + if n_text >= 32: + checks.insert(2, ("argmax match (text)", match >= 0.99, + f"{match*100:.2f}% >= 99% (n={n_text})")) + else: + print(f"\n note: only {n_text} text positions — the argmax gate is " + f"reported as informational, not binding " + f"({match*100:.2f}% tie-adjusted)") + print("\n" + "=" * 68) + for name, ok, detail in checks: + print(f" [{'PASS' if ok else 'FAIL'}] {name:24s} {detail}") + allok = all(c[1] for c in checks) + print(f"\nGATE 1: {'PASS' if allok else 'FAIL'}") + print("=" * 68) + return 0 if allok else 1 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/check_chameleon_thor_precision.py b/scripts/check_chameleon_thor_precision.py new file mode 100644 index 00000000..31d20c4b --- /dev/null +++ b/scripts/check_chameleon_thor_precision.py @@ -0,0 +1,203 @@ +"""Real-image precision gate for standalone Chameleon-7B on Thor. + +Compares: + * FlashRT FP16 vs HF BF16 last-token logits (cosine, top-k overlap, + greedy next-token equality) + * FlashRT dynamic FP8 vs FlashRT FP16 last-token logits (same metrics) + * optional final-hidden cosine + +Inputs are always real images (from a user-supplied directory), never +synthetic token ids, per the standalone Chameleon-7B optimization plan. + +Usage +----- + PYTHONPATH=. python scripts/check_chameleon_thor_precision.py \\ + --checkpoint /path/to/Chameleon_7B_mGPT \\ + --image-dir /path/to/images \\ + --prompt "Describe the image." \\ + --output /tmp/chameleon_thor_precision.json +""" + +from __future__ import annotations + +import argparse +import json +import pathlib +import sys + +import numpy as np + + +def _load_hf_bf16(checkpoint_dir: pathlib.Path): + """Load the plain HF ChameleonForConditionalGeneration at bf16. + + HF ``from_pretrained`` silently mis-loads q/k norm weights on this + checkpoint's old [1,128] shape; use a naked model + manual + ``load_state_dict`` instead. Requires the model's reference + ``modeling_chameleon`` implementation (optional; ``--skip-hf`` + skips this comparison). + """ + import torch + from transformers import AutoConfig + from transformers import ChameleonForConditionalGeneration as _Cls + from safetensors.torch import load_file + + cfg = AutoConfig.from_pretrained(str(checkpoint_dir)) + cfg.rope_scaling = None + if not hasattr(cfg, "rope_theta") or cfg.rope_theta is None: + cfg.rope_theta = 10000.0 + + model = _Cls(cfg) + sd = {} + for shard in sorted(checkpoint_dir.glob("model-*-of-*.safetensors")): + sd.update(load_file(str(shard))) + missing, unexpected = model.load_state_dict(sd, strict=False, assign=False) + non_vq_missing = [k for k in missing if "vqmodel" not in k] + if non_vq_missing: + print(f"[_load_hf_bf16] WARNING: {len(non_vq_missing)} " + f"non-VQVAE keys missing from ckpt") + if unexpected: + print(f"[_load_hf_bf16] WARNING: {len(unexpected)} unexpected keys") + model = model.to(torch.bfloat16).cuda().eval() + return model + + +def _hf_last_logits(model, input_ids: list[int]) -> np.ndarray: + import torch + ids = torch.tensor([input_ids], dtype=torch.long, device="cuda") + with torch.no_grad(): + out = model(input_ids=ids, use_cache=False) + logits = out.logits[0, -1].float().cpu().numpy() + return logits + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + images = [Image.open(p).convert("RGB") for p in paths] + return images, [str(p) for p in paths] + + +def _cosine(a: np.ndarray, b: np.ndarray) -> float: + a = a.astype(np.float64).ravel() + b = b.astype(np.float64).ravel() + denom = (np.linalg.norm(a) * np.linalg.norm(b)) + if denom == 0: + return 0.0 + return float(np.dot(a, b) / denom) + + +def _topk_overlap(a: np.ndarray, b: np.ndarray, k: int = 10) -> float: + top_a = set(np.argsort(-a)[:k].tolist()) + top_b = set(np.argsort(-b)[:k].tolist()) + return len(top_a & top_b) / float(k) + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true", + help="Use TensorRT VQGAN if compatible engines exist " + "(recommended when available; default is eager VQGAN)") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--topk", type=int, default=10) + ap.add_argument("--output", default="/tmp/chameleon_thor_precision.json") + ap.add_argument("--skip-hf", action="store_true", + help="Skip HF BF16 comparison (FP16 vs FP8 only)") + args = ap.parse_args() + + checkpoint_dir = pathlib.Path(args.checkpoint) + image_dir = pathlib.Path(args.image_dir) + images, image_paths = _load_real_images(image_dir, args.max_images) + print(f"[check] loaded {len(images)} real image(s) from {image_dir}: " + f"{image_paths}") + + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + result: dict = { + "checkpoint": str(checkpoint_dir), + "image_dir": str(image_dir), + "image_paths": image_paths, + "prompt": args.prompt, + "target_size": args.target_size, + "use_trt_vqgan": bool(args.use_trt_vqgan), + "trt_vqgan_engine_dir": args.trt_vqgan_engine_dir, + "vqgan_backend_requested": "trt" if args.use_trt_vqgan else "eager", + } + + print("[check] running FlashRT FP16 reference path...") + fe_fp16 = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=False, use_cuda_graph=False, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir) + out_fp16 = fe_fp16.prefill(args.prompt, images) + logits_fp16 = out_fp16["logits"].numpy().ravel() + hidden_fp16 = out_fp16["hidden"].numpy() + ids_fp16 = out_fp16["input_ids"] + result["vqgan_backend_actual_fp16"] = out_fp16.get("vqgan_backend") + del fe_fp16 + import torch + torch.cuda.empty_cache() + + print("[check] running FlashRT dynamic-FP8 path...") + fe_fp8 = ChameleonTorchFrontendThor( + str(checkpoint_dir), use_fp8=True, use_cuda_graph=False, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir) + out_fp8 = fe_fp8.prefill(args.prompt, images) + logits_fp8 = out_fp8["logits"].numpy().ravel() + hidden_fp8 = out_fp8["hidden"].numpy() + result["vqgan_backend_actual_fp8"] = out_fp8.get("vqgan_backend") + del fe_fp8 + torch.cuda.empty_cache() + + fp8_vs_fp16 = { + "logits_cosine": _cosine(logits_fp8, logits_fp16), + "topk_overlap": _topk_overlap(logits_fp8, logits_fp16, args.topk), + "greedy_token_match": bool( + int(np.argmax(logits_fp8)) == int(np.argmax(logits_fp16))), + "hidden_cosine": _cosine(hidden_fp8, hidden_fp16), + } + result["flashrt_fp8_vs_flashrt_fp16"] = fp8_vs_fp16 + print(f"[check] FlashRT FP8 vs FP16: {fp8_vs_fp16}") + + if not args.skip_hf: + print("[check] loading HF BF16 model (this may take a while)...") + try: + hf_model = _load_hf_bf16(checkpoint_dir) + except (ImportError, ModuleNotFoundError) as e: + print(f"[check] HF BF16 reference unavailable ({e}); " + f"rerun with --skip-hf for the FP16-vs-FP8 check only") + hf_model = None + if hf_model is not None: + logits_hf = _hf_last_logits(hf_model, ids_fp16) + del hf_model + torch.cuda.empty_cache() + + fp16_vs_hf = { + "logits_cosine": _cosine(logits_fp16, logits_hf), + "topk_overlap": _topk_overlap(logits_fp16, logits_hf, args.topk), + "greedy_token_match": bool( + int(np.argmax(logits_fp16)) == int(np.argmax(logits_hf))), + } + result["flashrt_fp16_vs_hf_bf16"] = fp16_vs_hf + print(f"[check] FlashRT FP16 vs HF BF16: {fp16_vs_hf}") + + with open(args.output, "w") as f: + json.dump(result, f, indent=2) + print(f"[check] wrote {args.output}") + + +if __name__ == "__main__": + main() diff --git a/scripts/profile_chameleon_thor.py b/scripts/profile_chameleon_thor.py new file mode 100644 index 00000000..40854ccb --- /dev/null +++ b/scripts/profile_chameleon_thor.py @@ -0,0 +1,115 @@ +"""Nsight Systems profiling helper for standalone Chameleon Thor. + +Use with CUDA profiler capture range, for example: + + nsys profile --force-overwrite=true \ + -o /tmp/chameleon_prefill_only_512 \ + --capture-range=cudaProfilerApi -t cuda,nvtx \ + env PYTHONPATH=. python scripts/profile_chameleon_thor.py \ + --checkpoint /path/to/Chameleon_7B_mGPT \ + --image-dir /path/to/images \ + --target-size 512 --reuse-input-ids --iters 5 +""" + +from __future__ import annotations + +import argparse +import pathlib + + +def _load_real_images(image_dir: pathlib.Path, max_images: int): + from PIL import Image + exts = (".jpg", ".jpeg", ".png", ".bmp") + paths = sorted(p for p in image_dir.iterdir() if p.suffix.lower() in exts) + if not paths: + raise FileNotFoundError(f"No real images found under {image_dir}") + paths = paths[:max_images] + return [Image.open(p).convert("RGB") for p in paths] + + +def _pad_ids(input_ids: list[int], pad_id: int = 1) -> tuple[list[int], int]: + real_len = len(input_ids) + padded = list(input_ids) + rem = len(padded) % 16 + if rem: + padded.extend([pad_id] * (16 - rem)) + return padded, real_len + + +def _run_prefill_body(fe, prompt: str, images, cached_ids, *, use_graph: bool): + import torch + + if cached_ids is None: + ids = fe.encode_prompt(prompt, images) + else: + ids = cached_ids + padded, real_len = _pad_ids(ids) + fe._real_len = real_len + fe.Se = len(padded) + fe._last_input_ids = padded + if fe._use_autotune: + fe._autotune_gemms(fe.Se) + fe._embed_ids(padded) + if use_graph: + fe._capture_graph(fe.Se) + fe._infer_graph.replay() + else: + fe._run_backbone(fe.Se) + fe._project_last() + torch.cuda.synchronize() + + +def main() -> None: + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--image-dir", required=True, + help="Directory of real input images") + ap.add_argument("--prompt", default="Describe the image.") + ap.add_argument("--max-images", type=int, default=1) + ap.add_argument("--target-size", type=int, default=512) + ap.add_argument("--use-trt-vqgan", action="store_true") + ap.add_argument("--trt-vqgan-engine-dir", default=None) + ap.add_argument("--use-fp16", action="store_true") + ap.add_argument("--no-graph", action="store_true") + ap.add_argument("--reuse-input-ids", action="store_true") + ap.add_argument("--warmup", type=int, default=3) + ap.add_argument("--iters", type=int, default=5) + args = ap.parse_args() + + import torch + from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + images = _load_real_images(pathlib.Path(args.image_dir), args.max_images) + fe = ChameleonTorchFrontendThor( + args.checkpoint, + use_fp8=not args.use_fp16, + use_cuda_graph=not args.no_graph, + target_size=args.target_size, + use_trt_vqgan=args.use_trt_vqgan, + trt_vqgan_engine_dir=args.trt_vqgan_engine_dir, + ) + cached_ids = fe.encode_prompt(args.prompt, images) if args.reuse_input_ids else None + + for _ in range(args.warmup): + _run_prefill_body(fe, args.prompt, images, cached_ids, use_graph=not args.no_graph) + + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStart() + for _ in range(args.iters): + _run_prefill_body(fe, args.prompt, images, cached_ids, use_graph=not args.no_graph) + torch.cuda.synchronize() + torch.cuda.cudart().cudaProfilerStop() + + print({ + "Se": fe.Se, + "real_len": fe._real_len, + "vqgan_backend": fe.vqgan_backend, + "use_fp8": not args.use_fp16, + "graph": not args.no_graph, + "reuse_input_ids": args.reuse_input_ids, + "iters": args.iters, + }) + + +if __name__ == "__main__": + main() diff --git a/tests/test_chameleon_thor_vqgan_backend.py b/tests/test_chameleon_thor_vqgan_backend.py new file mode 100644 index 00000000..9562a84a --- /dev/null +++ b/tests/test_chameleon_thor_vqgan_backend.py @@ -0,0 +1,18 @@ +from __future__ import annotations + +import inspect +import os + +from flash_rt.frontends.torch.chameleon_thor import ChameleonTorchFrontendThor + + +def test_chameleon_trt_vqgan_is_opt_in_by_default(): + sig = inspect.signature(ChameleonTorchFrontendThor.__init__) + assert sig.parameters["use_trt_vqgan"].default is False + + +def test_chameleon_fa4_attn_is_opt_in_by_default(): + sig = inspect.signature(ChameleonTorchFrontendThor.__init__) + assert sig.parameters["use_fa4_attn"].default is None + os.environ.pop("FLASHRT_CHAMELEON_FA4_ATTN", None) + assert bool(os.environ.get("FLASHRT_CHAMELEON_FA4_ATTN", "0") in ("1", "true", "on")) is False From f1bba6ecba436dfca055a015b17759db2373a2fd Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:10:08 +0000 Subject: [PATCH 4/7] fix(chameleon): rename leftover fork env var and clean fork terminology - FLASHRT_RYNNVLA2_FP4_LAYERS -> FLASHRT_CHAMELEON_FP4_LAYERS (the env var was inherited from the RynnVLA port with its old name) - replace "001"/"002"/"vendor bf16" comments with plain Chameleon / HF-reference wording in pipeline_thor.py and chameleon_thor.py --- flash_rt/frontends/torch/chameleon_thor.py | 4 ++-- flash_rt/models/chameleon/pipeline_thor.py | 14 +++++++------- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/flash_rt/frontends/torch/chameleon_thor.py b/flash_rt/frontends/torch/chameleon_thor.py index 2ff9ec9a..cc74116f 100644 --- a/flash_rt/frontends/torch/chameleon_thor.py +++ b/flash_rt/frontends/torch/chameleon_thor.py @@ -203,7 +203,7 @@ def _validate_config(self) -> None: def _build_image_token_mask(self) -> None: """Image-codebook vocab ids to suppress for text generation. - Mirrors vendor ``ChameleonForConditionalGeneration``'s + Mirrors HF ``ChameleonForConditionalGeneration``'s ``mask_image_logits``: without this, greedy decode on a text prompt can emit VQGAN codebook ids (garbage BPE decode) because those ids are heavily represented in training and the raw @@ -436,7 +436,7 @@ def _ensure_trt_vqgan_loaded(self) -> bool: return available def _preprocess_image_for_trt(self, pil_image, out_hw): - """PIL uint8 -> CUDA [1,3,H,W] float32 [-1,1] (vendor-matched).""" + """PIL uint8 -> CUDA [1,3,H,W] float32 [-1,1] (matches the HF reference preprocessing).""" H_out, W_out = out_hw if pil_image.size != (W_out, H_out): pil_image = pil_image.resize((W_out, H_out), resample=PIL.Image.BICUBIC) diff --git a/flash_rt/models/chameleon/pipeline_thor.py b/flash_rt/models/chameleon/pipeline_thor.py index 0741b91c..c76d0d38 100644 --- a/flash_rt/models/chameleon/pipeline_thor.py +++ b/flash_rt/models/chameleon/pipeline_thor.py @@ -34,10 +34,10 @@ def _parse_fp4_layer_policy() -> frozenset: - """FP4 FFN layer policy from FLASHRT_RYNNVLA2_FP4_LAYERS env var. + """FP4 FFN layer policy from FLASHRT_CHAMELEON_FP4_LAYERS env var. Values: - - unset / "": default = L0-L2 FP4, L3-L31 FP8 (safe, matches 001). + - unset / "": default = L0-L2 FP4, L3-L31 FP8 (safe default). - "0-7": FP4 for L0..L7 inclusive. - "0-14,20-31": FP4 for L0..L14 + L20..L31 (skip outlier L15-L19). This is the SM120-sweep-validated aggressive setting (13ms savings on @@ -47,7 +47,7 @@ def _parse_fp4_layer_policy() -> frozenset: Returns the frozenset of FP8 layer indices (complement of FP4 set). """ import os as _os - val = _os.environ.get("FLASHRT_RYNNVLA2_FP4_LAYERS", "").strip() + val = _os.environ.get("FLASHRT_CHAMELEON_FP4_LAYERS", "").strip() if not val: return frozenset(range(3, 32)) # default: L0-L2 FP4 @@ -340,7 +340,7 @@ def chameleon_forward( # ═══ End dynamic per-tensor FP8 branch ═══ # ── QKV FP8 GEMMs ── - # 002 has no attention bias; pass zero-buffer for fp8_nn_bias epilogue + # Chameleon has no attention bias; pass zero-buffer for fp8_nn_bias epilogue if alpha_host is not None: alpha_qkv = float(alpha_host[li * 4 + 0]) zero_bias_ptr = int(bufs['zero_bias_d']) @@ -746,7 +746,7 @@ def chameleon_forward_fp16( Ported from pipeline_rtx.chameleon_forward. All 32 layers run pure FP16 GEMMs via ``gemm.fp16_nn`` — same on Thor as RTX. - Precision-optimal path (cosine target ≥ 0.99 vs vendor bf16) at the + Precision-optimal path (cosine target ≥ 0.99 vs HF bf16) at the cost of ~2× the FP8 path latency in the LLM. Recommended when downstream ActionHead is sensitive to accumulated FP8 error. @@ -809,7 +809,7 @@ def chameleon_forward_fp16( Se, D, 1e-5, int(stream), ) - # Q / K / V GEMMs (no bias in 002). + # Q / K / V GEMMs (no bias in Chameleon). # NOTE: Q_ptr aliases xn_ptr on Thor (chameleon slots["Q_O"] = # bufs['xn'].data_ptr()). Because gemm.fp16_nn reads A (xn) and # writes D (Q) at the SAME fp16 dtype and SAME buffer, cuBLAS @@ -979,7 +979,7 @@ def chameleon_forward_calibrate( Se * D, int(stream), ) - # ── 3. Q/K/V FP8 GEMMs (no bias for 002) ── + # ── 3. Q/K/V FP8 GEMMs (no bias in Chameleon) ── q_w_ptr = int(weights['q_w'][li]) k_w_ptr = int(weights['k_w'][li]) v_w_ptr = int(weights['v_w'][li]) From b340a1bd61f0ecfc1a674df224677755d4468b4d Mon Sep 17 00:00:00 2001 From: DXICM <185532351+DXICM@users.noreply.github.com> Date: Thu, 6 Aug 2026 13:53:03 +0800 Subject: [PATCH 5/7] fix(chameleon): address upstream review feedback - Add #!/usr/bin/env python3 shebangs to 4 scripts and 1 benchmark - Rename _chameleon_spec.py to _chameleon_rtx_sm87_spec.py (hardware suffix per adding_new_model.md convention) and update the import - Add config="chameleon" to docs/stable_api.md (config enum, redirect bullet, resolve_pipeline_class registration) - Translate docs/chameleon_thor_sm110.md from Chinese to English - Remove all internal "derivative repo" / "RynnVLA" provenance references from both engineering docs (42 occurrences) - Move fp4_chameleon_layer16 benchmark from tests/ to benchmarks/ (it has no test_ functions; misfiled in the kernels branch) --- benchmarks/chameleon_thor_latency.py | 1 + benchmarks/fp4_chameleon_layer16.py | 177 ++++++++ docs/chameleon7b_rtx_sm87.md | 121 ++---- docs/chameleon_thor_sm110.md | 377 +++++++++--------- docs/stable_api.md | 11 +- ...on_spec.py => _chameleon_rtx_sm87_spec.py} | 0 .../frontends/torch/chameleon_rtx_sm87.py | 2 +- scripts/bench_chameleon_thor.py | 1 + scripts/chameleon_orin_check.py | 1 + scripts/check_chameleon_thor_precision.py | 1 + scripts/profile_chameleon_thor.py | 1 + 11 files changed, 424 insertions(+), 269 deletions(-) create mode 100644 benchmarks/fp4_chameleon_layer16.py rename flash_rt/frontends/torch/{_chameleon_spec.py => _chameleon_rtx_sm87_spec.py} (100%) diff --git a/benchmarks/chameleon_thor_latency.py b/benchmarks/chameleon_thor_latency.py index 04b95506..c72521c8 100644 --- a/benchmarks/chameleon_thor_latency.py +++ b/benchmarks/chameleon_thor_latency.py @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 """Chameleon-7B (Thor sm_110) latency benchmark. Measures standalone Chameleon-7B prefill latency on real images with clean diff --git a/benchmarks/fp4_chameleon_layer16.py b/benchmarks/fp4_chameleon_layer16.py new file mode 100644 index 00000000..54bc0745 --- /dev/null +++ b/benchmarks/fp4_chameleon_layer16.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python3 +"""Verify FP4 Gate+Up substitution at Chameleon layer-16 FFN shape. + +Compares three paths driven by identical fp16 weights/activation: + + REF : pure fp16 matmul + silu*mul + matmul (fp32 accumulate) + FP8 : full FP8 path (used by current chameleon_forward) + MIX : FP4 Gate+Up + (existing) silu_mul_split_fp8_fp16 + FP8 Down + ALL4: FP4 Gate+Up + fp16 silu*mul + FP4 Down (upper bound) + +For each path: cosine similarity vs REF + microbenchmark latency. +""" +import pytest + +torch = pytest.importorskip("torch") +fp4 = pytest.importorskip( + "flash_rt.flash_rt_fp4", + reason="flash_rt_fp4 requires an NVFP4 (sm_120+) build") +import numpy as np +import flash_rt.flash_rt_kernels as fvk +from flash_rt.executors.fp4_utils import ( + quant_weight_nvfp4, FP4ActScratch, quant_act_nvfp4, fp4_gemm, pick_variant, +) + + +def fp16_t(*shape, scale=1.0): + return (torch.randn(*shape, dtype=torch.float16, device='cuda') * scale).contiguous() + + +def cuda_time(fn, iters=100, warmup=20): + s = torch.cuda.current_stream() + for _ in range(warmup): fn() + s.synchronize() + e0 = torch.cuda.Event(enable_timing=True); e1 = torch.cuda.Event(enable_timing=True) + e0.record() + for _ in range(iters): fn() + e1.record(); s.synchronize() + return e0.elapsed_time(e1) / iters * 1000 # μs + + +def amax_scale(t: torch.Tensor) -> float: + return max(t.abs().max().item() / 448.0, 1e-9) + + +def make_scale_buf(scale: float) -> torch.Tensor: + return torch.tensor([scale], dtype=torch.float32, device='cuda') + + +def quant_fp8(W: torch.Tensor, scale: float): + out = torch.empty_like(W, dtype=torch.uint8) + sb = make_scale_buf(scale) + fvk.quantize_fp8_static_fp16(W.data_ptr(), out.data_ptr(), + sb.data_ptr(), W.numel(), 0) + return out, sb + + +def cos_vs(a, b): + return torch.nn.functional.cosine_similarity( + a.flatten().float().unsqueeze(0), + b.flatten().float().unsqueeze(0)).item() + + +def main(): + print(f"FP4 enabled: {fp4.has_nvfp4()}; variants: {fp4.cutlass_fp4_gemm_num_variants()}") + + Se, D, Dff = 1216, 4096, 11008 + + torch.manual_seed(0) + W_g = fp16_t(Dff, D, scale=0.02) + W_u = fp16_t(Dff, D, scale=0.02) + W_d = fp16_t(D, Dff, scale=0.02) + X = fp16_t(Se, D, scale=1.0) + + # ---- REF ---- + gate_ref = (X.float() @ W_g.float().T).half() + up_ref = (X.float() @ W_u.float().T).half() + h_ref = (torch.nn.functional.silu(gate_ref.float()) * up_ref.float()).half() + out_ref = (h_ref.float() @ W_d.float().T).half() + print(f"REF: |gate|max={gate_ref.abs().max():.2f} |up|max={up_ref.abs().max():.2f}" + f" |h|max={h_ref.abs().max():.2f} |out|max={out_ref.abs().max():.2f}") + + # ---- Pre-compute calibrated scales (per-tensor amax/448) ---- + s_x = amax_scale(X) + s_wg = amax_scale(W_g) + s_wu = amax_scale(W_u) + s_wd = amax_scale(W_d) + s_h = amax_scale(h_ref) # post-silu*up → fp8 input to Down + print(f"scales: x={s_x:.3e} w_g={s_wg:.3e} w_u={s_wu:.3e} w_d={s_wd:.3e} h={s_h:.3e}") + + gemm = fvk.GemmRunner() + + # FP8 weights + activation + # NB: fp8_nn_dev is NN (no transpose), so B must be [K, N] row-major. + # We store HF-style W as [N, K]; transpose before fp8 quant. + Wg_fp8, sg = quant_fp8(W_g.t().contiguous(), s_wg) # [D, Dff] + Wu_fp8, su = quant_fp8(W_u.t().contiguous(), s_wu) # [D, Dff] + Wd_fp8, sd = quant_fp8(W_d.t().contiguous(), s_wd) # [Dff, D] + sx_buf = make_scale_buf(s_x); sh_buf = make_scale_buf(s_h) + X_fp8 = torch.empty(Se, D, dtype=torch.uint8, device='cuda') + fvk.quantize_fp8_static_fp16(X.data_ptr(), X_fp8.data_ptr(), + sx_buf.data_ptr(), Se*D, 0) + + gate_out = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + up_out = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + gu_fp8 = torch.empty(Se, Dff, dtype=torch.uint8, device='cuda') + out_fp8 = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_fp8(): + gemm.fp8_nn_dev(X_fp8.data_ptr(), Wg_fp8.data_ptr(), gate_out.data_ptr(), + Se, Dff, D, sx_buf.data_ptr(), sg.data_ptr(), 0) + gemm.fp8_nn_dev(X_fp8.data_ptr(), Wu_fp8.data_ptr(), up_out.data_ptr(), + Se, Dff, D, sx_buf.data_ptr(), su.data_ptr(), 0) + fvk.silu_mul_split_fp8_fp16(gate_out.data_ptr(), up_out.data_ptr(), + gu_fp8.data_ptr(), Se*Dff, + sh_buf.data_ptr(), 0) + gemm.fp8_nn_dev(gu_fp8.data_ptr(), Wd_fp8.data_ptr(), out_fp8.data_ptr(), + Se, D, Dff, sh_buf.data_ptr(), sd.data_ptr(), 0) + + run_fp8(); torch.cuda.synchronize() + cos_fp8 = cos_vs(out_fp8, out_ref) + fp8_us = cuda_time(run_fp8) + + # ---- MIX (FP4 Gate+Up, FP8 Down) ---- + qg = quant_weight_nvfp4(W_g) + qu = quant_weight_nvfp4(W_u) + sc_x = FP4ActScratch(max_M=Se, K=D) + var_gu = pick_variant(Dff, D) + out_mix = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_mix(): + quant_act_nvfp4(X, sc_x, Se, stream=0) + fp4_gemm(sc_x, qg, gate_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fp4_gemm(sc_x, qu, up_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fvk.silu_mul_split_fp8_fp16(gate_out.data_ptr(), up_out.data_ptr(), + gu_fp8.data_ptr(), Se*Dff, + sh_buf.data_ptr(), 0) + gemm.fp8_nn_dev(gu_fp8.data_ptr(), Wd_fp8.data_ptr(), out_mix.data_ptr(), + Se, D, Dff, sh_buf.data_ptr(), sd.data_ptr(), 0) + + run_mix(); torch.cuda.synchronize() + cos_mix = cos_vs(out_mix, out_ref) + mix_us = cuda_time(run_mix) + + # ---- ALL-FP4 (Gate+Up+Down all FP4, fp16 silu*mul) ---- + qd = quant_weight_nvfp4(W_d) + sc_h = FP4ActScratch(max_M=Se, K=Dff) + var_dn = pick_variant(D, Dff) + h_buf = torch.empty(Se, Dff, dtype=torch.float16, device='cuda') + out_all4 = torch.empty(Se, D, dtype=torch.float16, device='cuda') + + def run_all4(): + quant_act_nvfp4(X, sc_x, Se, stream=0) + fp4_gemm(sc_x, qg, gate_out, Se, Dff, D, variant_idx=var_gu, stream=0) + fp4_gemm(sc_x, qu, up_out, Se, Dff, D, variant_idx=var_gu, stream=0) + # fp16 silu*up via torch (bench-only) + torch.mul(torch.nn.functional.silu(gate_out), up_out, out=h_buf) + quant_act_nvfp4(h_buf, sc_h, Se, stream=0) + fp4_gemm(sc_h, qd, out_all4, Se, D, Dff, variant_idx=var_dn, stream=0) + + run_all4(); torch.cuda.synchronize() + cos_all4 = cos_vs(out_all4, out_ref) + all4_us = cuda_time(run_all4) + + print() + print("="*72) + print("Chameleon layer-16 FFN block (Se=1216, D=4096, Dff=11008)") + print("="*72) + fmt = " {:14s} cos_vs_ref = {:.6f} {:7.1f} μs speedup={}" + print(fmt.format("FP8 baseline", cos_fp8, fp8_us, "1.00x")) + print(fmt.format("MIX (FP4 GU)", cos_mix, mix_us, f"{fp8_us/mix_us:.2f}x")) + print(fmt.format("ALL-FP4", cos_all4, all4_us, f"{fp8_us/all4_us:.2f}x")) + delta = fp8_us - mix_us + print(f"\n Per-layer MIX saves {delta:6.1f} μs → 32 layers ≈ {delta*32/1000:5.2f} ms") + + +if __name__ == '__main__': + main() diff --git a/docs/chameleon7b_rtx_sm87.md b/docs/chameleon7b_rtx_sm87.md index 2c369f16..1f0b77c1 100644 --- a/docs/chameleon7b_rtx_sm87.md +++ b/docs/chameleon7b_rtx_sm87.md @@ -10,11 +10,7 @@ > overflow (§4.6). > > This is the authoritative document for **upstream Chameleon-7B as an image+text -> → text VLM** on Orin SM87. It is a knowledge migration from the RynnVLA-001 -> port in the derivative repo this was migrated from (same Chameleon-7B -> backbone, but *prefill-only* + action head); the cross-model mechanism -> analysis and the decode-regime methodology on the same platform that it built -> on are documented in that repo. +> → text VLM** on Orin SM87. --- @@ -48,25 +44,6 @@ granularity, smoothing, or a per-layer precision fallback — and because the rotation preserves per-row scales it reuses the stock CUTLASS INT8 GEMMs, so it costs nothing (§4.5). -### What this port had to add over RynnVLA-001 - -RynnVLA-001 shares this exact backbone but is **prefill-only** -(the RynnVLA pipeline in the derivative repo sets `layer_stride_llm = 0`, so all -32 layers share one K/V scratch — no KV cache, lm_head loaded but idle, no -sampler). So: - -* **Prefill / TTFT** — RynnVLA's whole Orin tuning ladder (L2 swizzle Id4, - stages-5, t256x128, vec8 elementwise) lives in the kernels and transferred for - free; measured 91 % of ceiling with zero new tuning. -* **Decode (M=1)** — a different regime (weight-bandwidth-bound, not - FLOPs-bound) with different levers. This was the new work, and it still needed - **zero new CUDA kernels** for the QK-Norm / RoPE / KV-write path (§3.1): the - K/V GEMMs write straight into the cache slab and the RoPE position rides in - the cos/sin table pointer. - -The one kernel this port *did* add is the INT8 twin of the existing FHT pack -(§4.5) — a device-side function inside `csrc/kernels/fht_int4.cu`, no new GEMM. - ## 1. Platform | Field | Value | @@ -87,9 +64,8 @@ picking the wrong one produces >100 % "efficiency" nonsense: | single-stream vectorized reduce (`roofline.py --measure-bw`) | 99 GB/s | nothing — undersaturated | | **best achieved by a real weight-streaming kernel** (int8 gate GEMM @ M=1) | **173.3 GB/s** = 85 % of spec | **the decode roofline denominator** | -The RynnVLA-001 documentation in the derivative repo quotes 166 GB/s for a D2D -copy; we measure 124.5 in this container. Treat 173 GB/s (kernel-achieved, -read-dominated) as the decode ceiling. +A D2D copy measures 124.5 GB/s in this container. Treat 173 GB/s +(kernel-achieved, read-dominated) as the decode ceiling. > ⚠️ `/sys/devices/gpu.0/devfreq/*/cur_freq` is **not readable in this container**, so clocks > cannot be locked or even observed. Every number below is warm (≥30 warmup iters) and a median @@ -108,8 +84,8 @@ num_splits = fa2_num_splits_heuristic_causal(batch*num_heads_q*num_m_blocks, num ``` Chameleon decode: `1*32*1 = 32` vs `0.8 * (16*2) = 25.6` → **`num_splits = 1`**. Passing the -accumulators does nothing. RynnBrain's +12 % split-KV win (documented in the -derivative repo) worked *only* because Qwen3-VL-2B has **16** Q heads. +accumulators does nothing. A split-KV win works *only* when the model has few +enough Q heads (e.g. **16**) for the heuristic to engage splitting. `num_sms` is a pure heuristic knob in this wrapper, so biasing it selects the split count. Measured (q=1, kv=1040, 32 Q heads, head_dim 128, fp16): @@ -194,9 +170,9 @@ match on a deterministic 512×512 input, versus a full-fp32 reference: | **fp16 convs + fp32 distance/argmin** | 99.02 % | The fp32-argmin fix helps (`z²+e²−2ez` is cancellation-prone in fp16; -`modeling_chameleon.py:850-861`) and costs <0.1 ms. This probe used random noise; the -~92 % divergence recorded by the derivative repo's layerwise precision tests was on **real -images**, so re-measure on real content at Phase 5 before declaring the fix sufficient. +`modeling_chameleon.py:850-861`) and costs <0.1 ms. This probe used random noise; +divergence on **real images** can be much higher (~92 % has been observed), so +re-measure on real content at Phase 5 before declaring the fix sufficient. ### 2.5 R8 — decode-graph primitives are graph-safe: **PASS** @@ -235,9 +211,9 @@ Therefore: * **decode** — point them at `+ pos*4096*2` and call the same kernel with `seq_len=1` and cos/sin pre-offset by `pos*128*2`. -Also required: **`Se` must not be even-padded** (the RynnVLA pipeline in the derivative repo -pads `Se` for FP8 GEMM alignment) — with a real KV cache the pad row is junk that decode *will* -attend to, and CUTLASS constrains only `K`. +Also required: **`Se` must not be even-padded** (e.g. for FP8 GEMM alignment) — +with a real KV cache the pad row is junk that decode *will* attend to, and +CUTLASS constrains only `K`. Attention correctness: FA2 causal is **bottom-right aligned** (`fa2_wrapper_causal.cu:126-138`), so `q=1, kv=N` attends all N keys. The cuBLAS fallback @@ -257,24 +233,23 @@ rather than degrade to it. Decode always uses **dynamic per-row** activation quant — never the prefill static calibration, which was fitted at M=Se and does not describe a single decode row. -### 3.3 Token contract (upstream Chameleon — do NOT reuse RynnVLA's ids) +### 3.3 Token contract `[BOS 0] + n_img × ([8197 ] + [8711 ]×1024 + [8196 ]) + text + [8710 sep]`, so `S = 1 + n_img*1026 + n_text + 1`. Image token id = **VQ codebook index + 4**, exactly, for all 8192 codes; the 1024 tokens are a raster scan of the 32×32 latent grid. -> ⚠️ **Trap:** the RynnVLA port in the derivative repo hardcodes -> `: 8710, : 8720`. Both are **wrong** for upstream Chameleon (8710 is the -> `sep_token`). Its `_init_special_token_ids` also sets ids 65536-65539, which are out of range -> for `vocab_size=65536`. This is one of three reasons the Chameleon frontend is standalone -> rather than a subclass of the RynnVLA frontend — see §4. +> ⚠️ **Trap:** do not hardcode `: 8710, : 8720`. Both are +> **wrong** for upstream Chameleon (8710 is the `sep_token`). Likewise, special +> ids 65536-65539 are out of range for `vocab_size=65536`. This is one of three +> reasons the Chameleon frontend is standalone rather than a subclass — see §4. > > ⚠️ `config.json` says `bos_token_id: 1`, which is **stale** (``); `tokenizer.json` gives > ` = 0` and that is what the processor emits. ### 3.4 Why the frontend is standalone (not a subclass) -Three of the most attractive inheritable helpers from the derivative repo's RynnVLA port are +Three of the most attractive inheritable helpers from a VLA-style frontend are *actively wrong* for upstream Chameleon: its `_preprocess_image` is bicubic/384/`x*2-1` where Chameleon needs PIL **LANCZOS**/512/`u8*0.0078-1.0` → `[-1, +0.989]`; its `_vqgan_encode` emits a grid+newline token layout instead of a bare 1024 raster; its `_load_tokenizer` / @@ -335,10 +310,10 @@ Total GPU 281.1 ms (measured before the clamp restriction in §4.2.3): ### 4.2.2 The real GEMM ceiling is 64.4 TOPS, not 84.8 — GEMM tuning is spent -The 84.8 TOPS quoted by the RynnVLA-001 documentation in the derivative repo is -the **raw `mma.s8` issue rate** from a register-only probe. What CUTLASS actually -achieves on its best-case shape is lower: big-square probes measure **58.7 TOPS -at 4096^3 and 64.4 TOPS at 8192^3**. Against that realistic ceiling: +The often-quoted 84.8 TOPS figure is the **raw `mma.s8` issue rate** from a +register-only probe. What CUTLASS actually achieves on its best-case shape is +lower: big-square probes measure **58.7 TOPS at 4096^3 and 64.4 TOPS at +8192^3**. Against that realistic ceiling: | shape | ms/call | TOPS | vs 64.4 ceiling | |---|---|---|---| @@ -350,18 +325,17 @@ at 4096^3 and 64.4 TOPS at 8192^3**. Against that realistic ceiling: The prefill GEMMs are at **91 % of the achievable CUTLASS ceiling**, and the isolated probe reproduces the in-pipeline time to within 0.5 % (0.567 vs 0.568 ms on Q/K/V/O) — so there is no pipeline overhead left to recover. This -independently confirms RynnVLA-001's conclusion that the GEMM ladder (swizzle -Id4 / stages-5 / t256x128) is spent. Only `gate/up` at 85 % shows slack, and -RynnVLA already swept tiles there and measured 256x128 as "only ~2 % better, not -worth a 4th instantiation". +confirms that the GEMM ladder (swizzle Id4 / stages-5 / t256x128) is spent. +Only `gate/up` at 85 % shows slack, and tile sweeps there measured 256x128 as +"only ~2 % better, not worth a 4th instantiation". > WARNING: **the roofline probe itself had a DVFS bug**, found here. Whichever > shape was measured *first* was penalised by clock ramp: Q/K/V/O reported > **28.3 TOPS** measured first versus **61.1** for the identical shape after > adding a 3-second saturating pre-ramp, and the per-shape TOPS ascended purely > in measurement order (28.3 -> 53.8 -> 60.5). `_ramp_clocks()` now runs before -> any timing in the roofline script (in the derivative repo, -> `scripts/bench/orin_int8_roofline.py`). Any earlier per-shape number from that +> any timing in the roofline script +> (`scripts/bench/orin_int8_roofline.py`). Any earlier per-shape number from that > script is suspect. ### 4.2.3 Clamp restricted to the last 4 layers: -6.3 ms @@ -444,10 +418,11 @@ Measured on the same prompt, all three tiers: | 1032 | INT4+down | 0.9334 | 0.9313 | — | — | 0.998881 | **0/16** | | **1032** | **INT8+Hadamard** | **0.9989** | **0.9989** | **0.99972** | **0.99945** | **0.999968** | **16/16** | -> This **contradicts the RynnVLA-001 Orin conclusion** ("both INT4 tiers beat -> INT8 at every layer probe on every frame"). That doc is not wrong — it -> measured a *prefill-only VLA at fixed Se*; the verdict is ISL-dependent, and -> a VLM's production ISL sits in the opposite regime. +> This may appear to **contradict an earlier prefill-only conclusion** ("both +> INT4 tiers beat INT8 at every layer probe on every frame"). That measurement +> is not wrong — it was taken on a *prefill-only workload at fixed short Se*; +> the verdict is ISL-dependent, and a VLM's production ISL sits in the opposite +> regime. So the INT8-vs-INT4 verdict *inverts with sequence length* — at short ISL the sink row is 1/7 of the tensor and rotation dominates; at long ISL it is 1/1032 @@ -478,7 +453,7 @@ hand-written ceiling on 16-SM Orin measured just 41 TOPS. A per-layer FP16 fallback would also have worked, but it is checkpoint-specific tuning that permanently costs throughput — the rotation is free and generalizes. -### 4.6 SOLVED — the FP16 overflow, via RynnVLA-002's `ffn_down_clamp` +### 4.6 SOLVED — the FP16 overflow, via `ffn_down_clamp` With a real image the reference's L31 residual reaches **max|x| = 89088**, above FP16's 65504, so FlashRT stored `inf` and the final RMSNorm turned that row's @@ -486,15 +461,7 @@ logits into `nan`. It affects **both** precision tiers — it is a property of t residual *dtype*, not of the quantization. The first instinct (a BF16 residual stream, ~1 new kernel) was **wrong** — the -answer already existed in the lineage. The derivative repo's Chameleon -acceleration-methodology documentation records this exact failure for the -Chameleon backbone, and its FP8 optimization playbook had already flagged the -missing clamp as a *latent, unverified* risk for RynnVLA-001: - -> **可复用到 001**：001 当前**无 clamp**，是潜在的 inf 风险点（尤其长序列）。直接移植 -> `clamp_inplace_fp16` 即可。 - -This port empirically confirmed that prediction. +answer is a clamp. This port empirically confirmed that a clamp is sufficient. **Why a clamp is sufficient** — measured per-layer magnitudes in the bf16 reference (ISL=1032). The explosion is confined to **exactly one layer**: @@ -511,9 +478,9 @@ Because the pre-L31 residual is only ~2032, clamping the down **output** at (already in `flash_rt_kernels`, CUDA-Graph safe) removes the overflow with **zero new kernels and no dtype change**. -Unlike RynnVLA-002's Thor path we do **not** need to clamp the down *input*: -ours is BF16 (`cutlass_int8_silu_gated_bf16out`), whose range absorbs 151552 -without issue. The clamp is applied on every layer because L0-L30 are three +We do **not** need to clamp the down *input*: ours is BF16 +(`cutlass_int8_silu_gated_bf16out`), whose range absorbs 151552 without issue. +The clamp is applied on every layer because L0-L30 are three orders of magnitude below it and therefore untouched; cost is 32 extra elementwise launches (<0.3 % of the decode budget, ~0.7 % of prefill). @@ -527,14 +494,13 @@ elementwise launches (<0.3 % of the decode budget, ~0.7 % of prefill). | last-row logit cosine | 0.999916 | 0.999916 | | greedy prefix vs HF | 8/16 | 8/16 | -Exposed as `ffn_down_clamp` (default 60000, env `FLASHRT_CHAMELEON_DOWN_CLAMP`), -named after the corresponding RynnVLA-002 env var in the derivative repo. +Exposed as `ffn_down_clamp` (default 60000, env `FLASHRT_CHAMELEON_DOWN_CLAMP`). ⚠️ **The clamp did not change the text divergence** (still 8/16). That confirms the overflow was confined to the sink row's post-L31 residual, which feeds only that row's final norm — so it was never the cause of the divergence. The remaining gap is ordinary INT8 error at high-confidence text decisions and is -still open; see §6 for the ranked options inherited from the RynnVLA lineage. +still open; see §6 for the ranked options. Also note the **gate itself was wrong** at first: an absolute "max|x| < 30000" threshold fails by construction on a backbone whose reference @@ -584,7 +550,7 @@ FHT norm absorb part of it. A prediction that is close *and* slightly pessimisti is the sign the bottleneck model is right (a large gap in either direction would mean the model of the bottleneck is wrong, not that there is tuning left). -**Prefill**: predicted ~255 ms by scaling RynnVLA's measured Se=1214 numbers to +**Prefill**: predicted ~255 ms by scaling a measured Se=1214 prefill to Se=1032; **measured 273.8 ms warm** (within 7 %), of which GEMM is 229.2 ms at **91 % of the achievable CUTLASS ceiling** (§4.2.2). Image tokenize adds ~53 ms (PyTorch VQ-GAN; ~27 ms with a 512x512 TRT engine, not built). @@ -602,7 +568,7 @@ Se=1032; **measured 273.8 ms warm** (within 7 %), of which GEMM is 229.2 ms at | # | lever | outcome | |---|---|---| | 1 | **W8A8 + Hadamard (QuaRot at 8 bits)** | **DONE — this closed it.** greedy 8/16 → **16/16**, worst layer 0.9946 → 0.9986, last-row logit 0.999916 → 0.999968, at no throughput cost. Default tier. | -| 2 | `ffn_down_clamp` (ported from RynnVLA-002) | **DONE** — removed the L31 FP16 `inf` (§4.6) | +| 2 | `ffn_down_clamp` | **DONE** — removed the L31 FP16 `inf` (§4.6) | | ~~3~~ | ~~Tier-3 FP16 fallback for L31~~ | **not needed** — the rotation fixed the same layer at 8 bits. A per-layer precision fallback is checkpoint-specific tuning and costs throughput permanently; prefer the quantization method. | | ~~4~~ | ~~AWQ / SmoothQuant per-K smoothing~~ | **not needed** — and principle #17's measured ladder on this backbone puts smoothing (0.641) far below rotation (0.9914). Kept only as a fallback if a future checkpoint defeats rotation. | | 5 | ISL-adaptive tier selection | **obsolete** — W8A8+Hadamard wins at both short and long ISL, so there is nothing to switch between | @@ -619,7 +585,7 @@ Se=1032; **measured 273.8 ms warm** (within 7 %), of which GEMM is 229.2 ms at | 5 | INT8 Q/K/V/O at 67 % of ceiling (small-N tail: 4096/128 = 32 tiles on 16 SMs) | up to +12 % if it reached 100 % | high (hand GEMV) | open | | 6 | Devpos kernel + fp16 seqused-splitkv FA2 → *one* decode graph | 0 % throughput; removes capture cost + `max_new_tokens` cap | high (1 `.cu` + FA2 rebuild) | deferred | | 7 | Reduced lm_head (drop rows 4..8195) | +0.5–1 % | low | deferred | -| 8 | INT8 KV cache | +3.8 % @1040, **+12 % @4096** | high (needs a 32Q/32KV variant; RynnBrain measured break-even) | S≥4096 only | +| 8 | INT8 KV cache | +3.8 % @1040, **+12 % @4096** | high (needs a 32Q/32KV variant; break-even measured) | S≥4096 only | | ~~9~~ | ~~M=1 up-projection split~~ | ~~+7 %~~ → **measured 0.9 %, net ≈0** | — | **DEAD (§2.2)** | | 10 | Speculative decode | 1.5–2× | N/A — no draft model | — | @@ -665,8 +631,7 @@ PYTHONPATH=. python3 scripts/chameleon_orin_check.py \ # ... --vq-fp16-argmin measure VQ index drift instead of avoiding it # M=1 decode roofline (no checkpoint) — the "do we need a GEMV?" gate -# (the roofline probe script lives in the derivative repo this was migrated -# from: scripts/bench/orin_int8_roofline.py) +# (roofline probe script: scripts/bench/orin_int8_roofline.py) python3 scripts/bench/orin_int8_roofline.py --decode # large-M prefill roofline at the production shape @@ -731,7 +696,7 @@ entry yet (out of scope this round), so these are the frontend kwargs. |---|---| | `precision_tier` / `precision_spec()` / `get_model_info()` | report the resolved configuration; `timing` carries `prompt_ms` / `prefill_ms` / `decode_tok_s` | -**Not accepted** (unlike the RynnVLA frontends): `use_fp8`, `use_fp4`, +**Not accepted**: `use_fp8`, `use_fp4`, `use_fp8_attn`, `use_awq_v_proj`, `num_views`, `action_dim`, `action_chunk_size`, `state_dim` — SM87 has no FP8/FP4 tensor cores and this is not a VLA. Unknown kwargs are swallowed by `**_ignored`. diff --git a/docs/chameleon_thor_sm110.md b/docs/chameleon_thor_sm110.md index af665b6b..dfd25de5 100644 --- a/docs/chameleon_thor_sm110.md +++ b/docs/chameleon_thor_sm110.md @@ -1,99 +1,98 @@ -# 标准 Chameleon-7B @ Thor SM110 — 权威文档 +# Chameleon-7B on Thor SM110 -**平台**: Jetson AGX Thor (SM110, aarch64) · CUDA 13.0 · transformers 4.43+ -**模型**: 标准/独立 Chameleon-7B(纯 LLM 主干 + VQGAN 图像 tokenizer,**无 ActionHead / ActionVAE**) -**生产方案**: 全 32 层**运行时动态 per-tensor FP8**(实现位于 Chameleon 专用的 `flash_rt/models/chameleon/pipeline_thor.py::chameleon_forward`,迁移自 derivative repo 的 RynnVLA-002 Thor 移植)+ 通用 eager Chameleon VQGAN 默认路径 + cuBLASLt 逐 shape autotune + L31 selective clamp;TensorRT VQGAN 仅显式 opt-in -**版本**: v1.4 (2026-08,新增 KV-cache 增量解码 `generate_greedy`:30.4 tok/s,逐 token 与全前缀重算 oracle 一致) - -> 与 RynnVLA-002 Thor 的关系:本文档描述的是**独立/标准 Chameleon-7B**(纯文本+图像对话骨干,直接 frontend `encode_prompt`/`prefill`/`generate_greedy`,不是 VLA `predict()` 接口),没有 ActionHead/ActionVAE。其动态 FP8 Chameleon 主干实现与 Thor attention backend 迁移自 derivative repo 中 RynnVLA-002 的移植。 +**Platform**: Jetson AGX Thor (SM110, aarch64) · CUDA 13.0 · transformers 4.43+ +**Model**: Standalone Chameleon-7B (LLM backbone + VQGAN image tokenizer, **no ActionHead / ActionVAE**) +**Production path**: All 32 layers with **runtime dynamic per-tensor FP8** (implemented in the Chameleon-specific `flash_rt/models/chameleon/pipeline_thor.py::chameleon_forward`) + generic eager Chameleon VQGAN default + cuBLASLt per-shape autotune + L31 selective clamp; TensorRT VQGAN is explicit opt-in only +**Version**: v1.4 (2026-08, added KV-cache incremental decode `generate_greedy`: 30.4 tok/s, token-exact vs full-prefix recompute oracle) --- -## 0. 结论先行 +## 0. Summary -- **资产路径**:`/path/to/Chameleon_7B_mGPT`(注意实际目录名是 `mGPT` 不是 `mGP`)。包含权重 shards、tokenizer、`original_tokenizers/vqgan.{yaml,ckpt}`。 -- **HF 直接加载会失败**:当前 `transformers` 的 `ChameleonForConditionalGeneration.from_pretrained` 在该 checkpoint 上因 `q_norm`/`k_norm` 形状为旧式 `[1,128]`(而不是新式 `[32,128]`)而报错/静默错载。**变通方案**:生产路径直接走 FlashRT 自己的声明式 `WeightLoader`(完全绕开 HF `from_pretrained`);或如 `scripts/check_chameleon_thor_precision.py` 那样,用裸 `ChameleonForConditionalGeneration` 构造器 + `load_state_dict(strict=False)` 加载 HF 参考模型(上游没有 vendor 目录,脚本直接从 `transformers` 导入 `ChameleonForConditionalGeneration`,可用 `--skip-hf` 跳过 HF 参考比对)。 -- **精度已验证(真实图片,非合成 token id)**: - - FlashRT FP16 vs HF BF16(last-token logits cosine,mask_image_logits 后):**0.9999997**,greedy next-token 完全一致。 - - FlashRT 动态 FP8 vs FlashRT FP16:**0.99999999**,greedy next-token 完全一致,top-10 overlap 1.0。 -- **VQGAN backend policy(框架定位)**:FlashRT 面向**通用/标准 Chameleon**时保持框架通用性——VQGAN **默认走 eager** Chameleon tokenization(`use_trt_vqgan=False`),不默认依赖 RynnVLA/TensorRT engine,保证框架自身能力可独立运行。**若部署环境存在可用 TRT engine,建议显式开启**(`use_trt_vqgan=True` 或脚本 `--use-trt-vqgan`;实测 VQGAN 74.9→17.3ms,TRT E2E ~121ms vs eager ~190ms)。这与此前 RynnVLA 等**专用模型**的策略不同:专用模型以降低模型耗时为核心目标,可默认/直接使用 TRT 等加速路径;通用模型必须保持 FlashRT 框架能力为默认,TRT 只是显式 opt-in 的**建议加速项**。输出 JSON 记录实际 backend(`eager`/`trt`)。 -- **最新端到端性能(真实图片 `hand_1.jpg`,prompt "Describe the image.",target_size=512,stage-aware benchmark,含 §4.11 融合 kernel + §4.12 FA4 之后)**: +- **Asset path**: `/path/to/Chameleon_7B_mGPT` (note the actual directory name is `mGPT`, not `mGP`). Contains weight shards, tokenizer, and `original_tokenizers/vqgan.{yaml,ckpt}`. +- **HF direct loading fails**: The current `transformers` `ChameleonForConditionalGeneration.from_pretrained` errors or silently misloads on this checkpoint because `q_norm`/`k_norm` shapes are legacy `[1,128]` (not the newer `[32,128]`). **Workaround**: The production path uses FlashRT's own declarative `WeightLoader` (bypassing HF `from_pretrained` entirely); alternatively, as in `scripts/check_chameleon_thor_precision.py`, use the bare `ChameleonForConditionalGeneration` constructor + `load_state_dict(strict=False)` for the HF reference model (the script imports `ChameleonForConditionalGeneration` directly from `transformers`; use `--skip-hf` to skip the HF reference comparison). +- **Precision validated (real images, not synthetic token ids)**: + - FlashRT FP16 vs HF BF16 (last-token logits cosine, after mask_image_logits): **0.9999997**, greedy next-token exact match. + - FlashRT dynamic FP8 vs FlashRT FP16: **0.99999999**, greedy next-token exact match, top-10 overlap 1.0. +- **VQGAN backend policy (framework positioning)**: For **generic/standard Chameleon**, FlashRT preserves framework generality — VQGAN **defaults to eager** Chameleon tokenization (`use_trt_vqgan=False`), with no default dependency on TensorRT engines, ensuring the framework's own capabilities run independently. **If the deployment environment has compatible TRT engines, explicitly opt in** (`use_trt_vqgan=True` or script `--use-trt-vqgan`; measured VQGAN 74.9→17.3 ms, TRT E2E ~121 ms vs eager ~190 ms). Output JSON records the actual backend (`eager`/`trt`). +- **Latest end-to-end performance (real image `hand_1.jpg`, prompt "Describe the image.", target_size=512, stage-aware benchmark, including §4.11 fused kernels + §4.12 FA4)**: - | 口径 | VQGAN backend | FlashRT FP8 p50/mean | 说明 | + | Scope | VQGAN backend | FlashRT FP8 p50/mean | Notes | |---|---|--:|---| - | 默认 E2E | eager | **~190 ms** | VQGAN 74.9ms 主导;eager 无 TRT 时瓶颈在 VQGAN | - | 显式 opt-in E2E | TRT | **121.1 / 121.2 ms** | TRT VQGAN 17.5ms + transformer 103.5ms(含 FA4) | - | transformer-prefill-only(FA4) | eager ids reused | **101.9 / 102.0 ms** | HF-comparable,不含 VQGAN,50 iter | - - > **2026-08-05 复测(单热窗口,20 iter,`benchmarks/chameleon_thor_latency.py`)**: - > transformer-only FA4 off **111.2 ms** / FA4 on **104.2 ms**(−7.0);E2E eager+FA4 **177.3 ms**; - > E2E TRT+FA4 **120.2 ms**。与下表历史值差异在热噪声(±5%)内;PR 面文档 - > (`docs/chameleon_usage.md`、`docs/benchmark_comparison.md`、USAGE.md)统一用复测值。 - - Roofline 结论(详见 §4.10-4.12):Se=1056/1072 时理论工作量约 **14.3-14.5 TFLOP**;按 240 TFLOP/s 计,乐观 compute floor 约 **59-60 ms**。per-shape GEMM 微测(§4.11)证实 GEMM tactic 已接近 Thor 实测天花板(32 层 GEMM-only ≈61.9ms),因此与 floor 的差距主要来自非 GEMM 工作。§4.11 融合 RMSNorm/SwiGLU+amax(117.5→110.9ms)、§4.12 引入 FA4 attention(110.9→**101.9ms**,58.3% of 240TFLOP/s,**1.71× floor**)。剩余空间集中在 O-projection 量化(无自然融合点)与 KV-cache 增量解码(后者已在 §4.13 落地)。 -- **FA4 attention(显式 opt-in)**:参考上游 PR [`flashrt-project/FlashRT#163`](https://github.com/flashrt-project/FlashRT/pull/163)(GROOT N1.7 Thor NVFP4+FA4,单图 51.6→29.9ms,1.70×)。Chameleon 形状(Se=1056,32 head,HD=128,causal)实测 FA4 比仓库 CUTLASS causal FMHA **快 2.75×**(450.5→163.9 µs/层),输出 cos=0.99999994;集成后 transformer-only FP8 **-8.4ms**。依赖 `pip install .[thor-fa4]`(nvidia-cutlass-dsl==4.5.1 + quack-kernels==0.4.1),通过 `FLASHRT_CHAMELEON_FA4_ATTN=1` 或构造参数 `use_fa4_attn=True` 开启,backend 不可用时自动回退 CUTLASS FMHA。 -- **KV-cache 增量解码(2026-08 新增,详见 §4.13)**:`generate_greedy` 现为一次 prefill + M=1 增量 decode(`chameleon_decode_step`),稳态 **30.4 tok/s**(32.9 ms/token),墙钟约 **2.8×** 于全前缀重算;逐 token 与 eager 全前缀重算 oracle 完全一致(32-token 生成 38/38)。新增 bottom-right 对齐 causal FMHA 符号 `fmha_fp16_causal_br`(decode 时 SQ=1