import contextlib import hashlib import os import pathlib import time import traceback from dataclasses import dataclass import torch from loguru import logger os.environ.setdefault("TT_METAL_PINNED_MEMORY_CACHE_LIMIT_BYTES", "0") import ttnn from models.autoports.cloudflare_clef.tt import precision_defaults precision_defaults.apply() from models.autoports.cloudflare_clef.tt import encode as clef_encode from models.autoports.cloudflare_clef.tt import head as clef_head from models.autoports.cloudflare_clef.tt import vision as clef_vision from models.autoports.cloudflare_clef.tt.loader import ClefModelArgs from models.demos.blackhole.qwen36.tt import mlp as qwen36_mlp from models.demos.blackhole.qwen36.tt import tp_common as tpc from models.demos.blackhole.qwen36.tt.model import Qwen36Model from models.tt_transformers.tt.common import Mode BLOCK_SIZE = 64 ALIGN = 128 MAX_BUCKET = 1024 BUCKETS = tuple(b for b in Qwen36Model._PREFILL_MASK_BUCKETS if b <= MAX_BUCKET) TILE = ttnn.TILE_LAYOUT RM = ttnn.ROW_MAJOR_LAYOUT LM_HEAD_STUB_ROWS_PER_DEVICE = 32 CACHE_RELOAD_MODES = ("off", "host_to_device", "stock") VISION_MAX_SEQ_LEN = 4096 MEDIA_KEYS = ("pixel_values", "image_grid_thw", "pixel_values_videos", "video_grid_thw") SPARE_BLOCKS = MAX_BUCKET // BLOCK_SIZE REPLAY_COST_S = {128: 0.173, 256: 0.187, 512: 0.220, 1024: 0.338} EAGER_COST_S = {128: 0.172, 256: 0.191, 512: 0.226, 1024: 0.345} TRACE_GUARD_BYTES = 2 << 30 REFERENCE_GRIDS = "1,16,20;1,22,38;1,26,36;1,28,36;1,28,38;1,40,50;2,16,20" WARM_GRID_HELP = ( "CLEF_TRACED=1 serves only image grids compiled before the traces were captured: set " "CLEF_VISION_WARM_GRID to a ;-separated list of t,h,w grids (for the 8 reference images and the video: " f"{REFERENCE_GRIDS}), or run eager with CLEF_TRACED=0 (the default), which accepts any grid" ) def single_grid_message(n_grids): return ( f"this request carries {n_grids} image or video grids; a traced server (CLEF_TRACED=1) takes one grid " "per request, because the tower joins several images with a concat program compiled per image count " "and no vision program may compile while traces are live: send one image (or one video) per request, " "or serve with CLEF_TRACED=0 (the default), which accepts several" ) def parse_grids(text): return [tuple(int(g) for g in part.split(",")) for part in text.split(";") if part.strip()] def aligned_pieces(length): out = [] for b in sorted(BUCKETS, reverse=True): while length >= b: out.append(b) length -= b assert length == 0, f"length {length + sum(out)} is not a multiple of {ALIGN}" return out @dataclass class VisionRequest: tokens: object grid: torch.Tensor video: bool n_rows: int n_patches: int seconds: float rows_host: torch.Tensor = None def release(self): if self.tokens is not None: ttnn.deallocate(self.tokens) self.tokens = None self.rows_host = None @property def marker(self): return self.tokens if self.tokens is not None else self.rows_host @dataclass class StateHandle: S: int S0: int suffix_ids: torch.Tensor slot: int key: str = None state_ids: torch.Tensor = None vision: VisionRequest = None def media_key(media): if not media: return "" digest = hashlib.sha1() for name in MEDIA_KEYS: value = media.get(name) if value is None: continue tensor = torch.as_tensor(value).contiguous() digest.update(name.encode()) digest.update(str(tuple(tensor.shape)).encode()) digest.update(tensor.numpy().tobytes()) return digest.hexdigest() def load_media_paths(record): from PIL import Image def load(item): if isinstance(item, (str, os.PathLike)): return Image.open(item).convert("RGB") return item request = dict(record) if request.get("images"): request["images"] = [load(item) for item in request["images"]] if request.get("videos"): request["videos"] = [[load(frame) for frame in frames] for frames in request["videos"]] return request @contextlib.contextmanager def as_tensor_cache_policy(mode): if mode not in CACHE_RELOAD_MODES: raise ValueError(f"cache_reload={mode!r}; expected one of {CACHE_RELOAD_MODES}") if mode == "stock": yield return original = ttnn.as_tensor def patched( tensor, dtype=None, *, layout=ttnn.ROW_MAJOR_LAYOUT, device=None, memory_config=None, cache_file_name=None, preprocess=None, mesh_mapper=None, ): if mode == "off": cache_file_name = None elif cache_file_name is not None and device is not None: dtype_name = dtype.name if dtype is not None else "None" layout_name = layout.name if layout is not None else "None" full = pathlib.Path(f"{cache_file_name}_dtype_{dtype_name}_layout_{layout_name}.tensorbin") if full.is_file(): host = ttnn.load_tensor(full) return ttnn.to_device(host, device, memory_config=memory_config) return original( tensor, dtype, layout=layout, device=device, memory_config=memory_config, cache_file_name=cache_file_name, preprocess=preprocess, mesh_mapper=mesh_mapper, ) ttnn.as_tensor = patched try: yield finally: ttnn.as_tensor = original @contextlib.contextmanager def prefill_only_mlp_weights(enabled): if not enabled: yield return original = qwen36_mlp.load_mlp_weights def load_prefill_only(mesh_device, state_dict, tensor_cache_path=None, args=None, use_gateup_agmm=True): tp = getattr(args, "num_devices", 1) if args is not None else 1 interleaved = tp > 1 and getattr(args, "mlp_1d_decode", False) if not (interleaved and use_gateup_agmm and tpc.mlp_gateup_agmm_enabled(tp)): return original(mesh_device, state_dict, tensor_cache_path, args=args, use_gateup_agmm=use_gateup_agmm) def cache(name, tag=""): return str(tensor_cache_path / f"mlp.{name}.weight{tag}.tp") if tensor_cache_path else None gate_up = qwen36_mlp._build_gate_up( state_dict["gate_proj.weight"], state_dict["up_proj.weight"], mesh_device, tp, cache("gate_up", ".swiglu") ) down = tpc.shard_w( state_dict["down_proj.weight"], mesh_device, dim=0, memory_config=ttnn.DRAM_MEMORY_CONFIG, cache_path=cache("down_proj"), dtype=qwen36_mlp.MLP_DOWN_DTYPE, ) return qwen36_mlp.MLPWeights(w1=None, w2=down, w3=None, w_gate_up=gate_up) qwen36_mlp.load_mlp_weights = load_prefill_only try: yield finally: qwen36_mlp.load_mlp_weights = original PARENTS = {"1x4": ("FABRIC_1D", (1, 4)), "2x2": ("FABRIC_2D", (2, 2))} @contextlib.contextmanager def tp2_mesh(parent="1x4", trace_region_size=0, l1_small_size=24576): fabric_name, shape = PARENTS[parent] ttnn.set_fabric_config(getattr(ttnn.FabricConfig, fabric_name)) parent_mesh = ttnn.open_mesh_device( mesh_shape=ttnn.MeshShape(*shape), l1_small_size=l1_small_size, num_command_queues=2, trace_region_size=trace_region_size, ) sub = parent_mesh.create_submesh(ttnn.MeshShape(1, 2), offset=ttnn.MeshCoordinate(0, 0)) sub.enable_program_cache() logger.info( f"tp2_mesh: {fabric_name} parent {list(parent_mesh.shape)} chips {list(parent_mesh.get_device_ids())}, " f"submesh {list(sub.shape)} chips {list(sub.get_device_ids())}" ) try: yield sub finally: for child in parent_mesh.get_submeshes(): ttnn.close_mesh_device(child) ttnn.close_mesh_device(parent_mesh) ttnn.set_fabric_config(ttnn.FabricConfig.DISABLED) @contextlib.contextmanager def tp_proj_dtype(dtype): if dtype is None or dtype == ttnn.bfloat8_b: yield return original = tpc.shard_w def shard_w(torch_tensor, mesh, dim, memory_config, cache_path, dtype=ttnn.bfloat8_b, _forced=dtype): chosen = _forced if dtype == ttnn.bfloat8_b else dtype return original(torch_tensor, mesh, dim, memory_config, cache_path, dtype=chosen) tpc.shard_w = shard_w try: yield finally: tpc.shard_w = original def mesh_dram_free_bytes(mesh): view = ttnn.get_memory_view(mesh, ttnn.BufferType.DRAM) return int(view.total_bytes_free_per_bank) * int(view.num_banks) class ClefEngine: def __init__( self, mesh, args_cls=ClefModelArgs, max_state_len=16384, chunk_size=1024, snapshot_slots=4, n_layers=None, max_tail_len=4096, cache_reload=None, kv_reserve_bytes=2 << 30, prefill_only_mlp=True, vision=None, traced=None, trace_guard_bytes=None, vision_warm_grid=None, planner=None, ): assert chunk_size == MAX_BUCKET, f"chunk_size must be {MAX_BUCKET} at TP=2 (stage 0 finding 2)" self.mesh = mesh self._rep = ttnn.ReplicateTensorToMesh(mesh) self.num_devices = mesh.get_num_devices() self.chunk_size = chunk_size self.max_state_len = max_state_len self.max_tail_len = max_tail_len self.max_len = max_state_len + max_tail_len self.cache_reload = cache_reload or os.environ.get("CLEF_CACHE_RELOAD", "off") self.prefill_only_mlp = prefill_only_mlp self.precision = precision_defaults.active() self.traced = (os.environ.get("CLEF_TRACED", "0") == "1") if traced is None else bool(traced) self.planner = (os.environ.get("CLEF_PLANNER", "1") == "1") if planner is None else bool(planner) if self.traced: os.environ.setdefault("QWEN_GDN_CONV", "fir") self.gdn_conv_impl = os.environ.get("QWEN_GDN_CONV", "kda") if self.traced else "fir" self.cost_table = REPLAY_COST_S if self.traced else EAGER_COST_S self.trace_guard_bytes = ( int(os.environ.get("CLEF_TRACE_GUARD_BYTES", str(TRACE_GUARD_BYTES))) if trace_guard_bytes is None else int(trace_guard_bytes) ) self.vision_warm_grid = vision_warm_grid self.vision_warmed_grids = [] self.traces = {} self.trace_bytes = 0 self.bufs = {} self.counters = {"replays": 0, "host_writes": 0, "readbacks": 0, "tower_runs": 0} self.timings = {} t0 = time.perf_counter() args = args_cls(mesh_device=mesh, max_batch_size=1, max_seq_len=self.max_len) if n_layers is not None: args.n_layers = n_layers args.attention_type_list = args.attention_type_list[:n_layers] self.args = args self.cache_dir = args.weight_cache_path() logger.info( f"ClefEngine args={type(args).__name__} n_layers={args.n_layers} max_len={self.max_len} " f"cache={self.cache_dir} cache_reload={self.cache_reload} prefill_only_mlp={prefill_only_mlp} " f"precision={self.precision}" ) state_dict = args.load_state_dict() self.timings["host_state_dict_s"] = time.perf_counter() - t0 state_dict["output.weight"] = torch.zeros( self.num_devices * LM_HEAD_STUB_ROWS_PER_DEVICE, args.dim, dtype=torch.bfloat16 ) t1 = time.perf_counter() from models.demos.blackhole.qwen36.tt.precision import PROJ_DTYPE with as_tensor_cache_policy(self.cache_reload), prefill_only_mlp_weights(prefill_only_mlp), tp_proj_dtype( PROJ_DTYPE ): self.model = Qwen36Model(mesh, args, state_dict, tensor_cache_path=self.cache_dir) del state_dict ttnn.synchronize_device(mesh) self.timings["model_build_s"] = time.perf_counter() - t1 self.timings["load_total_s"] = time.perf_counter() - t0 self.model._lm_head = self._no_lm_head self.device_dtypes = self._device_dtypes() logger.info(f"ClefEngine device dtypes and fidelities: {self.device_dtypes}") self.dram_free_after_text_weights = mesh_dram_free_bytes(mesh) self.vision = None self.last_vision = None if vision is None: vision = os.environ.get("CLEF_VISION", "1") == "1" if vision: t2 = time.perf_counter() vision_args = clef_vision.ClefVisionArgs( mesh, args.snapshot, max_batch_size=1, max_seq_len=VISION_MAX_SEQ_LEN ) self.vision = clef_vision.ClefVision(mesh, vision_args) self.vision.attach(self.model) self.model._alloc_vision_merge_buffers(mesh, MAX_BUCKET) ttnn.synchronize_device(mesh) self.timings["vision_build_s"] = time.perf_counter() - t2 self.timings["load_total_s"] = time.perf_counter() - t0 logger.info( f"ClefEngine vision tower attached in {self.timings['vision_build_s']:.1f} s: {self.vision.describe()}" ) if self.vision is not None and self.traced: self._warm_vision_tower() self.dram_free_after_weights = mesh_dram_free_bytes(mesh) self.blocks_per_slot = self.max_len // BLOCK_SIZE self.spare_blocks = 0 if self.traced: padded = ((self.blocks_per_slot + SPARE_BLOCKS + 31) // 32) * 32 self.spare_blocks = padded - self.blocks_per_slot self.gdn_layers = [layer.attention for layer in self.model.layers if not layer.is_full_attention] self.n_attention_layers = sum(1 for layer in self.model.layers if layer.is_full_attention) self.snapshot_slots = self._fit_slots(snapshot_slots, kv_reserve_bytes) kv_shape = [ self.snapshot_slots * self.blocks_per_slot + self.spare_blocks, args.n_local_kv_heads, BLOCK_SIZE, args.head_dim, ] self.model.allocate_kv_caches(kv_shape, ttnn.bfloat16, batch_size=1) self.page_tables = [ torch.arange(s * self.blocks_per_slot, (s + 1) * self.blocks_per_slot, dtype=torch.int32).unsqueeze(0) for s in range(self.snapshot_slots) ] self.snapshots = [ [(self._zeros_like(dn.rec_state), self._zeros_like(dn.conv_carry)) for dn in self.gdn_layers] for _ in range(self.snapshot_slots) ] self.device_dtypes["gdn_rec_state"] = str(self.gdn_layers[0].rec_state.dtype) self.prefix_hidden = [None] * self.snapshot_slots self.handles = [None] * self.snapshot_slots ttnn.synchronize_device(mesh) self.dram_free_after_slots = mesh_dram_free_bytes(mesh) self._tokenizer = None self._head = None self._processor = None if self.traced: self._setup_traces() self.timings["load_total_s"] = time.perf_counter() - t0 logger.info( f"ClefEngine ready: load {self.timings['load_total_s']:.1f} s (state dict " f"{self.timings['host_state_dict_s']:.1f} s, build {self.timings['model_build_s']:.1f} s), " f"DRAM free after weights {self.dram_free_after_weights / 2**30:.2f} GiB, after slots " f"{self.dram_free_after_slots / 2**30:.2f} GiB, slots {self.snapshot_slots}, " f"blocks/slot {self.blocks_per_slot}, traced={self.traced} traces={len(self.traces)} " f"trace_region_used={self.trace_bytes / 2**20:.1f} MiB, gdn_conv={self.gdn_conv_impl}, " f"planner={self.planner}, warm_grids={[g for g, _ in self.vision_warmed_grids]}" ) def _dev(self, t, dtype, layout): return ttnn.from_torch(t, dtype=dtype, layout=layout, device=self.mesh, mesh_mapper=self._rep) def _host(self, t, dtype, layout, mapper=None): return ttnn.from_torch(t, dtype=dtype, layout=layout, mesh_mapper=mapper or self._rep) def _write(self, t, buf, dtype, layout, mapper=None): ttnn.copy_host_to_device_tensor(self._host(t, dtype, layout, mapper), buf) self.counters["host_writes"] += 1 @contextlib.contextmanager def _misses_allowed(self): if not self.traces: yield return self._misses_depth = getattr(self, "_misses_depth", 0) + 1 if self._misses_depth == 1: self.mesh.set_program_cache_misses_allowed(True) try: yield finally: self._misses_depth -= 1 if self._misses_depth == 0: self.mesh.set_program_cache_misses_allowed(False) def _warm_vision_tower(self): grids = self.vision_warm_grid if grids is None: text = os.environ.get("CLEF_VISION_WARM_GRID") if text is None: raise ValueError(WARM_GRID_HELP) grids = parse_grids(text) elif grids and isinstance(grids[0], int): grids = [tuple(grids)] if not grids: raise ValueError(WARM_GRID_HELP) t0 = time.perf_counter() warmed = [] for grid in grids: grid_t = torch.tensor([list(grid)], dtype=torch.long) n_patches = int(grid_t.prod()) pixels = torch.randn( n_patches, 3 * 2 * 16 * 16, dtype=torch.float32, generator=torch.Generator().manual_seed(0) ) if int(grid[0]) > 1: tokens = self.model.get_video_features(pixels, grid_t) else: tokens = self.model.get_image_features(pixels, grid_t) ttnn.synchronize_device(self.mesh) ttnn.deallocate(tokens) warmed.append((tuple(grid), n_patches)) self.model._req_image_grid_thw = None self.model._req_video_grid_thw = None self.timings["vision_warm_s"] = time.perf_counter() - t0 self.vision_warmed_grids = warmed logger.info(f"vision tower warmed on grids {warmed} in {self.timings['vision_warm_s']:.1f} s") def _alloc_guard(self, nbytes): if nbytes <= 0: return None rows = max(32, (nbytes // (2 * 1024) // 32) * 32) try: guard = ttnn.zeros( [1, 1, rows, 1024], dtype=ttnn.bfloat16, layout=TILE, device=self.mesh, memory_config=ttnn.DRAM_MEMORY_CONFIG, ) except Exception as error: logger.warning(f"ttnn.zeros guard failed ({error}); uploading host zeros instead") guard = self._dev(torch.zeros(1, 1, rows, 1024, dtype=torch.bfloat16), ttnn.bfloat16, TILE) self.trace_guard = {"bytes": rows * 1024 * 2, "address": int(guard.buffer_address())} return guard def _setup_traces(self): mesh = self.mesh t0 = time.perf_counter() self._rep = ttnn.ReplicateTensorToMesh(mesh) self.model._build_request_rope(torch.zeros(1, 1, dtype=torch.long), None) self.csi = self._dev(torch.zeros(1, dtype=torch.int32), ttnn.int32, RM) first_spare = self.snapshot_slots * self.blocks_per_slot spare = torch.arange(first_spare, first_spare + self.spare_blocks, dtype=torch.int32).unsqueeze(0) self.full_tables = [torch.cat([pt, spare], dim=1).contiguous() for pt in self.page_tables] self.pt_full = self._dev(self.full_tables[0], ttnn.int32, RM) for b in BUCKETS: cos, sin = self.model._rope_tp_cos_sin_torch(0, b) self.bufs[b] = { "tok": self._dev(torch.zeros(1, b, dtype=torch.int32), ttnn.uint32, RM), "cpt": self._dev(self.full_tables[0][:, : b // BLOCK_SIZE].contiguous(), ttnn.int32, RM), "cos": self._dev(cos.contiguous(), ttnn.bfloat16, TILE), "sin": self._dev(sin.contiguous(), ttnn.bfloat16, TILE), "out": None, } guard = self._alloc_guard(self.trace_guard_bytes) for b in BUCKETS: for _ in range(2): self.model._reset_gdn_state_for_new_sequence() self._forward_body(b) self._copy_state(0, to_live=False) self._copy_state(0, to_live=True) self.model._reset_gdn_state_for_new_sequence() ttnn.synchronize_device(mesh) self.timings["trace_warm_s"] = time.perf_counter() - t0 t1 = time.perf_counter() mesh.set_program_cache_misses_allowed(False) for b in BUCKETS: self.traces[("fwd", b)] = self._capture(lambda b=b: self._forward_body(b)) for slot in range(self.snapshot_slots): self.traces[("restore", slot)] = self._capture(lambda s=slot: self._copy_state(s, to_live=True)) self.traces[("save", slot)] = self._capture(lambda s=slot: self._copy_state(s, to_live=False)) self.traces[("zero",)] = self._capture(self.model._reset_gdn_state_for_new_sequence) ttnn.synchronize_device(mesh) view = ttnn.get_memory_view(mesh, ttnn.BufferType.TRACE) self.trace_bytes = int(view.total_bytes_allocated_per_bank) * int(view.num_banks) self.timings["trace_capture_s"] = time.perf_counter() - t1 if guard is not None: ttnn.deallocate(guard) probe = self._dev(torch.zeros(1, 1024, dtype=torch.int32), ttnn.int32, RM) self.trace_guard["probe_address_after_release"] = int(probe.buffer_address()) ttnn.deallocate(probe) self.model._reset_gdn_state_for_new_sequence() ttnn.synchronize_device(mesh) self.dram_free_after_traces = mesh_dram_free_bytes(mesh) logger.info( f"{len(self.traces)} traces captured in {self.timings['trace_capture_s']:.1f} s after a " f"{self.timings['trace_warm_s']:.1f} s warmup; trace region used {self.trace_bytes / 2**20:.1f} MiB; " f"guard {getattr(self, 'trace_guard', None)}; DRAM free {self.dram_free_after_traces / 2**30:.2f} GiB" ) def _capture(self, body): tid = ttnn.begin_trace_capture(self.mesh, cq_id=0) try: body() except Exception: logger.error(f"trace capture failed:\n{traceback.format_exc()}") try: ttnn.end_trace_capture(self.mesh, tid, cq_id=0) ttnn.release_trace(self.mesh, tid) except Exception as error: logger.error(f"could not end the failed capture: {error}") raise ttnn.end_trace_capture(self.mesh, tid, cq_id=0) return tid def _replay(self, key): ttnn.execute_trace(self.mesh, self.traces[key], cq_id=0, blocking=False) self.counters["replays"] += 1 def _forward_body(self, b): B = self.bufs[b] self.model._chunked_chunk_size = b hidden = self.model._forward_prefill_chunk_tp(B["tok"], B["cos"], B["sin"], self.csi, self.pt_full, B["cpt"]) x = self.model.norm(hidden, mode=Mode.PREFILL) ttnn.deallocate(hidden) if B["out"] is None: B["out"] = x return ttnn.copy(x, B["out"]) ttnn.deallocate(x) def _copy_state(self, slot, to_live): for dn, (rec, carry) in zip(self.gdn_layers, self.snapshots[slot]): if to_live: ttnn.copy(rec, dn.rec_state) ttnn.copy(carry, dn.conv_carry) else: ttnn.copy(dn.rec_state, rec) ttnn.copy(dn.conv_carry, carry) def _stage_merge(self, tok_host, chunk_start, full_ids, vision): model = self.model if model._vis_buf is None: return if vision is None or vision.rows_host is None: ttnn.copy_host_to_device_tensor(model._vis_zero_mask_host, model._vis_mask_buf) self.counters["host_writes"] += 1 return offset = model._vis_row_offset_for(full_ids, chunk_start) cs = int(model._vis_buf.shape[-2]) flat = tok_host.reshape(-1) pos = torch.nonzero(flat[:cs] == model._vision_placeholder_token_id(), as_tuple=False).reshape(-1) n = int(pos.numel()) if n == 0: ttnn.copy_host_to_device_tensor(model._vis_zero_mask_host, model._vis_mask_buf) self.counters["host_writes"] += 1 return assert offset + n <= int( vision.rows_host.shape[0] ), f"vision splice out of range: {offset}+{n} > {vision.rows_host.shape[0]}" vis_full = torch.zeros(1, 1, cs, self.args.dim, dtype=torch.bfloat16) vis_full[0, 0, pos] = vision.rows_host[offset : offset + n] mask = torch.zeros(1, 1, cs, 1, dtype=torch.bfloat16) mask[0, 0, pos, 0] = 1.0 shard = ttnn.ShardTensor2dMesh(self.mesh, dims=(None, -1), mesh_shape=self.args.cluster_shape) self._write(vis_full, model._vis_buf, ttnn.bfloat16, TILE, mapper=shard) self._write(mask, model._vis_mask_buf, ttnn.bfloat16, TILE) def _replay_chunk(self, token_ids, chunk_start, slot, full_ids=None, vision=None): length = token_ids.shape[1] bucket = self.bucket_for(length) assert chunk_start % ALIGN == 0, f"chunk_start {chunk_start} is not {ALIGN}-aligned" B = self.bufs[bucket] tok = torch.zeros(1, bucket, dtype=torch.int32) tok[:, :length] = token_ids.to(torch.int32) self._write(tok, B["tok"], ttnn.uint32, RM) self._write(torch.tensor([chunk_start], dtype=torch.int32), self.csi, ttnn.int32, RM) self._write(self.full_tables[slot], self.pt_full, ttnn.int32, RM) blk0 = chunk_start // BLOCK_SIZE self._write( self.full_tables[slot][:, blk0 : blk0 + bucket // BLOCK_SIZE].contiguous(), B["cpt"], ttnn.int32, RM ) cos, sin = self.model._rope_tp_cos_sin_torch(chunk_start, bucket) self._write(cos.contiguous(), B["cos"], ttnn.bfloat16, TILE) self._write(sin.contiguous(), B["sin"], ttnn.bfloat16, TILE) self._stage_merge(tok, chunk_start, full_ids, vision) self._replay(("fwd", bucket)) out = ttnn.to_torch(B["out"], mesh_composer=ttnn.ConcatMeshToTensor(self.mesh, dim=0)) self.counters["readbacks"] += 1 return out[0].reshape(-1, self.args.dim)[:length].float() def _device_dtypes(self): out = {} gdn = next(layer for layer in self.model.layers if not layer.is_full_attention) attn = next((layer for layer in self.model.layers if layer.is_full_attention), None) mlp = gdn.feed_forward.weights out["mlp_gate_up_packed"] = str(mlp.w_gate_up.dtype) if mlp.w_gate_up is not None else None out["mlp_w1"] = str(mlp.w1.dtype) if mlp.w1 is not None else None out["mlp_down"] = str(mlp.w2.dtype) out["mlp_fidelity"] = str(gdn.feed_forward.compute_kernel_config.math_fidelity) tw = gdn.attention.tw out["gdn_qkvzab"] = str(tw["qkvz"].dtype) out["gdn_out_colpar"] = str(tw["out_colpar"].dtype) out["gdn_fidelity"] = str(gdn.attention.cfg.math_fidelity) out["gdn_gate_fp32"] = bool(gdn.attention._gate_fp32) out["gdn_dt_bias"] = str(tw["dt_bias"].dtype) out["gdn_rec_state"] = None if attn is not None: atw = attn.attention.tw key = "wqkv_fused" if "wqkv_fused" in atw else "wqkv" out["attn_qkv"] = str(atw[key].dtype) out["attn_wo"] = str(atw["wo"].dtype) out["attn_fidelity"] = str(attn.attention.compute_cfg.math_fidelity) out["embedding"] = str(self.model.embd.weights.dtype) return out def _no_lm_head(self, x): raise RuntimeError("ClefEngine never runs the LM head; output.weight is a stub") def _fit_slots(self, requested, reserve): kv_per_slot = ( 2 * self.n_attention_layers * self.blocks_per_slot * self.args.n_local_kv_heads * BLOCK_SIZE * self.args.head_dim * 2 ) rec_bytes = 4 if os.environ.get("QWEN35_GDN_STATE_BF16") != "1" else 2 gdn_per_slot = sum( dn.Nv * dn.Dk * dn.Dv * rec_bytes + ttnn.TILE_SIZE * dn.qkv_dim_tp * 2 for dn in self.gdn_layers ) per_slot = kv_per_slot + gdn_per_slot free = self.dram_free_after_weights slots = max(1, min(requested, (free - reserve) // per_slot)) self.slot_bytes = { "kv_per_slot": kv_per_slot, "gdn_per_slot": gdn_per_slot, "fit": int((free - reserve) // per_slot), } logger.info( f"KV per slot {kv_per_slot / 2**30:.3f} GiB + GDN snapshot {gdn_per_slot / 2**20:.1f} MiB, DRAM free " f"{free / 2**30:.2f} GiB, reserve {reserve / 2**30:.2f} GiB: {slots} slot(s) of {requested} requested " f"(max fit {self.slot_bytes['fit']})" ) return slots def _zeros_like(self, t): return ttnn.from_torch( torch.zeros(list(t.shape), dtype=torch.float32 if t.dtype == ttnn.float32 else torch.bfloat16), dtype=t.dtype, layout=TILE, device=self.mesh, mesh_mapper=ttnn.ReplicateTensorToMesh(self.mesh), ) @property def tokenizer(self): if self._tokenizer is None: self._tokenizer = clef_encode.load_tokenizer(self.args.snapshot) return self._tokenizer @property def head(self): if self._head is None: self._head = clef_head.load_head(self.args.snapshot) return self._head @property def processor(self): if self._processor is None: self._processor = clef_encode.load_processor(self.args.snapshot) return self._processor def _vision_request(self, media): if not media: return None pixels = media.get("pixel_values") video_pixels = media.get("pixel_values_videos") if pixels is None and video_pixels is None: return None if self.vision is None: raise ValueError("this engine was built without the vision tower (CLEF_VISION=0)") if pixels is not None and video_pixels is not None: raise ValueError("one media kind per request: the backbone stages one grid (images or videos)") video = video_pixels is not None grid = torch.as_tensor(media["video_grid_thw" if video else "image_grid_thw"]).reshape(-1, 3) if self.traced: self.check_warm_grids(grid) t0 = time.perf_counter() with self._misses_allowed(): if video: tokens = self.model.get_video_features(video_pixels, grid) else: tokens = self.model.get_image_features(pixels, grid) ttnn.synchronize_device(self.mesh) n_rows = int(tokens.shape[0]) rows_host = None if self.traced: rows_host = ttnn.to_torch(tokens, mesh_composer=ttnn.ConcatMeshToTensor(self.mesh, dim=1)).to( torch.bfloat16 ) ttnn.deallocate(tokens) tokens = None self.counters["tower_runs"] += 1 request = VisionRequest( tokens=tokens, grid=grid, video=video, n_rows=n_rows, n_patches=int(grid.prod(dim=1).sum()), seconds=time.perf_counter() - t0, rows_host=rows_host, ) self.last_vision = dict( video=video, grid=grid.tolist(), n_patches=request.n_patches, n_rows=request.n_rows, seconds=round(request.seconds, 4), tower=dict(self.vision.last_run), ) return request def check_warm_grids(self, grid): rows = torch.as_tensor(grid).reshape(-1, 3).tolist() if len(rows) > 1: raise ValueError(single_grid_message(len(rows))) warmed = {tuple(g) for g, _ in self.vision_warmed_grids} for row in rows: if tuple(row) not in warmed: raise ValueError( f"image grid {tuple(row)} is not in this traced server's warm list {sorted(warmed)}: " "resize the image to a warmed grid, add the grid to CLEF_VISION_WARM_GRID, " "or serve with CLEF_TRACED=0" ) def _stage_rope(self, token_ids, vision): if vision is None: self.model._build_request_rope(token_ids, None) return self.model._req_image_grid_thw = None if vision.video else vision.grid self.model._req_video_grid_thw = vision.grid if vision.video else None self.model._build_request_rope(token_ids, vision.marker) def _release_slot(self, slot): handle = self.handles[slot] self.handles[slot] = None if handle is not None and handle.vision is not None: handle.vision.release() @staticmethod def bucket_for(length): bucket = Qwen36Model._mask_bucket_for(length) assert bucket in BUCKETS, f"segment of {length} tokens needs bucket {bucket} > {MAX_BUCKET}" return bucket def _begin_sequence(self, token_ids, vision=None): if self.traced: self._replay(("zero",)) else: self.model._reset_gdn_state_for_new_sequence() self._stage_rope(token_ids, vision) def _run_chunk(self, token_ids, chunk_start, slot, full_ids=None, vision=None): if self.traced: return self._replay_chunk(token_ids, chunk_start, slot, full_ids, vision) length = token_ids.shape[1] bucket = self.bucket_for(length) assert chunk_start % ALIGN == 0, f"chunk_start {chunk_start} is not {ALIGN}-aligned" buf = torch.zeros(1, bucket, dtype=torch.int32) buf[:, :length] = token_ids.to(torch.int32) if self.model._vis_buf is not None: if vision is None: self.model._set_vision_merge(buf, None) else: offset = self.model._vis_row_offset_for(full_ids, chunk_start) self.model._set_vision_merge(buf, vision.tokens, offset) return self.model._forward_prefill_chunk_masked_tp(buf, length, chunk_start, self.page_tables[slot], bucket) def _read_hidden(self, hidden, length): x = self.model.norm(hidden, mode=Mode.PREFILL) ttnn.deallocate(hidden) out = ttnn.to_torch(x, mesh_composer=ttnn.ConcatMeshToTensor(self.mesh, dim=0)) ttnn.deallocate(x) return out[0].reshape(-1, self.args.dim)[:length].float() def _segments(self, token_ids, start, exact=False): T = token_ids.shape[1] if exact and self.traced: full = (T // self.chunk_size) * self.chunk_size cs = 0 for b in [self.chunk_size] * (full // self.chunk_size) + aligned_pieces(T - full): yield token_ids[:, cs : cs + b], start + cs cs += b return for cs in range(0, T, self.chunk_size): ce = min(cs + self.chunk_size, T) yield token_ids[:, cs:ce], start + cs def _run_rows(self, token_ids, start, slot, full_ids=None, vision=None, exact=False): rows = [] for piece, chunk_start in self._segments(token_ids, start, exact): if self.traced: rows.append(self._replay_chunk(piece, chunk_start, slot, full_ids, vision)) continue hidden = self._run_chunk(piece, chunk_start, slot, full_ids, vision) rows.append(self._read_hidden(hidden, piece.shape[1])) if not rows: return torch.empty(0, self.args.dim, dtype=torch.float32) return torch.cat(rows, dim=0) def _save_state(self, slot): if self.traced: self._replay(("save", slot)) return self._copy_state(slot, to_live=False) def _restore_state(self, slot): if self.traced: self._replay(("restore", slot)) return self._copy_state(slot, to_live=True) def prefill_hidden(self, token_ids, slot=0, media=None): T = token_ids.shape[1] assert token_ids.shape[0] == 1 assert T <= self.max_len, f"T={T} exceeds max_len={self.max_len}" assert 0 <= slot < self.snapshot_slots, f"slot {slot} out of range" self._release_slot(slot) self.prefix_hidden[slot] = None vision = self._vision_request(media) try: self._begin_sequence(token_ids, vision) return self._run_rows(token_ids, 0, slot, token_ids, vision) finally: if vision is not None: vision.release() def _chunk_cost(self, length, exact): if length <= 0: return 0.0 if exact: pieces = [self.chunk_size] * (length // self.chunk_size) + aligned_pieces(length % self.chunk_size) else: full, rest = divmod(length, self.chunk_size) pieces = [self.chunk_size] * full + ([self.bucket_for(rest)] if rest else []) return sum(self.cost_table[p] for p in pieces) def plan_prefix(self, S, tail_len=None): S0_max = (S // ALIGN) * ALIGN if not self.planner or tail_len is None: return S0_max exact = self.traced base_tail = self._chunk_cost(S - S0_max + tail_len, exact=False) best, best_cost = S0_max, self._chunk_cost(S0_max, exact=exact) + base_tail S0 = S0_max - ALIGN while S0 >= 0 and S0 > S0_max - self.chunk_size: tail_cost = self._chunk_cost(S - S0 + tail_len, exact=False) if tail_cost > base_tail: break cost = self._chunk_cost(S0, exact=exact) + tail_cost if cost < best_cost: best, best_cost = S0, cost S0 -= ALIGN return best def prefill_state(self, state_ids, slot=0, key=None, media=None, tail_len=None): S = state_ids.shape[1] S0 = self.plan_prefix(S, tail_len) assert state_ids.shape[0] == 1 assert S <= self.max_state_len, f"S={S} exceeds max_state_len={self.max_state_len}" assert 0 <= slot < self.snapshot_slots, f"slot {slot} out of range" self._release_slot(slot) vision = self._vision_request(media) self._begin_sequence(state_ids, vision) self.prefix_hidden[slot] = self._run_rows(state_ids[:, :S0], 0, slot, state_ids, vision, exact=True) self._save_state(slot) handle = StateHandle(S, S0, state_ids[:, S0:S].clone(), slot, key, state_ids.clone(), vision) self.handles[slot] = handle return handle def schema_hidden(self, handle, schema_ids): tail = torch.cat([handle.suffix_ids, schema_ids], dim=1) L = tail.shape[1] assert handle.S0 + L <= self.max_len, f"S0+L={handle.S0 + L} exceeds max_len={self.max_len}" assert self.handles[handle.slot] is handle, "slot no longer holds this state" self._restore_state(handle.slot) full_ids = torch.cat([handle.state_ids[:, : handle.S0], tail], dim=1) self._stage_rope(full_ids, handle.vision) return self._run_rows(tail, handle.S0, handle.slot, full_ids, handle.vision) def cached_hidden(self, state_ids, tail_ids, slot=0, media=None): key = clef_encode.cache_key(state_ids[0].tolist()) + media_key(media) handle = self.handles[slot] hit = handle is not None and handle.key == key if not hit: handle = self.prefill_state(state_ids, slot, key, media=media, tail_len=int(tail_ids.shape[1])) tail_hidden = self.schema_hidden(handle, tail_ids) return torch.cat([self.prefix_hidden[slot], tail_hidden], dim=0), hit def probs_for_request(self, record, mode="full", slot=0, max_length=None): if mode not in ("full", "cached"): raise ValueError(f"mode={mode!r}; expected 'full' or 'cached'") t0 = time.perf_counter() request = load_media_paths(record) has_media = bool(request.get("images") or request.get("videos")) encoded = clef_encode.encode( self.tokenizer, request, processor=self.processor if has_media else None, max_length=max_length or min(16384, self.max_len), ) media = encoded.media if has_media else None ids = torch.tensor([list(encoded.input_ids)], dtype=torch.long) t_enc = time.perf_counter() - t0 hit = None self.last_vision = None if mode == "full": hidden = self.prefill_hidden(ids, slot, media=media) else: state_part, tail_part, _ = clef_encode.split_for_cache(encoded, self.tokenizer, request) hidden, hit = self.cached_hidden( torch.tensor([state_part], dtype=torch.long), torch.tensor([tail_part], dtype=torch.long), slot, media=media, ) t_dev = time.perf_counter() - t0 - t_enc probs = clef_head.probs_for_record(self.head, hidden, ids[0], encoded, self.args.load_lm_head_rows) answers = { qid: clef_encode.release_module().systemone_answer(request["questions"][qid], dist) for qid, dist in probs.items() } total = time.perf_counter() - t0 timing = { "encode_s": round(t_enc, 3), "device_s": round(t_dev, 3), "head_s": round(total - t_enc - t_dev, 3), } if self.last_vision is not None: timing["vision_s"] = self.last_vision["seconds"] result = { "model": request.get("model", "clef"), "answers": answers, "usage": {"input_tokens": len(encoded.input_ids), "output_tokens": 0}, "probs": probs, "input_tokens": len(encoded.input_ids), "seconds": round(total, 3), "timing": timing, "mode": mode, "cache_hit": hit, } if self.last_vision is not None: result["vision"] = dict(self.last_vision) return result