QSA KV cache host-RAM offload for ExLlamaV3 =========================================== Keeps the QSA attention K/V planes, and the indexer's raw key plane, in pinned device-mapped host memory instead of VRAM, so the context a model can hold stops being bounded by the VRAM left over after its weights. Only the pooled indexer plane stays on the device -- it is the one plane whose traffic is linear in context length, since the indexer scores every block of it on every step. VRAM 27.75 -> 0.75 KiB/token host RAM 0 -> 27.0 KiB/token Measured on 3x RTX 3090 (70.68 GiB) with a 4.05bpw EXL3 quant of Qwen3.8-Flash-Next, whose weights alone take 64.6 GiB: baseline 65,536 tokens does not load this patch 3,145,728 tokens, decode 53.6 t/s at 262k, 50.7 t/s at 1M (64.3 t/s at zero context, so -21% for a 1M context) Enable with -kvo / --kv_offload, or EXL3_QSA_KV_OFFLOAD=1. QSA models only. Base: turboderp-org/exllamav3 ca13bdd ("Bump to v1.4.7") Apply: git apply qsa-kv-offload.patch New files exllamav3/cache/qsa_offload.py CacheLayer_qsa_offload exllamav3/util/qsa_kvo_stats.py EXL3_QSA_KVO_STATS=1 instrumentation Changed exllamav3/modules/attention_fn/qsa_triton.py staged prefill gather exllamav3/modules/qsa_indexer.py dispatch + stats exllamav3/modules/attn.py cache layer selection + stats exllamav3/model/config.py infer_params.qsa_kv_offload exllamav3/model_init.py -kvo, -ccs exclusion exllamav3/generator/cpu_cache.py -ccs exclusion Environment EXL3_QSA_KV_OFFLOAD 0 same as -kvo EXL3_QSA_KVO_RAW 1 0 keeps the raw indexer plane in VRAM EXL3_QSA_KVO_ARENA 32MiB prefill staging arena per device EXL3_QSA_KVO_STAGE 1 0 disables staged prefill (direct gather) EXL3_QSA_KVO_STAGE_ROWS 1024 query rows per staged kernel launch EXL3_QSA_KVO_STATS 0 1 prints per-layer offload statistics at exit Note: at a context that leaves little VRAM free, use a prefill chunk of 2048 (the generator's default); 4096 does not fit alongside these weights, with or without this patch. diff --git a/exllamav3/generator/cpu_cache.py b/exllamav3/generator/cpu_cache.py index 731f209..79ca7cf 100644 --- a/exllamav3/generator/cpu_cache.py +++ b/exllamav3/generator/cpu_cache.py @@ -78,8 +78,19 @@ class CPUPageCache: # the workers and the ranks hold the host buffers themselves. self.segments = [] offset = 0 + from ..cache.qsa_offload import CacheLayer_qsa_offload for cache in local_caches: for layer in cache.layers.values(): + # Exclusive with the KV offload. Living in host memory is not the same as + # surviving eviction -- an offloaded page is still recycled for the next + # sequence -- so the tier would have to keep copying those planes into its own + # slots to restore anything, which is the same host RAM spent twice. Narrowing + # the tier to the planes still in VRAM is not an option either: a page image + # without K/V restores the wrong K/V + assert not isinstance(layer, CacheLayer_qsa_offload), \ + "CPU page cache tier cannot be combined with QSA KV offload (--kv_offload): " \ + "the offloaded planes are already host-resident, and a page image without " \ + "them could not restore a page correctly." for t in layer.get_tensors(): if t is None: continue diff --git a/exllamav3/model/config.py b/exllamav3/model/config.py index a8f3114..3a9563d 100644 --- a/exllamav3/model/config.py +++ b/exllamav3/model/config.py @@ -57,6 +57,14 @@ class InferParams: # per-forward row gathers instead of loading the whole table into system RAM (tens of # GB). Set before loading the model self.ngram_stream_from_disk = os.environ.get("EXL3_NGRAM_STREAM", "1") != "0" + # Keep QSA attention's K/V planes, and the indexer's raw key plane, in pinned + # device-mapped host memory instead of VRAM, leaving only the pooled plane on the device. + # Both are read in amounts fixed by the indexer budget rather than by context length -- + # the gather takes a top-k selection, the pool rebuild only touches the write head -- so + # the PCIe cost is constant while the VRAM saved is linear, which is what buys the very + # long contexts. The pooled plane is NOT offloaded: it is scanned in full every step. Set + # before loading the model + self.qsa_kv_offload = os.environ.get("EXL3_QSA_KV_OFFLOAD", "0") != "0" def use_mgemm(self, K: int, out_features: int, mul1: bool = False, device = None) -> bool: # Unfusing only pays when the separate GEMV calls can actually take the int8 path, which diff --git a/exllamav3/model_init.py b/exllamav3/model_init.py index b511795..73e75d0 100755 --- a/exllamav3/model_init.py +++ b/exllamav3/model_init.py @@ -57,6 +57,7 @@ def add_args( parser.add_argument("-tp", "--tensor_parallel", action = "store_true", help = "Load model in Tensor-parallel mode, attempts to respect --gpu_split") parser.add_argument("-mcl", "--moe_cpu_offload", type = int, help = "Experimental: run the routed experts of the first N block-sparse MoE layers on the CPU, with expert weights in system RAM. Layer-split mode only; requires mul1-codebook experts (ineligible layers fall back to the GPU)", default = 0) parser.add_argument("-mcs", "--moe_cpu_split", type = int, help = "Experimental: per-layer expert split — run the TAIL N routed experts of every eligible block-sparse MoE layer on the CPU, overlapping the CPU GEMMs with each layer's own GPU expert compute. Dynamic hot/cold expert placement is on by default (EXL3_MOE_CPU_SWAP=0 for static placement). Mutually exclusive with --moe_cpu_offload. Layer-split mode only; requires mul1-codebook experts", default = 0) + parser.add_argument("-kvo", "--kv_offload", action = "store_true", help = "Store QSA attention K/V and the indexer's raw key plane in pinned host memory (read zero-copy over PCIe) instead of VRAM, leaving only the pooled plane on the device (0.75 vs 27.75 KiB/token of VRAM). Trades a constant per-token PCIe read for a VRAM saving linear in context length. EXL3_QSA_KVO_RAW=0 keeps the raw plane in VRAM. QSA models only; not compatible with a quantized cache or with --cpu_cache_size") parser.add_argument("-mct", "--moe_cpu_threads", type = int, help = "Worker thread count for --moe_cpu_offload / --moe_cpu_split (default: EXL3_MOE_CPU_THREADS env, else cpu_count/2)", default = None) parser.add_argument("-ngr", "--ngram_ram", action = "store_true", help = "Load an n-gram embedding table (PLE models, e.g. Qwen3.8-Flash-Next) fully into system RAM instead of streaming rows from disk per forward (tens of GB of RAM; avoids per-token disk reads)") parser.add_argument("-tpb", "--tp_backend", type = str, help = "Tensor-parallel backend, either 'native' (default) or 'nccl'", default = "native") @@ -209,6 +210,12 @@ def init( config.infer_params.moe_cpu_threads = args.moe_cpu_threads if getattr(args, "ngram_ram", False): config.infer_params.ngram_stream_from_disk = False + if getattr(args, "kv_offload", False): + # Both would keep a copy of every page in system RAM, and the CPU tier's page images + # would not include the K/V it no longer owns + assert not getattr(args, "cpu_cache_size", 0), \ + "--kv_offload and --cpu_cache_size are mutually exclusive: K/V already lives in host memory" + config.infer_params.qsa_kv_offload = True if override_dynamic_seq_len: config.override_dynamic_seq_len(override_dynamic_seq_len) dmcl = getattr(args, "draft_moe_cpu_layers", 0) dmclt = getattr(args, "moe_cpu_threads", None) diff --git a/exllamav3/modules/attention_fn/qsa_triton.py b/exllamav3/modules/attention_fn/qsa_triton.py index 33dae67..7c3b2cf 100644 --- a/exllamav3/modules/attention_fn/qsa_triton.py +++ b/exllamav3/modules/attention_fn/qsa_triton.py @@ -29,6 +29,7 @@ paged/BC form runs decode rows (q_len == 1), the flat form any (B * S) row set. runtime arguments or derived on device, so the kernels are CUDA-graph-safe. """ +import os import torch try: @@ -247,6 +248,179 @@ if has_triton: return _sm_counts[dev.index] + # Default staging arena per device, in bytes (EXL3_QSA_KVO_ARENA). Deliberately small: this + # path exists for configurations where VRAM is the binding constraint, and on one packed to + # the brim by autosplit even a couple of hundred MB per device is the difference between + # loading and not. A smaller arena only means more buckets, and the transfer volume -- the + # thing being fixed here -- is identical either way; only the kernel's own re-scan of the + # selection is repeated. Raise it when there is headroom to spare + ARENA_BYTES = int(os.environ.get("EXL3_QSA_KVO_ARENA", 32 * 1024 ** 2)) + + # EXL3_QSA_KVO_STAGE=0 forces the direct (per-row, straight-from-host) gather, for A/B + STAGING_ENABLED = os.environ.get("EXL3_QSA_KVO_STAGE", "1") != "0" + + # Query rows per kernel launch. The launch's partial buffers are linear in this, and they are + # charged against exactly the VRAM the offload exists to free, so a long prefill chunk is + # served a block at a time. The blocks run inside the staging loop, against an arena that is + # already filled, so splitting them costs launches and not a single extra byte over PCIe + STAGE_ROWS = int(os.environ.get("EXL3_QSA_KVO_STAGE_ROWS", 1024)) + + + def qsa_staging_worthwhile(rows: int, k_pad: int, num_pages_used: int, page_size: int) -> bool: + """Is staging the page range cheaper than letting the gather read host memory directly? + + The direct path reads rows * k_pad cache positions (every row fetches its own selection + over PCIe, and rows overlap heavily); staging reads each page of the history exactly once, + i.e. num_pages_used * page_size positions. Prefill chunks sit far on the staging side of + that comparison -- 2048 rows x 2080 selections against a history of at most ~4096 pages -- + while a one-row decode fallback sits on the other.""" + return STAGING_ENABLED and rows * k_pad > 2 * num_pages_used * page_size + + + def qsa_sparse_attend_rows_staged( + q: torch.Tensor, # (R, n_q_heads, head_dim) fp16, normed + roped + k_cache: torch.Tensor, # (pages, page_size, n_kv_heads, head_dim) fp16, HOST-resident + v_cache: torch.Tensor, + indices: torch.Tensor, # (R, K_pad) int32 cache positions, -1 padded + sm_scale: float, + block_table: torch.Tensor, # (num_pages,) int32, the sequence's page table + num_pages_used: int, # pages of it the selection can reach + arena_bytes: int = 0, + layer_idx: int | None = None, + ) -> torch.Tensor: + """Sparse gather over a host-resident K/V cache, staging the history through a VRAM arena + instead of letting every query row pull its own selection over PCIe. + + Rows in a prefill chunk each pick ~2048 cache positions, and the union of those picks + approaches the whole history, so the direct path re-reads the same pages thousands of + times per chunk -- constant per row, but multiplied by the row count. Here the page range + is walked once: each bucket of pages is copied host -> arena in one pass, every query row + is scored against it, and the bucket's partial (o, m, l) is folded into a running softmax. + Each page therefore crosses PCIe exactly once per chunk per layer, whatever the row count. + + Staging is the outer loop and query rows the inner one, deliberately: blocking the rows + bounds the kernel's partial buffers (they are charged against the same VRAM the offload + exists to free), and doing it inside the bucket keeps that free -- the other nesting would + re-stage the whole page range once per row block. + + All rows must share one page table (bsz == 1), which is what prefill and every + single-sequence eager fallback give us. + """ + from ...util import qsa_kvo_stats + + R, H, hd = q.shape + kvh = k_cache.shape[2] + page_size = k_cache.shape[1] + group = H // kvh + BLOCK_H = 16 + BLOCK_N = 32 + h_blocks = triton.cdiv(group, BLOCK_H) + K_pad = indices.shape[1] + dev = q.device + assert q.is_contiguous() and indices.is_contiguous() + assert K_pad % BLOCK_N == 0, "staged path needs K_pad aligned to the kernel's BLOCK_N" + + # Every workspace here is allocated per call and released with it, like the direct path's + # partials. Keeping them in the shared tensor cache instead would subtract permanently + # from the allocator's pool, and on a device autosplit has packed to the brim what has to + # fit is the largest single module's peak, not the sum of everyone's + page_bytes = page_size * kvh * hd * 2 + arena_pages = max(1, (arena_bytes or ARENA_BYTES) // (2 * page_bytes)) + arena_pages = min(arena_pages, num_pages_used) + k_arena = torch.empty((arena_pages, page_size, kvh, hd), dtype = torch.half, device = dev) + v_arena = torch.empty((arena_pages, page_size, kvh, hd), dtype = torch.half, device = dev) + k_flat = k_arena.view(-1, kvh, hd) + v_flat = v_arena.view(-1, kvh, hd) + + rb = min(STAGE_ROWS, R) + po = torch.empty((rb * kvh * h_blocks * BLOCK_H * hd,), dtype = torch.float, device = dev) + pml = torch.empty((rb * kvh * h_blocks * BLOCK_H * 2,), dtype = torch.float, device = dev) + out = torch.empty((R, H, hd), dtype = torch.half, device = dev) + + # One bucket is the common case (the whole reachable history fits the arena), and there is + # then nothing to fold: the kernel's own partial IS the result, and no selection has to be + # masked because every page is present + single = num_pages_used <= arena_pages + if not single: + acc = torch.zeros((R, kvh, group, hd), dtype = torch.float, device = dev) + m = torch.full((R, kvh, group), -float("inf"), dtype = torch.float, device = dev) + l = torch.zeros((R, kvh, group), dtype = torch.float, device = dev) + idx_buf = torch.empty((rb, K_pad), dtype = torch.int32, device = dev) + neg1 = torch.tensor(-1, dtype = torch.int32, device = dev) + zero = torch.zeros((), dtype = torch.float, device = dev) + + pages = torch.arange(block_table.shape[0], dtype = torch.int32, device = dev) + + for p0 in range(0, num_pages_used, arena_pages): + p1 = min(p0 + arena_pages, num_pages_used) + n = p1 - p0 + + # Host -> arena, one contiguous read per page. Stream-ordered ahead of the kernels + sel = block_table[p0 : p1].long() + torch.index_select(k_cache, 0, sel, out = k_arena[:n]) + torch.index_select(v_cache, 0, sel, out = v_arena[:n]) + if qsa_kvo_stats.enabled() and layer_idx is not None: + qsa_kvo_stats.record_stage(layer_idx, 2 * n * page_bytes) + + # A page table that redirects this bucket's pages to their arena slots. Entries + # outside it are never loaded (their selections are masked to -1 below), so the + # clamped values they hold do not matter + bt_arena = (pages - p0).clamp_(0, max(n - 1, 0)).unsqueeze(0) + lo, hi = p0 * page_size, p1 * page_size + + for r0 in range(0, R, rb): + r1 = min(r0 + rb, R) + rows = r1 - r0 + programs = rows * kvh * h_blocks + idx_r = indices[r0 : r1] + if single: + idx_use = idx_r + else: + # Pages map to contiguous position ranges, so restricting the selection to + # this bucket is a range test on the cache positions themselves + idx_use = idx_buf[:rows] + torch.where((idx_r >= lo) & (idx_r < hi), idx_r, neg1, out = idx_use) + + with torch.cuda.device(dev): + _qsa_sparse_split_kernel[(programs, 1)]( + q[r0 : r1], k_flat, v_flat, bt_arena, idx_use, po, pml, + K_pad, 0, 1, K_pad, + n_q_heads = H, n_kv_heads = kvh, page_size = page_size, + head_dim = hd, K_pad = K_pad, scale = float(sm_scale), + BLOCK_H = BLOCK_H, BLOCK_N = BLOCK_N, PAGED = 1, + num_warps = 4, num_stages = 2, + ) + + # pid = (row * n_kv_heads + kv_head) * h_blocks + h_block, and within a program + # row r is q head kv_head * group + h_block * BLOCK_H + r, so the program/row axes + # unflatten straight into (rows, kv_head, q head in group) once BLOCK_H's padding + # past group is dropped + po_v = po[: programs * BLOCK_H * hd] \ + .view(rows, kvh, h_blocks * BLOCK_H, hd)[:, :, :group] + pml_v = pml[: programs * BLOCK_H * 2] \ + .view(rows, kvh, h_blocks * BLOCK_H, 2)[:, :, :group] + if single: + o = po_v / pml_v[..., 1].clamp_min(1e-30).unsqueeze(-1) + out[r0 : r1] = o.reshape(rows, H, hd).to(torch.half) + continue + + # Fold this bucket into the running softmax. A bucket a row selected nothing from + # arrives as m = -inf / l = 0 and contributes nothing + a_r, m_r, l_r = acc[r0 : r1], m[r0 : r1], l[r0 : r1] + m_b, l_b = pml_v[..., 0], pml_v[..., 1] + m_new = torch.maximum(m_r, m_b) + alpha = torch.where(m_r > -float("inf"), torch.exp(m_r - m_new), zero) + beta = torch.where(m_b > -float("inf"), torch.exp(m_b - m_new), zero) + a_r.mul_(alpha.unsqueeze(-1)).add_(po_v * beta.unsqueeze(-1)) + l_r.mul_(alpha).add_(l_b * beta) + m_r.copy_(m_new) + + if not single: + o = acc / l.clamp_min(1e-30).unsqueeze(-1) + out.copy_(o.reshape(R, H, hd)) + return out + + def qsa_sparse_attend_rows( q: torch.Tensor, # (R, n_q_heads, head_dim) fp16, normed + roped k: torch.Tensor, # (rows, n_kv_heads, head_dim) fp16 (paged: flat cache view) diff --git a/exllamav3/modules/attn.py b/exllamav3/modules/attn.py index ea7c161..2f89b39 100644 --- a/exllamav3/modules/attn.py +++ b/exllamav3/modules/attn.py @@ -10,6 +10,7 @@ from .multilinear import MultiLinear from ..ext import exllamav3_ext as ext from ..model.model_tp_alloc import TPAllocation from ..util import profile_opt +from ..util import qsa_kvo_stats import os from .attention_fn.bc_attn import bc_attn_enable as _bc_attn_enable, build_bc_attn, MAX_BSZ as _bc_max_bsz, MAX_QLEN as _bc_max_qlen @@ -832,6 +833,9 @@ class Attention(Module): from ..cache.qsa import CacheLayer_qsa assert default is CacheLayer_fp16, \ "QSA attention currently supports only the fp16 cache layer" + if self.config is not None and getattr(self.config.infer_params, "qsa_kv_offload", False): + from ..cache.qsa_offload import CacheLayer_qsa_offload + return CacheLayer_qsa_offload, kwargs return CacheLayer_qsa, kwargs @@ -932,6 +936,15 @@ class Attention(Module): o = self.bc_attn_step(x, cache, params, block_table, cache_seqlens, host_seqlens = qsa_seqlens_cpu) if o is not None: + if qsa_sparse and qsa_kvo_stats.enabled(): + # The selection is built inside the capture, so only the padded bound is + # visible from here. Same widths as select_indices_paged / _qsa_sparse_geometry + idx = self.qsa_indexer + cr = idx.compress_ratio + k_pad = -(-(idx.block_topk * cr + cr - 1) // 32) * 32 + qsa_kvo_stats.record_sparse( + self.layer_idx, bsz * seqlen, k_pad, self.num_kv_heads, self.head_dim, + decode = True) return o if self.qsa_indexer is not None: diff --git a/exllamav3/modules/qsa_indexer.py b/exllamav3/modules/qsa_indexer.py index 4ac9309..c3992c1 100644 --- a/exllamav3/modules/qsa_indexer.py +++ b/exllamav3/modules/qsa_indexer.py @@ -647,15 +647,41 @@ class QSAIndexer(Module): update_planes. Returns (bsz, seq, num_q_heads, head_dim) fp16. """ from .attention_fn.qsa_triton import qsa_sparse_attend_rows + from ..util import qsa_kvo_stats bsz, seq = q.shape[:2] indices = self.select_indices_paged(layer, q_idx, block_table, cache_seqlens_cpu) + if qsa_kvo_stats.enabled(): + qsa_kvo_stats.record_sparse( + attn.layer_idx, bsz * seq, indices.shape[1], attn.num_kv_heads, attn.head_dim, + indices = indices) + qf = q.reshape(bsz * seq, attn.num_q_heads, attn.head_dim).contiguous() + page_size = layer.k.shape[1] + + # Host-resident K/V: every row of a prefill chunk would otherwise pull its own ~2048 + # selections across PCIe, re-reading the same history thousands of times per chunk. Stage + # the page range through a VRAM arena instead, so each page crosses once. Needs one shared + # page table, so batched/verify calls keep the direct path (their row counts are small + # enough that it is the cheaper one anyway) + from ..cache.qsa_offload import CacheLayer_qsa_offload + if isinstance(layer, CacheLayer_qsa_offload) and bsz == 1: + from .attention_fn.qsa_triton import ( + qsa_staging_worthwhile, qsa_sparse_attend_rows_staged) + npu = -(-(int(cache_seqlens_cpu.max().item()) + seq) // page_size) + npu = min(npu, block_table.shape[1]) + if qsa_staging_worthwhile(bsz * seq, indices.shape[1], npu, page_size): + o = qsa_sparse_attend_rows_staged( + qf, layer.k, layer.v, indices, attn.sm_scale, + block_table[0].int().contiguous(), npu, layer_idx = attn.layer_idx, + ) + return o.view(bsz, seq, attn.num_q_heads, attn.head_dim) + bt_rows = block_table.int().unsqueeze(1).expand(bsz, seq, -1) \ .reshape(bsz * seq, -1).contiguous() o = qsa_sparse_attend_rows( - q.reshape(bsz * seq, attn.num_q_heads, attn.head_dim).contiguous(), + qf, layer.k.view(-1, attn.num_kv_heads, attn.head_dim), layer.v.view(-1, attn.num_kv_heads, attn.head_dim), indices, attn.sm_scale, - block_table = bt_rows, page_size = layer.k.shape[1], + block_table = bt_rows, page_size = page_size, ) return o.view(bsz, seq, attn.num_q_heads, attn.head_dim) diff --git a/exllamav3/cache/qsa_offload.py b/exllamav3/cache/qsa_offload.py new file mode 100644 index 0000000..12d6094 --- /dev/null +++ b/exllamav3/cache/qsa_offload.py @@ -0,0 +1,196 @@ +from __future__ import annotations +from typing_extensions import override +import mmap +import os +import numpy as np +import torch +from .qsa import CacheLayer_qsa +from ..ext import exllamav3_ext as ext +from ..model.model_tp_cuda import ( + cuda_host_register, + cuda_host_unregister, + CUDA_HOST_REGISTER_PORTABLE, + CUDA_HOST_REGISTER_MAPPED, +) + +# EXL3_QSA_KVO_RAW=0 keeps the indexer's raw key plane in VRAM (the phase-1 placement). It is +# offloaded by default: it is read only around the write head -- the pool kernel rebuilds just the +# blocks an append touches -- so it costs a few KiB per step over PCIe while freeing 3 KiB per +# token of VRAM, five times what the K/V planes leave behind +_OFFLOAD_RAW = os.environ.get("EXL3_QSA_KVO_RAW", "1") != "0" + +_announced = False + + +class CacheLayer_qsa_offload(CacheLayer_qsa): + """ + QSA cache layer with the K/V planes in pinned, device-mapped host memory and only the + indexer's side planes in VRAM. + + The gather kernels reach K/V through a base pointer plus computed offsets and never touch a + stride or a device property of those tensors, so handing them a zero-copy CUDA alias of host + memory needs no kernel change: `sparse_attend`, `get_kv` and `ext.paged_kv_cache_update` all + keep working, and `build_bc_attn`'s `layer.k.device == module.device` check passes because + the alias really is a cuda:N tensor. + + What stays in VRAM is `pooled`, and that is not an optimisation but a requirement: the indexer + scores EVERY block of the pooled plane on every step, so that plane's traffic is linear in + context length (732 MiB per token at 1M across the full-attention layers). Everything else is + read in bounded amounts -- K/V through a top-k selection whose size is fixed by the indexer + budget, `raw_k` only around the write head, where the pool kernel rebuilds the blocks an append + touched -- so their traffic is constant in context length. That is the whole reason this trade + works, and why `pooled` is the one plane it cannot include. + + The slab is anonymous mmap memory registered with cudaHostRegister(PORTABLE | MAPPED) while + the layer's device is current, not a `pin_memory = True` tensor. Both give pinned, mapped + host memory, but torch resolves a device pointer's owning device through + cudaPointerGetAttributes, which reports whichever device was current when the region was + pinned -- and torch's caching host allocator recycles freed blocks across devices, so a slab + freed on one device (autosplit rolls a layer back and retries on the next device) comes back + bound to the wrong one and the alias is rejected. Registering the mapping ourselves ties each + slab to its layer's device deterministically. mmap also guarantees the page alignment + cudaHostRegister wants, and hands back zeroed pages without a 2 GiB memset. + """ + + def __init__( + self, + config, + attention, + cache_id: int, + max_num_tokens: int, + ): + super().__init__(config, attention, cache_id, max_num_tokens) + # The pinned host allocation backing self.k / self.v. The device aliases do NOT own it, + # so these references are what keep them valid + self.host_map = None # mmap object (owns the pages) + self.host_slab = None # CPU tensor over it (what pinned_cuda_view aliases) + self.host_ptr = 0 # registered base address, for unregister + self.offload_raw = _OFFLOAD_RAW + # CPU-side views of the same storage the device aliases point at, for copy_page + self.host_k = None + self.host_v = None + self.host_raw_k = None + + @override + def alloc(self, device: torch.device): + global _announced + + dev = torch.device(device) + assert dev.type == "cuda", \ + "QSA KV offload requires a CUDA device (the K/V aliases are device pointers)." + self.device = device + + n = int(np.prod(self.shape)) if self.shape else 0 + nr = int(np.prod(self.raw_k_shape)) if self.offload_raw else 0 + total = 2 * n + nr + + if total: + nbytes = total * torch.half.itemsize + if not _announced: + _announced = True + print(f" -- QSA KV offload: pinning {nbytes / 1024 ** 3:.2f} GiB of host memory " + f"per attention layer; the first touch of each is slow (page-table setup).") + # One slab for every offloaded plane: pinning is per-allocation work, and they are all + # allocated and freed together + idx = dev.index if dev.index is not None else 0 + self.host_map = mmap.mmap(-1, nbytes) + self.host_slab = torch.frombuffer(self.host_map, dtype = torch.half, count = total) + self.host_ptr = self.host_slab.data_ptr() + with torch.cuda.device(idx): + cuda_host_register(self.host_ptr, nbytes, + CUDA_HOST_REGISTER_PORTABLE | CUDA_HOST_REGISTER_MAPPED) + alias = ext.pinned_cuda_view(self.host_slab, idx) + else: + alias = None + + if self.shape is None: + self.k = None + self.v = None + else: + self.k = alias[:n].view(self.shape) + self.v = alias[n : 2 * n].view(self.shape) + self.host_k = self.host_slab[:n].view(self.shape) + self.host_v = self.host_slab[n : 2 * n].view(self.shape) + + if self.offload_raw: + self.raw_k = alias[2 * n :].view(self.raw_k_shape) + self.host_raw_k = self.host_slab[2 * n :].view(self.raw_k_shape) + else: + self.raw_k = torch.zeros(self.raw_k_shape, dtype = torch.half, device = device) + self.pooled = torch.zeros(self.pooled_shape, dtype = torch.half, device = device) + + @override + def free(self): + # Order matters: the aliases are non-owning views of the slab's storage, so every + # reference to them has to go before the slab does + self.k = None + self.v = None + self.raw_k = None + self.pooled = None + self.host_k = None + self.host_v = None + self.host_raw_k = None + self.device = None + if self.host_ptr: + cuda_host_unregister(self.host_ptr) + self.host_ptr = 0 + self.host_slab = None + if self.host_map is not None: + try: + self.host_map.close() + except BufferError: + # Something still holds a buffer view of the mapping; dropping the reference + # leaves the unmap to the collector, which is fine now that it is unregistered + pass + self.host_map = None + + @override + def copy_page(self, source: "CacheLayer_qsa_offload", from_page: int, to_page: int, + num_tokens: int): + # Host-resident planes are copied on the CPU rather than by a device kernel reading and + # writing the same memory across PCIe. It measures ~20% faster, but the reason to do it is + # that it leaves the PCIe link alone: that link is what the offload spends on the + # attention gather, and a page copy has no business competing for it. + # + # The device stream has to be drained first. A page being read may still be waiting on a + # write from the current forward, and a page being written may still be under a read, and + # a host memcpy is outside the stream's ordering entirely + assert self.shape == source.shape + torch.cuda.current_stream(torch.device(self.device)).synchronize() + nt = num_tokens + if self.shape is not None: + self.host_k[to_page, :nt].copy_(source.host_k[from_page, :nt]) + self.host_v[to_page, :nt].copy_(source.host_v[from_page, :nt]) + if self.offload_raw: + self.host_raw_k[to_page, :nt].copy_(source.host_raw_k[from_page, :nt]) + else: + self.raw_k[to_page, :nt].copy_(source.raw_k[from_page, :nt], non_blocking = True) + nb = (num_tokens + self.compress_ratio - 1) // self.compress_ratio + self.pooled[to_page, :nb].copy_(source.pooled[from_page, :nb], non_blocking = True) + + @override + def storage_size(self): + # VRAM only. Autosplit and the TP planner size device allocations from this, and counting + # the host-resident planes here would have them refuse a context that fits perfectly well + n = np.prod(self.pooled_shape) + if not self.offload_raw: + n += np.prod(self.raw_k_shape) + return n * torch.half.itemsize + + def host_size(self): + """Bytes of pinned host memory this layer holds. Not part of the CacheLayer interface; + used for reporting, since storage_size() deliberately hides it.""" + n = 2 * np.prod(self.shape) if self.shape else 0 + if self.offload_raw: + n += np.prod(self.raw_k_shape) + return n * torch.half.itemsize + + @override + def tp_export(self, plan): + return { + "cls": CacheLayer_qsa_offload, + "args": { + "cache_id": self.cache_id, + "max_num_tokens": self.max_num_tokens + } + } diff --git a/exllamav3/util/qsa_kvo_stats.py b/exllamav3/util/qsa_kvo_stats.py new file mode 100644 index 0000000..6b13e4d --- /dev/null +++ b/exllamav3/util/qsa_kvo_stats.py @@ -0,0 +1,145 @@ +from __future__ import annotations +import atexit +import os +import torch + +# Phase-0 instrumentation for the QSA KV host-offload path (EXL3_QSA_KVO_STATS=1). +# +# Every judgement in the offload plan rests on one claim: a QSA row reads a constant amount of +# K/V (4 * block_topk + tail, padded to the kernel's K_pad) no matter how long the context is. +# This module is what makes that claim checkable instead of assumed. It records, per attention +# layer, how many sparse rows were served, how many cache positions those rows actually selected +# (counted on-device, so no synchronisation is added to the hot path), the resulting host->device +# read volume, and how much was moved by explicit prefill staging (phase 2). +# +# Counting is deliberately split into an exact term and a bound: +# - the eager/prefill path builds its index list in python, so the valid (non -1) count is +# accumulated exactly into a device-side counter; +# - the graph-captured decode path builds indices inside the capture, where python cannot see +# them, so only the padded K_pad bound is recorded. The bound is the transfer figure the plan +# budgets against anyway, so both numbers are reported side by side. + +_ENABLED = os.environ.get("EXL3_QSA_KVO_STATS", "0") != "0" + + +def enabled() -> bool: + return _ENABLED + + +class _LayerStats: + __slots__ = ("sparse_calls", "rows", "k_pad_max", "bound_bytes", "counter", + "decode_calls", "decode_rows", "stage_calls", "stage_bytes") + + def __init__(self): + self.sparse_calls = 0 + self.rows = 0 + self.k_pad_max = 0 + self.bound_bytes = 0 + self.counter = None # device tensor, exact count of selected cache positions + self.decode_calls = 0 + self.decode_rows = 0 + self.stage_calls = 0 + self.stage_bytes = 0 + + +_layers: dict[int, _LayerStats] = {} +_reported = False + + +def _layer(layer_idx: int) -> _LayerStats: + st = _layers.get(layer_idx) + if st is None: + st = _layers[layer_idx] = _LayerStats() + return st + + +def record_sparse(layer_idx: int, rows: int, k_pad: int, kv_heads: int, head_dim: int, + indices: torch.Tensor | None = None, decode: bool = False): + """One sparse-attention launch over `rows` query rows, each reading up to `k_pad` cache + positions of K and V. `indices` (the (rows, k_pad) int32 selection, -1 padded) is counted + exactly when the caller has it in hand.""" + if not _ENABLED: + return + st = _layer(layer_idx) + st.sparse_calls += 1 + st.rows += rows + st.k_pad_max = max(st.k_pad_max, k_pad) + # K and V, fp16, one head_dim vector per kv head per selected position + st.bound_bytes += rows * k_pad * kv_heads * head_dim * 2 * 2 + if decode: + st.decode_calls += 1 + st.decode_rows += rows + if indices is not None: + n = (indices >= 0).sum() + if st.counter is None: + st.counter = torch.zeros((), dtype = torch.long, device = indices.device) + st.counter += n + + +def record_stage(layer_idx: int, nbytes: int): + """One explicit host->VRAM staging copy (phase 2 prefill arena).""" + if not _ENABLED: + return + st = _layer(layer_idx) + st.stage_calls += 1 + st.stage_bytes += nbytes + + +def report(): + global _reported + if not _ENABLED or _reported or not _layers: + return + _reported = True + try: + _report() + except Exception as e: + # atexit runs while CUDA may already be tearing down; a stats dump is never worth + # turning a clean exit into a traceback + print(f"QSA KV offload stats unavailable: {type(e).__name__}: {e}") + + +def _report(): + + MIB = 1024 ** 2 + rowspec = "{:>5} {:>9} {:>11} {:>7} {:>13} {:>13} {:>8} {:>11}" + print() + print("QSA KV offload stats (EXL3_QSA_KVO_STATS)") + print(rowspec.format("layer", "sp.calls", "sparse rows", "K_pad", + "sel/row (act)", "host MiB (bd)", "stages", "staged MiB")) + print("-" * 96) + + tot = _LayerStats() + tot_sel = 0 + any_exact = False + for layer_idx in sorted(_layers): + st = _layers[layer_idx] + sel = int(st.counter.item()) if st.counter is not None else None + per_row = f"{sel / st.rows:.1f}" if sel is not None and st.rows else "-" + if sel is not None: + any_exact = True + tot_sel += sel + print(rowspec.format( + layer_idx, st.sparse_calls, st.rows, st.k_pad_max, per_row, + f"{st.bound_bytes / MIB:.1f}", st.stage_calls, f"{st.stage_bytes / MIB:.1f}")) + tot.sparse_calls += st.sparse_calls + tot.rows += st.rows + tot.k_pad_max = max(tot.k_pad_max, st.k_pad_max) + tot.bound_bytes += st.bound_bytes + tot.decode_calls += st.decode_calls + tot.decode_rows += st.decode_rows + tot.stage_calls += st.stage_calls + tot.stage_bytes += st.stage_bytes + print("-" * 96) + print(rowspec.format( + "all", tot.sparse_calls, tot.rows, tot.k_pad_max, + f"{tot_sel / tot.rows:.1f}" if any_exact and tot.rows else "-", + f"{tot.bound_bytes / MIB:.1f}", tot.stage_calls, f"{tot.stage_bytes / MIB:.1f}")) + + if tot.decode_rows: + print(f"\ndecode: {tot.decode_calls} sparse launches over {tot.decode_rows} rows") + print(f"total host read bound: {tot.bound_bytes / 1024 ** 3:.2f} GiB" + f"{'' if not any_exact else f' (exact selections: {tot_sel})'}") + print() + + +atexit.register(report)