diff --git a/benchmark/bench_hisparse.py b/benchmark/bench_hisparse.py new file mode 100644 index 000000000..d0b855e10 --- /dev/null +++ b/benchmark/bench_hisparse.py @@ -0,0 +1,809 @@ +import argparse +import os +from itertools import product + +import pandas as pd +import torch +import triton +from sgl_kernel import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) + +DEVICE = "xpu" +SEED = 42 + +# Linear-layout MLA item: 512 kv-lora + 64 rope, bf16 (matches the DSA MLA cache). +DTYPE = torch.bfloat16 +ELEM = DTYPE.itemsize +KV_DIM = 576 +LINEAR_ITEM_BYTES = KV_DIM * ELEM + +# Page-padded C4 layout constants (must match c4_layout.hpp). +DSV4_PAGE_SIZE = 64 +DSV4_VALUE_BYTES = 576 +DSV4_SCALE_BYTES = 8 +DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES +DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 +DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE + +VRAM_BUDGET_FRACTION = 0.5 + +QUANTILES = [0.5, 0.25, 0.75] # median, fastest, slowest +RESULT_DIR = "bench_bmg_hisparse_res" + +all_results = [] + + +def _div_up(a, b): + return (a + b - 1) // b + + +def _release(*objs): + """Free device buffers; holding several configs live at once halves bandwidth.""" + for o in objs: + if isinstance(o, (list, dict)): + o.clear() + torch.xpu.synchronize() + torch.xpu.empty_cache() + + +def _gbps(nbytes, ms): + if not nbytes or ms <= 0: + return float("nan") + return nbytes * 1e-9 / (ms * 1e-3) + + +def _dsv4_views(cache): + """Expose a page-padded C4 cache as (value, scale) views, no copy. + + Slicing + ``view`` would fail (the slice is not contiguous, row stride is + DSV4_PAGE_BYTES), so build genuine views with the layout's own strides. + """ + pages = cache.shape[0] + value = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_VALUE_BYTES), + (DSV4_PAGE_BYTES, DSV4_VALUE_BYTES, 1), + ) + scale = cache.as_strided( + (pages, DSV4_PAGE_SIZE, DSV4_SCALE_BYTES), + (DSV4_PAGE_BYTES, DSV4_SCALE_BYTES, 1), + storage_offset=DSV4_SCALE_OFFSET, + ) + return value, scale + + +def _page_split(index): + return index // DSV4_PAGE_SIZE, index % DSV4_PAGE_SIZE + + +def _bench(fn, warmup=10): + for _ in range(warmup): + fn() + torch.xpu.synchronize() + return triton.testing.do_bench(fn, quantiles=QUANTILES) + + +def _bench_cold(make_state, call, reps=20): + warm = make_state() + call(warm) # warm the kernel and the allocator + warm.clear() + torch.xpu.synchronize() + + samples = [] + for _ in range(reps): + state = make_state() + start = torch.xpu.Event(enable_timing=True) + end = torch.xpu.Event(enable_timing=True) + start.record() + call(state) + end.record() + torch.xpu.synchronize() + samples.append(start.elapsed_time(end)) + state.clear() + + samples.sort() + last = len(samples) - 1 + return tuple(samples[min(int(q * len(samples)), last)] for q in QUANTILES) + + +# --------------------------------------------------------------------------- +# swap-in state +# --------------------------------------------------------------------------- + +# Every miss rep reuses the same host cache; only the *device* state has to be +# rebuilt. torch's caching host allocator never returns pinned blocks to the OS, +# so allocating one per rep starved the transfer suite of its 2.3 GB. +_HOST_CACHES = {} + + +def _host_cache(batch_size, num_top_k, hot_buffer_size, is_dsv4): + key = (batch_size, num_top_k, hot_buffer_size, is_dsv4) + cached = _HOST_CACHES.get(key) + if cached is not None: + return cached + + num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 + if is_dsv4: + shape = (_div_up(num_host_items, DSV4_PAGE_SIZE), DSV4_PAGE_BYTES) + cache = torch.empty(shape, dtype=torch.uint8, device="cpu").pin_memory() + cache.fill_(7) + else: + shape = (num_host_items, 1, KV_DIM) + cache = torch.empty(shape, dtype=DTYPE, device="cpu").pin_memory() + cache.fill_(1.0) + _HOST_CACHES[key] = cache + return cache + + +def release_host_caches(): + _release(_HOST_CACHES) + + +def _make_swapin_state(batch_size, num_top_k, hot_buffer_size, is_dsv4, regime, block): + slots_per_req = hot_buffer_size + 1 # +1 reserved newest slot + num_device_items = batch_size * slots_per_req + num_host_items = batch_size * (num_top_k + hot_buffer_size) + 1 + host_cache = _host_cache(batch_size, num_top_k, hot_buffer_size, is_dsv4) + + if is_dsv4: + device_buffer = torch.zeros( + (_div_up(num_device_items, DSV4_PAGE_SIZE), DSV4_PAGE_BYTES), + dtype=torch.uint8, + device=DEVICE, + ) + item_size_bytes = DSV4_ITEM_BYTES + else: + device_buffer = torch.zeros( + (num_device_items, 1, KV_DIM), dtype=DTYPE, device=DEVICE + ) + item_size_bytes = LINEAR_ITEM_BYTES + + # Each request owns a contiguous, disjoint span of device slots. + device_buffer_locs = ( + torch.arange(num_device_items, dtype=torch.int32, device=DEVICE) + .view(batch_size, slots_per_req) + .contiguous() + ) + host_cache_locs = ( + torch.arange(num_host_items, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + + # Request r asks for tokens [base_r, base_r + num_top_k). + base = torch.arange(batch_size, dtype=torch.int32, device=DEVICE).view(-1, 1) * ( + num_top_k + hot_buffer_size + ) + top_k_tokens = ( + base + torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) + ).contiguous() + + device_buffer_tokens = torch.full( + (batch_size, slots_per_req), -1, dtype=torch.int32, device=DEVICE + ) + if regime == "hit": + device_buffer_tokens[:, :num_top_k] = top_k_tokens + else: + device_buffer_tokens[:, :hot_buffer_size] = ( + base + + num_top_k + + torch.arange(hot_buffer_size, dtype=torch.int32, device=DEVICE).view( + 1, -1 + ) + ) + device_buffer_tokens[:, hot_buffer_size] = top_k_tokens[:, -1] + + lru_slots = ( + torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(batch_size, 1) + .contiguous() + ) + torch.xpu.synchronize() + + return { + "top_k_tokens": top_k_tokens, + "device_buffer_tokens": device_buffer_tokens, + "host_cache_locs": host_cache_locs, + "device_buffer_locs": device_buffer_locs, + "host_cache": host_cache, + "device_buffer": device_buffer, + "top_k_device_locs": torch.full_like(top_k_tokens, -1), + "req_pool_indices": torch.arange(batch_size, dtype=torch.int64, device=DEVICE), + # seq_len must exceed hot_buffer_size or the kernel takes the fast path. + "seq_lens": torch.full( + (batch_size,), num_top_k + hot_buffer_size, dtype=torch.int32, device=DEVICE + ), + "lru_slots": lru_slots, + "item_size_bytes": item_size_bytes, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "page_size": DSV4_PAGE_SIZE if is_dsv4 else 1, + "block_size": block, + "num_real_reqs": torch.tensor([batch_size], dtype=torch.int32, device=DEVICE), + } + + +def _swapin(state, is_dsv4): + fn = ( + load_cache_to_device_buffer_dsv4_mla + if is_dsv4 + else load_cache_to_device_buffer_mla + ) + fn(**state) + + +def _torch_swapin(state, is_dsv4): + """Pure-PyTorch equivalent of ``load_cache_to_device_buffer_*_mla``. + + A *timing* reference only -- the kernel's exact miss-to-slot assignment is + internal, so slot ids may differ; accuracy is covered by tests/. The dominant + cost is structural: indexing a pinned CPU tensor needs the gather indices on + the host, forcing a device->host sync per call. The kernel stays on device. + """ + top_k = state["top_k_tokens"] # [B, K] int32 + dbt = state["device_buffer_tokens"] # [B, S] int32 + dbl = state["device_buffer_locs"] # [B, S] int32 + hcl = state["host_cache_locs"] # [B, H] int64 + host_cache = state["host_cache"] + device_buffer = state["device_buffer"] + lru = state["lru_slots"] # [B, hot] int16 + + # Classify each requested token as hit or miss. + eq = top_k.unsqueeze(2) == dbt.unsqueeze(1) # [B, K, S] + is_hit = eq.any(dim=2) # [B, K] + hit_slot = eq.to(torch.uint8).argmax(dim=2) # [B, K] + slot_is_hit = eq.any(dim=1) # [B, S] + + # Pick evict slots in LRU order; the stable sort brings evictable slots + # first while preserving that order. + lru_long = lru.to(torch.int64) + evictable = ~slot_is_hit.gather(1, lru_long) # [B, hot] in LRU order + order = torch.argsort(~evictable, dim=1, stable=True) + evict_slots = lru_long.gather(1, order) # [B, hot] + + # Rank each miss, and match it to the evict slot of the same rank. + miss = ~is_hit + miss_rank = miss.cumsum(dim=1) - 1 # [B, K] + evict_pick = evict_slots.gather(1, miss_rank.clamp_(min=0)) + assigned = torch.where(miss, evict_pick, hit_slot) # [B, K] + + # Stream the misses in. + b_idx, k_idx = miss.nonzero(as_tuple=True) + if b_idx.numel(): + token = top_k[b_idx, k_idx].to(torch.int64) + host_loc = hcl[b_idx, token] + dev_loc = dbl[b_idx, assigned[b_idx, k_idx]].to(torch.int64) + # Pinned host cache must be indexed with CPU indices -> D2H sync. + host_loc_cpu = host_loc.cpu() + if is_dsv4: + h_val, h_scale = _dsv4_views(host_cache) + d_val, d_scale = _dsv4_views(device_buffer) + hp, ho = _page_split(host_loc_cpu) + dp, do = _page_split(dev_loc) + d_val[dp, do] = h_val[hp, ho].to(DEVICE, non_blocking=True) + d_scale[dp, do] = h_scale[hp, ho].to(DEVICE, non_blocking=True) + else: + staged = host_cache[host_loc_cpu].to(DEVICE, non_blocking=True) + device_buffer[dev_loc] = staged + + dbt[b_idx, assigned[b_idx, k_idx]] = top_k[b_idx, k_idx] + + state["top_k_device_locs"].copy_(dbl.gather(1, assigned)) + # Touched slots become most-recently-used; untouched keep LRU order first. + touched = torch.zeros_like(slot_is_hit) + touched.scatter_(1, assigned, True) + key = touched.gather(1, lru_long).to(torch.uint8) + lru.copy_(lru_long.gather(1, torch.argsort(key, dim=1, stable=True)).to(lru.dtype)) + + +# --------------------------------------------------------------------------- +# transfer state +# --------------------------------------------------------------------------- + +# One live state at a time. triton iterates providers innermost, so a shape's +# block sizes arrive back to back and reuse cuts allocation work 3x. +_TRANSFER_STATE = {"key": None, "value": None} + + +def _transfer_state(num_items, num_layers): + key = (num_items, num_layers) + if _TRANSFER_STATE["key"] == key: + return _TRANSFER_STATE["value"] + + release_transfer_state() + pages = _div_up(num_items, DSV4_PAGE_SIZE) + srcs = [ + torch.full((pages, DSV4_PAGE_BYTES), 3, dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + dsts = [ + torch.zeros((pages, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + for _ in range(num_layers) + ] + value = ( + srcs, + dsts, + # Raw addresses only -- srcs/dsts must stay alive alongside them. + torch.tensor([t.data_ptr() for t in srcs], dtype=torch.uint64, device=DEVICE), + torch.tensor([t.data_ptr() for t in dsts], dtype=torch.uint64, device=DEVICE), + torch.arange(num_items, dtype=torch.int64, device=DEVICE), + ) + torch.xpu.synchronize() + _TRANSFER_STATE.update(key=key, value=value) + return value + + +def _torch_transfer(srcs, dsts, src_indices, dst_indices): + """Pure-PyTorch equivalent of ``transfer_cache_dsv4_mla``. + + The kernel walks all layers in one launch; eager needs an indexed copy per + layer, which is the cost this comparison isolates. + """ + sp, so = _page_split(src_indices) + dp, do = _page_split(dst_indices) + for src, dst in zip(srcs, dsts): + s_val, s_scale = _dsv4_views(src) + d_val, d_scale = _dsv4_views(dst) + d_val[dp, do] = s_val[sp, so] + d_scale[dp, do] = s_scale[sp, so] + + +def release_transfer_state(): + value = _TRANSFER_STATE["value"] + if value is not None: + _release(value[0], value[1]) + _TRANSFER_STATE.update(key=None, value=None) + + +def _vram_budget(): + return int(torch.xpu.get_device_properties(0).total_memory * VRAM_BUDGET_FRACTION) + + +def _transfer_footprint(num_items, num_layers): + return 2 * num_layers * _div_up(num_items, DSV4_PAGE_SIZE) * DSV4_PAGE_BYTES + + +# --------------------------------------------------------------------------- +# swap-in benchmark +# --------------------------------------------------------------------------- + + +def swapin_providers(with_torch): + impls = ("sglang", "torch") if with_torch else ("sglang",) + return [ + f"{impl}-{layout}-{regime}" + for layout in ("linear", "dsv4") + for regime in ("hit", "miss") + for impl in impls + ] + + +_STYLES = [ + ("blue", "-"), + ("blue", "--"), + ("green", "-"), + ("green", "--"), + ("red", "-"), + ("red", "--"), + ("orange", "-"), + ("orange", "--"), +] + + +def benchmark_swapin( + batch_size, num_top_k, hot_buffer_size, provider, block_size, reps +): + impl, layout, regime = provider.split("-") + print( + f"benchmark load_cache_to_device_buffer {provider} batch_size={batch_size} " + f"num_top_k={num_top_k} hot_buffer_size={hot_buffer_size} " + f"block_size={block_size}" + ) + torch.xpu.manual_seed_all(SEED) + + is_dsv4 = layout == "dsv4" + item_bytes = DSV4_ITEM_BYTES if is_dsv4 else LINEAR_ITEM_BYTES + run = _torch_swapin if impl == "torch" else _swapin + + if regime == "miss": + # A miss reads one item from the host and writes one to the device buffer. + nbytes = batch_size * num_top_k * item_bytes * 2 + ms, fast_ms, slow_ms = _bench_cold( + lambda: _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "miss", block_size + ), + lambda st: run(st, is_dsv4), + reps=reps, + ) + else: + # Hits move no bytes: index resolution plus LRU bookkeeping only. + nbytes = 0 + state = _make_swapin_state( + batch_size, num_top_k, hot_buffer_size, is_dsv4, "hit", block_size + ) + ms, fast_ms, slow_ms = _bench(lambda: run(state, is_dsv4)) + state.clear() + _release() + + all_results.append( + { + "op": f"load_cache_to_device_buffer [{regime}]", + "impl": impl, + "layout": layout, + "batch_size": batch_size, + "num_top_k": num_top_k, + "hot_buffer_size": hot_buffer_size, + "block_size": block_size, + "us (median)": round(ms * 1000, 2), + "GB/s (median)": round(_gbps(nbytes, ms), 2), + "GB/s (min)": round(_gbps(nbytes, slow_ms), 2), + "GB/s (max)": round(_gbps(nbytes, fast_ms), 2), + } + ) + return ms * 1000, fast_ms * 1000, slow_ms * 1000 + + +def swapin_mark(configs, with_torch=False): + providers = swapin_providers(with_torch) + return triton.testing.Mark( + benchmark_swapin, + triton.testing.Benchmark( + x_names=["batch_size", "num_top_k", "hot_buffer_size"], + x_vals=configs, + line_arg="provider", + line_vals=providers, + line_names=[p.replace("-", " ") for p in providers], + styles=_STYLES[: len(providers)], + ylabel="us", + plot_name="hisparse-load-cache-to-device-buffer", + args={}, + ), + ) + + +# --------------------------------------------------------------------------- +# transfer benchmark +# --------------------------------------------------------------------------- + + +def benchmark_transfer(num_items, num_layers, provider): + is_torch = provider == "torch" + block = 0 if is_torch else int(provider) + print( + f"benchmark transfer_cache_dsv4_mla {provider} num_items={num_items} " + f"num_layers={num_layers}" + ) + torch.xpu.manual_seed_all(SEED) + + footprint = _transfer_footprint(num_items, num_layers) + budget = _vram_budget() + if footprint > budget: + # Say what was skipped: a silent drop reads as "covered" in the table. + print( + f" SKIPPED: needs {footprint / 1024**3:.1f} GiB of VRAM, " + f"budget is {budget / 1024**3:.1f} GiB" + ) + nan = float("nan") + return nan, nan, nan + + srcs, dsts, src_ptrs, dst_ptrs, idx = _transfer_state(num_items, num_layers) + nbytes = num_items * num_layers * DSV4_ITEM_BYTES * 2 # read + write + + if is_torch: + fn = lambda: _torch_transfer(srcs, dsts, idx, idx) # noqa: E731 + else: + fn = lambda: transfer_cache_dsv4_mla( # noqa: E731 + src_ptrs=src_ptrs, + dst_ptrs=dst_ptrs, + src_indices=idx, + dst_indices=idx, + block_size=block, + ) + ms, fast_ms, slow_ms = _bench(fn) + + all_results.append( + { + "op": "transfer_cache_dsv4_mla", + "impl": "torch" if is_torch else "sglang", + "num_items": num_items, + "num_layers": num_layers, + "block_size": block, + "us (median)": round(ms * 1000, 2), + "GB/s (median)": round(_gbps(nbytes, ms), 2), + "GB/s (min)": round(_gbps(nbytes, slow_ms), 2), + "GB/s (max)": round(_gbps(nbytes, fast_ms), 2), + } + ) + return _gbps(nbytes, ms), _gbps(nbytes, slow_ms), _gbps(nbytes, fast_ms) + + +def transfer_mark(configs, block_sizes, with_torch=False): + providers = [str(b) for b in block_sizes] + (["torch"] if with_torch else []) + return triton.testing.Mark( + benchmark_transfer, + triton.testing.Benchmark( + x_names=["num_items", "num_layers"], + x_vals=configs, + line_arg="provider", + line_vals=providers, + line_names=[ + "torch eager" if p == "torch" else f"block={p}" for p in providers + ], + styles=_STYLES[: len(providers)], + ylabel="GB/s", + plot_name="hisparse-transfer-cache-dsv4-mla", + args={}, + ), + ) + + +# --------------------------------------------------------------------------- +# speedup analysis +# --------------------------------------------------------------------------- + + +def speedup_analysis(df, transfer_block_size): + if "impl" not in df.columns or "torch" not in set(df["impl"]): + return + + print("\n" + "=" * 80) + print("SPEEDUP ANALYSIS (torch eager / sglang, higher = kernel is faster)") + print("=" * 80) + + for op, sub in df.groupby("op", sort=False): + sgl = sub[sub["impl"] == "sglang"] + ref = sub[sub["impl"] == "torch"] + if sgl.empty or ref.empty: + continue + if op == "transfer_cache_dsv4_mla": + sgl = sgl[sgl["block_size"] == transfer_block_size] + group, label = "num_layers", f"block={transfer_block_size}" + else: + group, label = "layout", "all layouts" + + keys = [ + c + for c in ("layout", "batch_size", "num_top_k", "num_items", "num_layers") + if c in sub.columns + ] + merged = sgl.merge(ref, on=keys, how="inner", suffixes=("_sgl", "_torch")) + if merged.empty: + continue + merged["speedup"] = merged["us (median)_torch"] / merged["us (median)_sgl"] + + print(f"\n### {op} ({label})\n") + print( + f" overall: avg={merged['speedup'].mean():.2f}x " + f"max={merged['speedup'].max():.2f}x " + f"min={merged['speedup'].min():.2f}x" + ) + print(f"\n by {group}:") + for value, rows in merged.groupby(group, sort=True): + print( + f" {value!s:>8}: avg={rows['speedup'].mean():6.2f}x " + f"max={rows['speedup'].max():6.2f}x " + f"min={rows['speedup'].min():6.2f}x" + ) + + +def check_nonzero(): + probe = torch.tensor([[0, 1, 0], [1, 0, 1]], dtype=torch.bool, device=DEVICE) + rows, cols = probe.nonzero(as_tuple=True) + ok = rows.tolist() == [0, 1, 1] and cols.tolist() == [1, 0, 2] + if not ok: + print( + "WARNING: torch.nonzero is wrong on this stack " + f"(got rows={rows.tolist()} cols={cols.tolist()}, " + "expected rows=[0, 1, 1] cols=[1, 0, 2]). " + "The torch-eager baseline is unreliable; speedups will be overstated." + ) + return ok + + +# --------------------------------------------------------------------------- +# regression tracking +# --------------------------------------------------------------------------- + +# Columns identifying one benchmarked configuration, for joining two runs. +_KEY_COLUMNS = [ + "op", + "impl", + "layout", + "batch_size", + "num_top_k", + "hot_buffer_size", + "num_items", + "num_layers", + "block_size", +] + + +def compare_results(df): + os.makedirs(RESULT_DIR, exist_ok=True) + previous_csv = os.path.join(RESULT_DIR, "previous.csv") + current_csv = os.path.join(RESULT_DIR, "current.csv") + + df.to_csv(current_csv, index=False) + print(f"\nCurrent results saved to: {current_csv}") + + if not os.path.exists(previous_csv): + print(f"No {previous_csv} found, nothing to compare against.") + print("Tip: copy current.csv to previous.csv to set a baseline.") + return + + try: + prev = pd.read_csv(previous_csv) + except Exception as e: # noqa: BLE001 -- a stale CSV must not kill the run + print(f"Error loading {previous_csv}: {e}") + return + print(f"Loaded previous results from: {previous_csv}") + + keys = [c for c in _KEY_COLUMNS if c in df.columns and c in prev.columns] + merged = df.merge(prev, on=keys, how="inner", suffixes=("", "_prev")) + if merged.empty: + print("No configurations in common with the previous run.") + return + + merged["delta %"] = ( + (merged["us (median)"] - merged["us (median)_prev"]) + / merged["us (median)_prev"] + * 100 + ).round(1) + report = merged[keys + ["us (median)_prev", "us (median)", "delta %"]].rename( + columns={"us (median)_prev": "previous us", "us (median)": "current us"} + ) + + print("\n" + "=" * 80) + print("REGRESSION vs previous.csv (positive delta = slower than before)") + print("=" * 80 + "\n") + print(report.dropna(axis=1, how="all").to_markdown(index=False)) + print( + f"\nWorst: {report['delta %'].max():+.1f}% " + f"Best: {report['delta %'].min():+.1f}% " + f"Mean: {report['delta %'].mean():+.1f}%" + ) + + +# --------------------------------------------------------------------------- +# main +# --------------------------------------------------------------------------- + + +def parse_args(): + p = argparse.ArgumentParser(description="HiSparse KV-cache benchmark for XPU") + p.add_argument( + "--suite", + nargs="+", + choices=["swapin", "transfer"], + default=["swapin", "transfer"], + help="Which ops to benchmark (default: both)", + ) + p.add_argument( + "--batch-sizes", + type=int, + nargs="+", + default=[1, 8, 32, 128], + metavar="B", + help="Swap-in batch sizes (default: 1 8 32 128)", + ) + p.add_argument( + "--top-k", + type=int, + nargs="+", + default=[64, 256], + metavar="K", + help="Swap-in num_top_k values; hot_buffer_size is set equal to each " + "(default: 64 256)", + ) + p.add_argument( + "--swapin-block-size", + type=int, + default=256, + help="block_size for the swap-in op (default: 256)", + ) + p.add_argument( + "--reps", + type=int, + default=20, + help="Reps for the cold (miss) swap-in regime (default: 20)", + ) + p.add_argument( + "--num-items", + type=int, + nargs="+", + default=[64, 512, 4096, 32768], + metavar="N", + help="Transfer item counts (default: 64 512 4096 32768)", + ) + p.add_argument( + "--num-layers", + type=int, + nargs="+", + default=[1, 8, 61], + metavar="L", + help="Transfer layer counts (default: 1 8 61)", + ) + p.add_argument( + "--block-sizes", + type=int, + nargs="+", + default=[256, 512, 1024], + choices=[256, 512, 1024], + metavar="B", + help="block_size values for the transfer op (default: 256 512 1024)", + ) + p.add_argument( + "--with-torch", + action="store_true", + help="Add a torch-eager provider and print a speedup analysis. Off by " + "default: the eager path allocates heavily and its allocator traffic " + "widens the kernels' own measured spread", + ) + p.add_argument( + "--save-csv", + action="store_true", + help="Save results to CSV and diff against the previous run " + "(default: only print results)", + ) + return p.parse_args() + + +def main(): + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + print("ERROR: no XPU device available.") + raise SystemExit(1) + + args = parse_args() + props = torch.xpu.get_device_properties(0) + print(f"Device : {torch.xpu.get_device_name(0)}") + print(f"VRAM : {props.total_memory / 1024**3:.1f} GiB") + print(f"dtype : {DTYPE}") + print(f"suites : {args.suite}") + print(f"torch : {'yes' if args.with_torch else 'no (kernels only)'}\n") + + if args.with_torch: + check_nonzero() + + all_results.clear() + + if "swapin" in args.suite: + configs = [(b, k, k) for b, k in product(args.batch_sizes, args.top_k)] + swapin_mark(configs, args.with_torch).run( + print_data=True, + show_plots=False, + save_path=None, + block_size=args.swapin_block_size, + reps=args.reps, + ) + release_host_caches() + + if "transfer" in args.suite: + configs = list(product(args.num_items, args.num_layers)) + transfer_mark(configs, args.block_sizes, args.with_torch).run( + print_data=True, show_plots=False, save_path=None + ) + release_transfer_state() + + if not all_results: + print("No results collected.") + return + + print("\nBenchmark finished!") + df = pd.DataFrame(all_results) + for op, sub in df.groupby("op", sort=False): + sub = sub.dropna(axis=1, how="all").reset_index(drop=True) + print(f"\n### {op}\n") + print(sub.drop(columns=["op"]).to_markdown(index=False)) + + speedup_analysis(df, max(args.block_sizes)) + + if args.save_csv: + compare_results(df) + + +if __name__ == "__main__": + main() diff --git a/cmake/BuildFlags.cmake b/cmake/BuildFlags.cmake index ed116fcca..a24250689 100644 --- a/cmake/BuildFlags.cmake +++ b/cmake/BuildFlags.cmake @@ -77,6 +77,10 @@ if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -ftemplate-backtrace-limit=0) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fno-sycl-unnamed-lambda) + # ATen/torch headers included in device-compiled FMHA runners require C++20. + # -sycl-std sets the SYCL spec version, not the C++ language standard, so the + # C++ standard must be set explicitly to match SYCL_HOST_FLAGS (-std=c++20). + set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -std=c++20) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -sycl-std=2020) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fhonor-nans) set(SYCL_KERNEL_OPTIONS ${SYCL_KERNEL_OPTIONS} -fhonor-infinities) diff --git a/include/sgl_kernel/hisparse/c4_layout.hpp b/include/sgl_kernel/hisparse/c4_layout.hpp new file mode 100644 index 000000000..0546605af --- /dev/null +++ b/include/sgl_kernel/hisparse/c4_layout.hpp @@ -0,0 +1,75 @@ +/** + * HiSparse C4 paged-cache layout helpers. + * Paged C4 cache layout (per page of kPageSize tokens): + * VALUE 0, VALUE 1, ..., VALUE 63, (kValueBytes each) + * SCALE 0, SCALE 1, ..., SCALE 63, (kScaleBytes each) + * [padding to align the page to a 576-byte boundary] + * + * FlashMLA requires each page to be aligned to 576 bytes. + */ + +#pragma once + +#include +#include + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +inline constexpr int kSubGroupSize = 32; + +// C4 paged layout constants (must match the model-side cache layout exactly). +inline constexpr int64_t kPageSize = 64; +inline constexpr int64_t kPageBits = 6; // log2(kPageSize) +inline constexpr int64_t kValueBytes = 576; +inline constexpr int64_t kScaleBytes = 8; +inline constexpr int64_t kItemBytes = kValueBytes + kScaleBytes; + +// div_ceil(kItemBytes * kPageSize, 576) * 576 -> page byte stride. +inline constexpr int64_t kPageBytes = ((kItemBytes * kPageSize + 576 - 1) / 576) * 576; +inline constexpr int64_t kScaleOffset = kValueBytes * kPageSize; + +// int64-word counts for the strided copy loops. +inline constexpr int kValueWords = static_cast(kValueBytes / 8); // 72 +inline constexpr int kScaleWords = static_cast(kScaleBytes / 8); // 1 + +static_assert(kValueBytes % 8 == 0, "kValueBytes must be a multiple of 8"); +static_assert(kScaleBytes % 8 == 0, "kScaleBytes must be a multiple of 8"); +static_assert((int64_t(1) << kPageBits) == kPageSize, "kPageBits must equal log2(kPageSize)"); + +struct PointerInfo { + int64_t* value_ptr; + int64_t* scale_ptr; +}; + +// Resolve the value/scale int64 pointers for a single token slot in a paged +// C4 cache. `index` is the logical token slot; the layout is page-padded. +inline PointerInfo get_pointer_paged(void* cache, int32_t index) { + const int32_t page_num = index >> kPageBits; + const int32_t page_offset = index & (kPageSize - 1); + char* base = static_cast(cache) + static_cast(page_num) * kPageBytes; + char* value_ptr = base + static_cast(page_offset) * kValueBytes; + char* scale_ptr = base + kScaleOffset + static_cast(page_offset) * kScaleBytes; + return {reinterpret_cast(value_ptr), reinterpret_cast(scale_ptr)}; +} + +// Copy one C4 item (value + scale) between page-padded caches, cooperatively +// across a sub-group. `lane_id`/`sg_size` are the sub-group local id and width; +// the strided loops make this correct for any Intel SIMD width (8/16/32). +inline void +transfer_item(int lane_id, int sg_size, void* dst_cache, void* src_cache, int32_t dst_index, int32_t src_index) { + const PointerInfo dst = get_pointer_paged(dst_cache, dst_index); + const PointerInfo src = get_pointer_paged(src_cache, src_index); + + for (int j = lane_id; j < kValueWords; j += sg_size) { + dst.value_ptr[j] = src.value_ptr[j]; + } + for (int j = lane_id; j < kScaleWords; j += sg_size) { + dst.scale_ptr[j] = src.scale_ptr[j]; + } +} + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp b/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp new file mode 100644 index 000000000..aa62e8e22 --- /dev/null +++ b/include/sgl_kernel/hisparse/load_cache_to_device_buffer.hpp @@ -0,0 +1,449 @@ +/** + * HiSparse: load_cache_to_device_buffer SYCL kernel (Intel XPU). + */ + +#pragma once + +#include +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +static constexpr int32_t kTokenHit = static_cast(0xFFFFFFFF); // -1 sentinel "already resident" +static constexpr int32_t kHashEmpty = -1; + +// Knuth multiplicative hash into an open-addressing table of size hash_size. +// hash_mask is hash_size-1 for power-of-two sizes (else 0), avoiding a modulo. +inline int hash_slot(int32_t key, int hash_size, int hash_mask) { + const uint32_t h = static_cast(key) * 2654435761u; + return hash_mask != 0 ? static_cast(h & static_cast(hash_mask)) + : static_cast(h % static_cast(hash_size)); +} + +// Linear-probe step: (slot + 1) % hash_size without the division. +inline int hash_probe_next(int slot, int hash_size) { + const int next = slot + 1; + return next == hash_size ? 0 : next; +} + +// Cooperative linear (non-paged) item copy across a sub-group, for the generic +// miss path where device + host are both linear with stride item_size_bytes. +inline void transfer_item_linear(int lane_id, int sg_size, const void* src, void* dst, int64_t item_size_bytes) { + const int64_t nwords = item_size_bytes / 8; + const int64_t* s = static_cast(src); + int64_t* d = static_cast(dst); + for (int64_t j = lane_id; j < nwords; j += sg_size) { + d[j] = s[j]; + } + const int64_t tail_start = nwords * 8; + const char* sc = static_cast(src) + tail_start; + char* dc = static_cast(dst) + tail_start; + for (int64_t j = lane_id; j < item_size_bytes - tail_start; j += sg_size) { + dc[j] = sc[j]; + } +} + +// Local-memory atomic CAS returning the previous value. +inline int32_t atomic_cas_local(int32_t* addr, int32_t compare, int32_t val) { + ::sycl::atomic_ref< + int32_t, + ::sycl::memory_order::relaxed, + ::sycl::memory_scope::work_group, + ::sycl::access::address_space::local_space> + ref(*addr); + int32_t expected = compare; + ref.compare_exchange_strong(expected, val); + // On success `expected` is unchanged (== compare); on failure it holds the + // current value. + return expected; +} + +// Single-sub-group inclusive prefix scan over the local-memory window +// [offset, offset+sg_size), threading a running accumulator. +inline int sub_group_inclusive_scan( + const ::sycl::sub_group& sg, int32_t* s_data, int lane_id, int sg_size, int offset, int count, int accumulator) { + const int idx = lane_id + offset; + int val = (idx < count) ? s_data[idx] : 0; + val = ::sycl::inclusive_scan_over_group(sg, val, ::sycl::plus()); + val += accumulator; + if (idx < count) { + s_data[idx] = val; + } + accumulator = ::sycl::group_broadcast(sg, val, sg_size - 1); + return accumulator; +} + +// Local-memory layout: an int32_t region followed by an int16_t region starting +// at int32 slot total_int32, so it inherits 4-byte alignment. num_top_k and +// hot_buffer_size are runtime values, so the host computes the layout and hands +// the offsets to the kernel. +struct SmemLayout { + int hash_size; + int hash_mask; // hash_size-1 if a power of two, else 0 + int num_buffer_chunks; + int num_token_chunks; + int total_int32; + int total_int16; + int total_int32_slots; // allocation size of the int32 local_accessor + + static SmemLayout make(int num_top_k, int hot_buffer_size) { + SmemLayout l{}; + l.hash_size = num_top_k * 2; + l.hash_mask = (l.hash_size & (l.hash_size - 1)) == 0 ? l.hash_size - 1 : 0; + l.num_buffer_chunks = (hot_buffer_size + kSubGroupSize - 1) / kSubGroupSize; + l.num_token_chunks = (num_top_k + kSubGroupSize - 1) / kSubGroupSize; + // int32 region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + // + {total_hits, newest_hit} + l.total_int32 = num_top_k + (l.num_buffer_chunks + 1) + (l.num_buffer_chunks + 1) + l.hash_size + 2; + // int16 region: lru_slots_out + hash_vals + l.total_int16 = hot_buffer_size + l.hash_size; + l.total_int32_slots = + l.total_int32 + static_cast((l.total_int16 * sizeof(int16_t) + sizeof(int32_t) - 1) / sizeof(int32_t)); + return l; + } + + size_t bytes() const { + return static_cast(total_int32_slots) * sizeof(int32_t); + } +}; + +// seq_lens and req_pool_indices are int32 or int64 depending on the caller. Each +// is read once per work-group, so a runtime dtype branch beats instantiating the +// kernel for both index types. +inline int64_t load_index(const void* base, bool is_i64, int bid) { + return is_i64 ? static_cast(base)[bid] : static_cast(static_cast(base)[bid]); +} + +// IsMLA / IsDsv4Layout stay compile-time because they select the inner copy loop. +// Aggregate-initialized so the launcher can name the fields: the list is long +// enough that positional arguments are easy to transpose silently. +template +struct LoadCacheToDeviceBufferKernel { + static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA)."); + + const int32_t* top_k_tokens_; + int32_t* device_buffer_tokens_; + const int64_t* host_cache_locs_; + const int32_t* device_buffer_locs_; + const void* host_cache_k_; + const void* host_cache_v_; + void* device_buffer_k_; + void* device_buffer_v_; + int32_t* top_k_device_locs_; + const void* req_pool_indices_; + const void* seq_lens_; + int16_t* lru_slots_; + const int32_t* num_real_reqs_; + bool req_pool_indices_is_i64_; + bool seq_lens_is_i64_; + int64_t buffer_stride_0_; + int64_t host_stride_; + int64_t lru_slot_stride_0_; + int64_t top_k_tokens_stride_; + int64_t top_k_device_locs_stride_; + int64_t item_size_bytes_; + // Runtime shape / layout, precomputed on the host (see SmemLayout). + int block_size_; + int num_sub_groups_; + int num_top_k_; + int hot_buffer_size_; + int hash_size_; + int hash_mask_; + int num_buffer_chunks_; + int num_token_chunks_; + int iters_per_sg_buffer_; + int iters_per_sg_token_; + int total_int32_; + ::sycl::local_accessor smem_; + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const int bid = static_cast(item.get_group(0)); + const int tid = static_cast(item.get_local_id(0)); + int32_t* req_top_k_device_locs = top_k_device_locs_ + bid * top_k_device_locs_stride_; + + // A graph-captured batch is padded to the captured size. Keep padded output + // rows invalid without a separate fill kernel. + if (bid >= num_real_reqs_[0]) { + for (int i = tid; i < num_top_k_; i += block_size_) { + req_top_k_device_locs[i] = -1; + } + return; + } + + const ::sycl::sub_group sg = item.get_sub_group(); + const int sg_id = static_cast(sg.get_group_linear_id()); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + const int64_t rid = load_index(req_pool_indices_, req_pool_indices_is_i64_, bid); + const int64_t seq_len = load_index(seq_lens_, seq_lens_is_i64_, bid); + + // Per-request base offsets. + const int32_t* req_top_k_tokens = top_k_tokens_ + bid * top_k_tokens_stride_; + + const int64_t buffer_offset = rid * buffer_stride_0_; + int32_t* req_device_buffer_tokens = device_buffer_tokens_ + buffer_offset; + const int32_t* req_device_buffer_locs = device_buffer_locs_ + buffer_offset; + const int64_t* req_host_cache_locs = host_cache_locs_ + rid * host_stride_; + int16_t* req_lru_slots = lru_slots_ + rid * lru_slot_stride_0_; + + // Fast path: short sequences have all tokens resident in device-buffer order. + if (seq_len <= hot_buffer_size_) { + const int count = (seq_len < num_top_k_) ? static_cast(seq_len) : num_top_k_; + for (int i = tid; i < num_top_k_; i += block_size_) { + int32_t device_loc = -1; + if (i < count) { + const int32_t token_pos = req_top_k_tokens[i]; + if (token_pos >= 0) { + device_loc = req_device_buffer_locs[token_pos]; + } + } + req_top_k_device_locs[i] = device_loc; + } + return; + } + + // Scratch is one int32_t accessor; the int16 region starts at slot + // total_int32_, so both regions stay 4-byte aligned. + int32_t* smem_i32 = &smem_[0]; + int32_t* s_top_k_tokens = smem_i32; // num_top_k + int32_t* s_chunk_offset = s_top_k_tokens + num_top_k_; // num_buffer_chunks + 1 + int32_t* s_evict_chunk_offset = s_chunk_offset + (num_buffer_chunks_ + 1); // num_buffer_chunks + 1 + int32_t* s_hash_keys = s_evict_chunk_offset + (num_buffer_chunks_ + 1); // hash_size + int32_t* s_total_hits_ptr = s_hash_keys + hash_size_; // 1 + int32_t* s_newest_hit_ptr = s_hash_keys + hash_size_ + 1; // 1 + + int16_t* smem_i16 = reinterpret_cast(smem_i32 + total_int32_); + int16_t* s_lru_slots_out = smem_i16; // hot_buffer_size + int16_t* s_hash_vals = s_lru_slots_out + hot_buffer_size_; // hash_size + + // Initialize counters, hash table, and prefix-sum offsets. + if (tid == 0) { + *s_total_hits_ptr = 0; + *s_newest_hit_ptr = 0; + } + for (int i = tid; i < hash_size_; i += block_size_) { + s_hash_keys[i] = kHashEmpty; + } + for (int i = tid; i < num_buffer_chunks_ + 1; i += block_size_) { + s_chunk_offset[i] = 0; + s_evict_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + const int newest_slot = hot_buffer_size_; + const int32_t newest_token = static_cast(seq_len - 1); + + // Insert top-k token positions into the local-memory hash table. + for (int i = tid; i < num_top_k_; i += block_size_) { + int32_t token_idx = req_top_k_tokens[i]; + if (token_idx == newest_token) { + // The latest token lives at newest_slot, outside LRU tracking: bind it + // directly and mark it a hit. + s_top_k_tokens[i] = kTokenHit; + req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot]; + *s_newest_hit_ptr = 1; + } else { + int slot = hash_slot(token_idx, hash_size_, hash_mask_); + while (true) { + int32_t old = atomic_cas_local(&s_hash_keys[slot], kHashEmpty, token_idx); + if (old == kHashEmpty || old == token_idx) { + s_hash_vals[slot] = static_cast(i); + break; + } + slot = hash_probe_next(slot, hash_size_); + } + s_top_k_tokens[i] = token_idx; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over hot-buffer slots: classify hits vs evictables and compact them. + int total_hit_count = 0; + int total_evict_count = 0; + for (int iter = 0; iter < iters_per_sg_buffer_; iter++) { + const int chunk_idx = sg_id + iter * num_sub_groups_; + const bool has_valid_chunk = chunk_idx < num_buffer_chunks_; + + const int slot_idx = chunk_idx * kSubGroupSize + lane_id; + const bool has_valid_slot = has_valid_chunk && (slot_idx < hot_buffer_size_); + const int16_t buf_slot = has_valid_slot ? req_lru_slots[slot_idx] : static_cast(-1); + int32_t my_buffer_token = (buf_slot >= 0) ? req_device_buffer_tokens[buf_slot] : -1; + int my_found_top_k_idx = -1; + if (my_buffer_token >= 0) { + int h = hash_slot(my_buffer_token, hash_size_, hash_mask_); + while (true) { + int32_t k = s_hash_keys[h]; + if (k == my_buffer_token) { + my_found_top_k_idx = static_cast(s_hash_vals[h]); + break; + } + if (k == kHashEmpty) break; + h = hash_probe_next(h, hash_size_); + } + } + const bool is_hit = my_found_top_k_idx >= 0; + const bool is_evictable = has_valid_slot && !is_hit; + + // Record hits: bind the top-k index to this resident slot's device loc. + if (is_hit) { + s_top_k_tokens[my_found_top_k_idx] = kTokenHit; + req_top_k_device_locs[my_found_top_k_idx] = req_device_buffer_locs[buf_slot]; + } + + int local_hit_offset = 0; + int local_evict_offset = 0; + if (has_valid_chunk) { + local_hit_offset = ::sycl::exclusive_scan_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + local_evict_offset = ::sycl::exclusive_scan_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + const int sg_hits = ::sycl::reduce_over_group(sg, is_hit ? 1 : 0, ::sycl::plus()); + const int sg_evicts = ::sycl::reduce_over_group(sg, is_evictable ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = sg_hits; + s_evict_chunk_offset[chunk_idx + 1] = sg_evicts; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (sg_id == 0) { + // Bound the scan window to num_sub_groups lanes: only that many entries + // were written this iteration, and letting the remaining lanes join with + // the wide count would fold stale values from earlier iterations into the + // accumulator and write it into slots future iterations read. + const int scan_count = ::std::min(chunk_idx + 1 + num_sub_groups_, num_buffer_chunks_ + 1); + total_hit_count = + sub_group_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_hit_count); + total_evict_count = sub_group_inclusive_scan( + sg, s_evict_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_evict_count); + if (tid == 0) { + *s_total_hits_ptr = total_hit_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Hits grow forward from index 0. + if (is_hit) { + int hit_offset = s_chunk_offset[chunk_idx] + local_hit_offset; + s_lru_slots_out[hit_offset] = buf_slot; + } + // Evictables grow backward from hot_buffer_size - 1. + if (is_evictable) { + int evict_offset = s_evict_chunk_offset[chunk_idx] + local_evict_offset; + s_lru_slots_out[hot_buffer_size_ - 1 - evict_offset] = buf_slot; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + // Reset offsets for the miss-counting phase (num_token_chunks + 1 entries). + for (int i = tid; i < num_token_chunks_ + 1; i += block_size_) { + s_chunk_offset[i] = 0; + } + item.barrier(::sycl::access::fence_space::local_space); + + // Pass over top-k tokens: identify misses and assign them evictable slots. + int total_misses = 0; + for (int iter = 0; iter < iters_per_sg_token_; iter++) { + const int chunk_idx = sg_id + iter * num_sub_groups_; + const bool has_valid_chunk = chunk_idx < num_token_chunks_; + + const int chunk_token_start = chunk_idx * kSubGroupSize; + const int my_token_idx = chunk_token_start + lane_id; + const bool has_valid_token = has_valid_chunk && (my_token_idx < num_top_k_); + + int32_t my_token = 0; + bool is_miss = false; + int local_miss_offset = 0; + + if (has_valid_token) { + is_miss = s_top_k_tokens[my_token_idx] != kTokenHit; + if (is_miss) { + my_token = s_top_k_tokens[my_token_idx]; + } + } + + if (has_valid_chunk) { + local_miss_offset = ::sycl::exclusive_scan_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + const int sg_miss_count = ::sycl::reduce_over_group(sg, is_miss ? 1 : 0, ::sycl::plus()); + if (lane_id == 0) { + s_chunk_offset[chunk_idx + 1] = sg_miss_count; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + if (sg_id == 0) { + // Same bounded window as the buffer pass above. + const int scan_count = ::std::min(chunk_idx + 1 + num_sub_groups_, num_token_chunks_ + 1); + total_misses = + sub_group_inclusive_scan(sg, s_chunk_offset, lane_id, sg_size, chunk_idx + 1, scan_count, total_misses); + } + item.barrier(::sycl::access::fence_space::local_space); + + if (is_miss) { + int miss_offset = s_chunk_offset[chunk_idx] + local_miss_offset; + int16_t evict_slot = s_lru_slots_out[hot_buffer_size_ - 1 - miss_offset]; + // Reuse s_top_k_tokens as miss scratch: miss_offset < my_token_idx always + // holds (hits are skipped), so compacted writes never overrun pending reads. + s_top_k_tokens[miss_offset] = my_token; + req_top_k_device_locs[my_token_idx] = req_device_buffer_locs[evict_slot]; + req_device_buffer_tokens[evict_slot] = my_token; + } + } + item.barrier(::sycl::access::fence_space::local_space); + + total_misses = num_top_k_ - *s_total_hits_ptr - *s_newest_hit_ptr; + // Rewrite LRU order: misses then remaining evictables at the front (LRU), + // hits at the back (MRU). + { + const int total_evictable = hot_buffer_size_ - *s_total_hits_ptr; + for (int i = tid; i < hot_buffer_size_; i += block_size_) { + if (i < total_misses) { + req_lru_slots[total_evictable - total_misses + i] = s_lru_slots_out[hot_buffer_size_ - 1 - i]; + } else if (i < total_evictable) { + req_lru_slots[i - total_misses] = s_lru_slots_out[hot_buffer_size_ - 1 - i]; + } else { + req_lru_slots[i] = s_lru_slots_out[i - total_evictable]; + } + } + } + + // Each sub-group copies one miss directly from host cache to device buffer. + for (int miss_idx = sg_id; miss_idx < total_misses; miss_idx += num_sub_groups_) { + const int32_t miss_token = s_top_k_tokens[miss_idx]; + const int16_t evict_slot = s_lru_slots_out[hot_buffer_size_ - 1 - miss_idx]; + + const int64_t src_loc = req_host_cache_locs[miss_token]; + const int64_t dst_loc = static_cast(req_device_buffer_locs[evict_slot]); + + if constexpr (IsDsv4Layout) { + // Page-padded C4 device layout + page-padded host layout, K-only. + transfer_item( + lane_id, + sg_size, + device_buffer_k_, + const_cast(host_cache_k_), + static_cast(dst_loc), + static_cast(src_loc)); + } else { + // Generic path: device + host both linear, stride == item_size_bytes. + const char* src_k = static_cast(host_cache_k_) + src_loc * item_size_bytes_; + char* dst_k = static_cast(device_buffer_k_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_k, dst_k, item_size_bytes_); + + if constexpr (!IsMLA) { + const char* src_v = static_cast(host_cache_v_) + src_loc * item_size_bytes_; + char* dst_v = static_cast(device_buffer_v_) + dst_loc * item_size_bytes_; + transfer_item_linear(lane_id, sg_size, src_v, dst_v, item_size_bytes_); + } + } + } + } +}; + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp b/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp new file mode 100644 index 000000000..2434689b5 --- /dev/null +++ b/include/sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp @@ -0,0 +1,58 @@ +/** + * HiSparse: transfer_cache_dsv4_mla SYCL kernel. + * + * Bulk-copies DSv4-MLA C4 tokens between two sets of page-padded C4 buffers, one + * set per model layer. One sub-group copies one item across all layers, with a + * global-stride loop over items. + * + * src_caches / dst_caches are device arrays of `num_layers` raw cache base + * pointers (uint64_t values), one per layer. + */ + +#pragma once + +#include +#include + +#include "c4_layout.hpp" + +namespace sgl { +namespace sycl_kernel { +namespace hisparse { + +// Aggregate-initialized so the launcher can name the fields. +template +struct TransferCacheDsv4MlaKernel { + static_assert(BLOCK_SIZE % kSubGroupSize == 0, "BLOCK_SIZE must be a multiple of the sub-group size (32)."); + static constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + + void** src_caches_; + void** dst_caches_; + const int64_t* src_indices_; + const int64_t* dst_indices_; + uint32_t num_items_; + uint32_t num_layers_; + uint32_t total_sub_groups_; + + [[sycl::reqd_sub_group_size(kSubGroupSize)]] void operator()(::sycl::nd_item<1> item) const { + const ::sycl::sub_group sg = item.get_sub_group(); + const int lane_id = static_cast(sg.get_local_linear_id()); + const int sg_size = static_cast(sg.get_max_local_range()[0]); + + // Global sub-group index: group * subgroups_per_group + local subgroup index. + const uint32_t global_sg = + static_cast(item.get_group(0)) * kNumSubGroups + static_cast(sg.get_group_linear_id()); + + for (uint32_t i = global_sg; i < num_items_; i += total_sub_groups_) { + const int32_t src_index = static_cast(src_indices_[i]); + const int32_t dst_index = static_cast(dst_indices_[i]); + for (uint32_t layer_id = 0; layer_id < num_layers_; ++layer_id) { + transfer_item(lane_id, sg_size, dst_caches_[layer_id], src_caches_[layer_id], dst_index, src_index); + } + } + } +}; + +} // namespace hisparse +} // namespace sycl_kernel +} // namespace sgl diff --git a/include/sgl_kernel_ops.h b/include/sgl_kernel_ops.h index 18000d455..7d3880f66 100644 --- a/include/sgl_kernel_ops.h +++ b/include/sgl_kernel_ops.h @@ -1258,4 +1258,33 @@ void causal_conv1d_update( const std::optional& conv_state_indices_, int64_t pad_slot_id); +/* + * HiSparse hierarchical sparse KV cache (DeepSeek DSA / V4) + */ +void transfer_cache_dsv4_mla( + const at::Tensor& src_ptrs, + const at::Tensor& dst_ptrs, + const at::Tensor& src_indices, + const at::Tensor& dst_indices, + int64_t block_size); + +void load_cache_to_device_buffer_mla( + const at::Tensor& top_k_tokens, + const at::Tensor& device_buffer_tokens, + const at::Tensor& host_cache_locs, + const at::Tensor& device_buffer_locs, + const at::Tensor& host_cache, + const at::Tensor& device_buffer, + const at::Tensor& top_k_device_locs, + const at::Tensor& req_pool_indices, + const at::Tensor& seq_lens, + const at::Tensor& lru_slots, + const std::optional& num_real_reqs, + int64_t item_size_bytes, + int64_t num_top_k, + int64_t hot_buffer_size, + int64_t page_size, + int64_t block_size, + bool is_dsv4_layout); + #pragma GCC visibility pop diff --git a/python/sgl_kernel/__init__.py b/python/sgl_kernel/__init__.py index 2ecbf45ee..03abb0e80 100644 --- a/python/sgl_kernel/__init__.py +++ b/python/sgl_kernel/__init__.py @@ -80,6 +80,11 @@ ) from sgl_kernel.grammar import apply_token_bitmask_inplace_cuda from sgl_kernel.hadamard import hadamard_transform +from sgl_kernel.hisparse import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) from sgl_kernel.inkling_attn_prologue import ( compile_inkling_attn_prologue, inkling_attn_prologue_decode, diff --git a/python/sgl_kernel/hisparse.py b/python/sgl_kernel/hisparse.py new file mode 100644 index 000000000..765612ea2 --- /dev/null +++ b/python/sgl_kernel/hisparse.py @@ -0,0 +1,194 @@ +# HiSparse hierarchical sparse KV-cache ops for Intel XPU, used by DeepSeek +# DSA / V4 hierarchical sparse attention. +# Paged C4 layout (per page of 64 tokens): 64 value slots of 576 B, then 64 scale +# slots of 8 B, then padding so each page starts on a 576-byte boundary. + +from typing import Optional + +import torch + +# Work-group size for transfer_cache_dsv4_mla; only 256/512/1024 are compiled. +# All three measure within noise on BMG (Xe2), so this is an escape hatch rather +# than a tunable. +_DEFAULT_TRANSFER_BLOCK_SIZE = 1024 + +# Work-group size for the swap-in kernel: a plain runtime value, so any multiple +# of 32 works. +_DEFAULT_SWAP_IN_BLOCK_SIZE = 256 + + +def transfer_cache_dsv4_mla( + src_ptrs: torch.Tensor, + dst_ptrs: torch.Tensor, + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + block_size: int = _DEFAULT_TRANSFER_BLOCK_SIZE, +) -> None: + """Transfer DSv4 C4 tokens between page-padded C4 buffers, all layers. + + Args: + src_ptrs: 1-D uint64 tensor of per-layer source cache base pointers. + dst_ptrs: 1-D uint64 tensor of per-layer destination cache base pointers. + src_indices: 1-D int64 tensor of source token slot indices. + dst_indices: 1-D int64 tensor of destination token slot indices. + block_size: work-group size; one of 256, 512, 1024. + """ + torch.ops.sgl_kernel.transfer_cache_dsv4_mla.default( + src_ptrs, + dst_ptrs, + src_indices, + dst_indices, + block_size, + ) + + +def _load_cache_to_device_buffer_mla( + is_dsv4_layout: bool, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int, + block_size: int, + num_real_reqs: Optional[torch.Tensor], +) -> None: + torch.ops.sgl_kernel.load_cache_to_device_buffer_mla.default( + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + num_real_reqs, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + is_dsv4_layout, + ) + + +def load_cache_to_device_buffer_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = _DEFAULT_SWAP_IN_BLOCK_SIZE, + num_real_reqs: Optional[torch.Tensor] = None, +) -> None: + """Generic MLA hisparse swap-in: device + host both linear (stride=item_size_bytes). + + Swaps each request's top-k tokens into a small hot device buffer, maintaining + the per-request LRU order, streaming misses in from the host cache and writing + every top-k token's device slot to ``top_k_device_locs``. + + Args: + top_k_tokens: (batch, num_top_k) int32 top-k token positions per request. + device_buffer_tokens: (num_reqs, hot_buffer_size + 1) int32 token position + resident in each device-buffer slot; updated in place. + host_cache_locs: (num_reqs, max_seq_len) int64 host cache slot per token. + device_buffer_locs: (num_reqs, hot_buffer_size + 1) int32 device cache slot + per buffer slot. Must share ``stride(0)`` with device_buffer_tokens. + host_cache: host-side KV cache tensor. + device_buffer: device-side hot KV buffer tensor. + top_k_device_locs: (batch, num_top_k) int32 output device slots. + req_pool_indices: (batch,) int32/int64 request-pool row per batch entry. + seq_lens: (batch,) int32/int64 sequence lengths. + lru_slots: (num_reqs, hot_buffer_size) int16 LRU order; updated in place. + item_size_bytes: bytes per KV item (one token, all heads). + num_top_k: top-k count; must be <= hot_buffer_size. + hot_buffer_size: device buffer capacity in tokens, excluding the extra + slot reserved for the newest token. + page_size: accepted for API parity; unused. + block_size: work-group size; a multiple of 32. + num_real_reqs: (1,) int32 count of non-padded requests, for graph-captured + padded batches. Defaults to the full batch. + """ + _load_cache_to_device_buffer_mla( + False, + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + num_real_reqs, + ) + + +def load_cache_to_device_buffer_dsv4_mla( + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + top_k_device_locs: torch.Tensor, + req_pool_indices: torch.Tensor, + seq_lens: torch.Tensor, + lru_slots: torch.Tensor, + item_size_bytes: int, + num_top_k: int, + hot_buffer_size: int, + page_size: int = 1, + block_size: int = _DEFAULT_SWAP_IN_BLOCK_SIZE, + num_real_reqs: Optional[torch.Tensor] = None, +) -> None: + """DSv4 hisparse swap-in: page-padded device + page-padded host C4 layout. + + Same as :func:`load_cache_to_device_buffer_mla`, except the miss copy walks + the paged C4 layout described at the top of this module, so ``host_cache`` and + ``device_buffer`` must both be page-padded C4 buffers. + """ + _load_cache_to_device_buffer_mla( + True, + top_k_tokens, + device_buffer_tokens, + host_cache_locs, + device_buffer_locs, + host_cache, + device_buffer, + top_k_device_locs, + req_pool_indices, + seq_lens, + lru_slots, + item_size_bytes, + num_top_k, + hot_buffer_size, + page_size, + block_size, + num_real_reqs, + ) diff --git a/src/sycl/HiSparse.cpp b/src/sycl/HiSparse.cpp new file mode 100644 index 000000000..e2ee88536 --- /dev/null +++ b/src/sycl/HiSparse.cpp @@ -0,0 +1,314 @@ +#include + +#include + +#include "Utils.h" +#include "comm/General.h" +#include "sgl_kernel/hisparse/load_cache_to_device_buffer.hpp" +#include "sgl_kernel/hisparse/transfer_cache_dsv4_mla.hpp" +#include "sgl_kernel_export.h" + +using namespace sgl::sycl_kernel::hisparse; + +namespace { + +// --------------------------------------------------------------------------- +// transfer_cache_dsv4_mla +// --------------------------------------------------------------------------- + +template +void launch_transfer_cache_dsv4_mla( + void** src_caches, + void** dst_caches, + const int64_t* src_indices, + const int64_t* dst_indices, + uint32_t num_items, + uint32_t num_layers) { + constexpr int kNumSubGroups = BLOCK_SIZE / kSubGroupSize; + const uint32_t num_groups = div_up(num_items, static_cast(kNumSubGroups)); + const uint32_t total_sub_groups = num_groups * kNumSubGroups; + + TransferCacheDsv4MlaKernel kernel{ + .src_caches_ = src_caches, + .dst_caches_ = dst_caches, + .src_indices_ = src_indices, + .dst_indices_ = dst_indices, + .num_items_ = num_items, + .num_layers_ = num_layers, + .total_sub_groups_ = total_sub_groups, + }; + + auto cgf = DPCPP_Q_CGF(cgh) { + cgh.parallel_for( + sycl::nd_range<1>(sycl::range<1>(static_cast(num_groups) * BLOCK_SIZE), sycl::range<1>(BLOCK_SIZE)), + kernel); + }; + dpcppGetCurrentQueue().submit(cgf); +} + +// Validate a uint64 pointer table the same way KVCacheIO.cpp does. +void check_ptr_table(const at::Tensor& tbl, int64_t num_layers, const char* name) { + TORCH_CHECK(tbl.scalar_type() == at::kUInt64, name, " must be a uint64 pointer table"); + TORCH_CHECK(tbl.is_contiguous(), name, " must be contiguous"); + TORCH_CHECK(tbl.numel() == num_layers, name, " must have num_layers entries, got ", tbl.numel()); +} + +// --------------------------------------------------------------------------- +// load_cache_to_device_buffer +// --------------------------------------------------------------------------- + +struct LoadCacheArgs { + const int32_t* top_k_tokens; + int32_t* device_buffer_tokens; + const int64_t* host_cache_locs; + const int32_t* device_buffer_locs; + const void* host_cache_k; + const void* host_cache_v; + void* device_buffer_k; + void* device_buffer_v; + int32_t* top_k_device_locs; + const void* req_pool_indices; + const void* seq_lens; + int16_t* lru_slots; + const int32_t* num_real_reqs; + bool req_pool_indices_is_i64; + bool seq_lens_is_i64; + int64_t buffer_stride_0; + int64_t host_stride; + int64_t lru_slot_stride_0; + int64_t top_k_tokens_stride; + int64_t top_k_device_locs_stride; + int64_t item_size_bytes; + int64_t batch_size; + int block_size; + int num_top_k; + int hot_buffer_size; +}; + +template +void launch_load_cache_to_device_buffer(const LoadCacheArgs& a, const SmemLayout& layout) { + const int num_sub_groups = a.block_size / kSubGroupSize; + + auto cgf = DPCPP_Q_CGF(cgh) { + sycl::local_accessor smem(sycl::range<1>(static_cast(layout.total_int32_slots)), cgh); + LoadCacheToDeviceBufferKernel kernel{ + .top_k_tokens_ = a.top_k_tokens, + .device_buffer_tokens_ = a.device_buffer_tokens, + .host_cache_locs_ = a.host_cache_locs, + .device_buffer_locs_ = a.device_buffer_locs, + .host_cache_k_ = a.host_cache_k, + .host_cache_v_ = a.host_cache_v, + .device_buffer_k_ = a.device_buffer_k, + .device_buffer_v_ = a.device_buffer_v, + .top_k_device_locs_ = a.top_k_device_locs, + .req_pool_indices_ = a.req_pool_indices, + .seq_lens_ = a.seq_lens, + .lru_slots_ = a.lru_slots, + .num_real_reqs_ = a.num_real_reqs, + .req_pool_indices_is_i64_ = a.req_pool_indices_is_i64, + .seq_lens_is_i64_ = a.seq_lens_is_i64, + .buffer_stride_0_ = a.buffer_stride_0, + .host_stride_ = a.host_stride, + .lru_slot_stride_0_ = a.lru_slot_stride_0, + .top_k_tokens_stride_ = a.top_k_tokens_stride, + .top_k_device_locs_stride_ = a.top_k_device_locs_stride, + .item_size_bytes_ = a.item_size_bytes, + .block_size_ = a.block_size, + .num_sub_groups_ = num_sub_groups, + .num_top_k_ = a.num_top_k, + .hot_buffer_size_ = a.hot_buffer_size, + .hash_size_ = layout.hash_size, + .hash_mask_ = layout.hash_mask, + .num_buffer_chunks_ = layout.num_buffer_chunks, + .num_token_chunks_ = layout.num_token_chunks, + .iters_per_sg_buffer_ = div_up(layout.num_buffer_chunks, num_sub_groups), + .iters_per_sg_token_ = div_up(layout.num_token_chunks, num_sub_groups), + .total_int32_ = layout.total_int32, + .smem_ = smem, + }; + cgh.parallel_for( + sycl::nd_range<1>( + sycl::range<1>(static_cast(a.batch_size) * a.block_size), + sycl::range<1>(static_cast(a.block_size))), + kernel); + }; + dpcppGetCurrentQueue().submit(cgf); +} + +void check_swap_in_tensor(const at::Tensor& t, at::ScalarType expected, const char* name) { + TORCH_CHECK(t.scalar_type() == expected, name, " must be ", expected, ", got ", t.scalar_type()); + TORCH_CHECK(t.device().is_xpu(), name, " must be on XPU, got ", t.device()); + // The kernel is given only stride(0); every row must be contiguous. + TORCH_CHECK(t.dim() < 2 || t.stride(-1) == 1, name, " must be row-contiguous (stride(-1) == 1)"); +} + +} // namespace + +SGL_KERNEL_EXPORT void transfer_cache_dsv4_mla( + const at::Tensor& src_ptrs, + const at::Tensor& dst_ptrs, + const at::Tensor& src_indices, + const at::Tensor& dst_indices, + int64_t block_size) { + TORCH_CHECK(src_indices.scalar_type() == at::kLong, "src_indices must be int64"); + TORCH_CHECK(dst_indices.scalar_type() == at::kLong, "dst_indices must be int64"); + TORCH_CHECK(src_indices.is_contiguous(), "src_indices must be contiguous"); + TORCH_CHECK(dst_indices.is_contiguous(), "dst_indices must be contiguous"); + TORCH_CHECK( + src_indices.numel() == dst_indices.numel(), + "index count mismatch: ", + src_indices.numel(), + " vs ", + dst_indices.numel()); + + const int64_t num_layers = src_ptrs.numel(); + check_ptr_table(src_ptrs, num_layers, "src_ptrs"); + check_ptr_table(dst_ptrs, num_layers, "dst_ptrs"); + + const int64_t num_items = src_indices.numel(); + if (num_items == 0 || num_layers == 0) return; // nothing to transfer + + auto** src_caches = reinterpret_cast(src_ptrs.data_ptr()); + auto** dst_caches = reinterpret_cast(dst_ptrs.data_ptr()); + const auto* src_idx = src_indices.data_ptr(); + const auto* dst_idx = dst_indices.data_ptr(); + const auto items = static_cast(num_items); + const auto layers = static_cast(num_layers); + + // block_size is a template parameter; 1024 is the default, the rest are + // escape hatches (all three measure within noise on Xe2). + switch (block_size) { + case 256: + launch_transfer_cache_dsv4_mla<256>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + case 512: + launch_transfer_cache_dsv4_mla<512>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + case 1024: + launch_transfer_cache_dsv4_mla<1024>(src_caches, dst_caches, src_idx, dst_idx, items, layers); + break; + default: + TORCH_CHECK(false, "block_size must be one of 256, 512, 1024, got ", block_size); + } +} + +SGL_KERNEL_EXPORT void load_cache_to_device_buffer_mla( + const at::Tensor& top_k_tokens, + const at::Tensor& device_buffer_tokens, + const at::Tensor& host_cache_locs, + const at::Tensor& device_buffer_locs, + const at::Tensor& host_cache, + const at::Tensor& device_buffer, + const at::Tensor& top_k_device_locs, + const at::Tensor& req_pool_indices, + const at::Tensor& seq_lens, + const at::Tensor& lru_slots, + const std::optional& num_real_reqs, + int64_t item_size_bytes, + int64_t num_top_k, + int64_t hot_buffer_size, + int64_t page_size, + int64_t block_size, + bool is_dsv4_layout) { + TORCH_CHECK(num_top_k > 0, "num_top_k must be positive, got ", num_top_k); + TORCH_CHECK( + hot_buffer_size >= num_top_k, "hot_buffer_size (", hot_buffer_size, ") must be >= num_top_k (", num_top_k, ")"); + TORCH_CHECK( + block_size > 0 && block_size % kSubGroupSize == 0, + "block_size must be a positive multiple of ", + kSubGroupSize, + ", got ", + block_size); + TORCH_CHECK(item_size_bytes > 0, "item_size_bytes must be positive, got ", item_size_bytes); + // int16_t slot indices are stored in the LRU array and the hash values. + TORCH_CHECK(hot_buffer_size < 32767, "hot_buffer_size must fit in int16, got ", hot_buffer_size); + + check_swap_in_tensor(top_k_tokens, at::kInt, "top_k_tokens"); + check_swap_in_tensor(device_buffer_tokens, at::kInt, "device_buffer_tokens"); + check_swap_in_tensor(host_cache_locs, at::kLong, "host_cache_locs"); + check_swap_in_tensor(device_buffer_locs, at::kInt, "device_buffer_locs"); + check_swap_in_tensor(top_k_device_locs, at::kInt, "top_k_device_locs"); + check_swap_in_tensor(lru_slots, at::kShort, "lru_slots"); + TORCH_CHECK( + req_pool_indices.scalar_type() == at::kInt || req_pool_indices.scalar_type() == at::kLong, + "req_pool_indices must be int32 or int64, got ", + req_pool_indices.scalar_type()); + TORCH_CHECK( + seq_lens.scalar_type() == at::kInt || seq_lens.scalar_type() == at::kLong, + "seq_lens must be int32 or int64, got ", + seq_lens.scalar_type()); + TORCH_CHECK(req_pool_indices.is_contiguous(), "req_pool_indices must be contiguous"); + TORCH_CHECK(seq_lens.is_contiguous(), "seq_lens must be contiguous"); + TORCH_CHECK(host_cache_locs.dim() >= 2, "host_cache_locs must be at least 2-D"); + + TORCH_CHECK( + device_buffer_tokens.stride(0) == device_buffer_locs.stride(0), + "device_buffer_tokens and device_buffer_locs must share stride(0), got ", + device_buffer_tokens.stride(0), + " vs ", + device_buffer_locs.stride(0)); + + const int64_t batch_size = top_k_tokens.size(0); + TORCH_CHECK( + req_pool_indices.numel() >= batch_size && seq_lens.numel() >= batch_size, + "req_pool_indices / seq_lens must cover the batch (", + batch_size, + "), got ", + req_pool_indices.numel(), + " / ", + seq_lens.numel()); + if (batch_size == 0) return; + + const SmemLayout layout = SmemLayout::make(static_cast(num_top_k), static_cast(hot_buffer_size)); + const size_t local_mem_size = dpcppGetCurrentQueue().get_device().get_info(); + TORCH_CHECK( + layout.bytes() <= local_mem_size, + "hisparse swap-in needs ", + layout.bytes(), + " bytes of shared local memory for num_top_k=", + num_top_k, + " hot_buffer_size=", + hot_buffer_size, + ", but the device provides only ", + local_mem_size); + + at::Tensor real_reqs = + num_real_reqs.has_value() ? *num_real_reqs : at::full({1}, batch_size, top_k_tokens.options().dtype(at::kInt)); + TORCH_CHECK(real_reqs.scalar_type() == at::kInt, "num_real_reqs must be int32, got ", real_reqs.scalar_type()); + TORCH_CHECK(real_reqs.device().is_xpu(), "num_real_reqs must be on XPU, got ", real_reqs.device()); + + LoadCacheArgs args{ + .top_k_tokens = top_k_tokens.data_ptr(), + .device_buffer_tokens = device_buffer_tokens.data_ptr(), + .host_cache_locs = host_cache_locs.data_ptr(), + .device_buffer_locs = device_buffer_locs.data_ptr(), + .host_cache_k = host_cache.data_ptr(), + .host_cache_v = nullptr, // MLA: K-only + .device_buffer_k = device_buffer.data_ptr(), + .device_buffer_v = nullptr, // MLA: K-only + .top_k_device_locs = top_k_device_locs.data_ptr(), + .req_pool_indices = req_pool_indices.data_ptr(), + .seq_lens = seq_lens.data_ptr(), + .lru_slots = lru_slots.data_ptr(), + .num_real_reqs = real_reqs.data_ptr(), + .req_pool_indices_is_i64 = req_pool_indices.scalar_type() == at::kLong, + .seq_lens_is_i64 = seq_lens.scalar_type() == at::kLong, + .buffer_stride_0 = device_buffer_tokens.stride(0), + .host_stride = host_cache_locs.size(1), + .lru_slot_stride_0 = lru_slots.stride(0), + .top_k_tokens_stride = top_k_tokens.stride(0), + .top_k_device_locs_stride = top_k_device_locs.stride(0), + .item_size_bytes = item_size_bytes, + .batch_size = batch_size, + .block_size = static_cast(block_size), + .num_top_k = static_cast(num_top_k), + .hot_buffer_size = static_cast(hot_buffer_size), + }; + (void)page_size; + + if (is_dsv4_layout) { + launch_load_cache_to_device_buffer(args, layout); + } else { + launch_load_cache_to_device_buffer(args, layout); + } +} diff --git a/src/torch_extension_sycl.cc b/src/torch_extension_sycl.cc index a87d03388..e0908a5a3 100644 --- a/src/torch_extension_sycl.cc +++ b/src/torch_extension_sycl.cc @@ -596,6 +596,23 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) { "Tensor out_loc, Tensor! kvcache, bool is_decode, int compress_ratio, int page_size, bool use_fp4, " "int preshuffle_size=0, bool use_bf16_store=False) -> ()"); m.impl("fused_norm_rope_store", torch::kXPU, &at::native::xpu::fused_norm_rope_store); + + /* + * HiSparse hierarchical sparse KV cache kernels + */ + m.def( + "transfer_cache_dsv4_mla(Tensor src_ptrs, Tensor(a!) dst_ptrs, " + "Tensor src_indices, Tensor dst_indices, int block_size) -> ()"); + m.impl("transfer_cache_dsv4_mla", torch::kXPU, &transfer_cache_dsv4_mla); + + m.def( + "load_cache_to_device_buffer_mla(Tensor top_k_tokens, Tensor(a!) device_buffer_tokens, " + "Tensor host_cache_locs, Tensor device_buffer_locs, Tensor host_cache, " + "Tensor(b!) device_buffer, Tensor(c!) top_k_device_locs, Tensor req_pool_indices, " + "Tensor seq_lens, Tensor(d!) lru_slots, Tensor? num_real_reqs, int item_size_bytes, " + "int num_top_k, int hot_buffer_size, int page_size, int block_size, " + "bool is_dsv4_layout) -> ()"); + m.impl("load_cache_to_device_buffer_mla", torch::kXPU, &load_cache_to_device_buffer_mla); } REGISTER_EXTENSION(common_ops) diff --git a/tests/run_suite.py b/tests/run_suite.py index 75a515f94..9c0b1ff57 100644 --- a/tests/run_suite.py +++ b/tests/run_suite.py @@ -75,6 +75,7 @@ class TestFile: TestFile("test_sconv_metadata_and_windows.py"), TestFile("test_sconv_update_sconv_cache.py"), TestFile("test_inkling_attn_prologue.py"), + TestFile("test_hisparse.py"), ], # Nightly suite: exercises the wheel installed in the intel/sgl-kernel-xpu-dev # nightly image. Populate with longer-running or full-shape tests that are diff --git a/tests/test_hisparse.py b/tests/test_hisparse.py new file mode 100755 index 000000000..88b17c70e --- /dev/null +++ b/tests/test_hisparse.py @@ -0,0 +1,604 @@ +""" +Accuracy tests for the XPU/SYCL HiSparse swap-in kernels. + +The kernels pin the sub-group to 32 lanes, so the slot<->lane mapping and +eviction ordering are fully determined and the expected values below are exact. + +Guarded failure modes (derived-property + bug-regression): + - LRU hit/evict compaction and MRU/LRU write-back ordering. + - Miss classification, evict-slot reuse, and host->device miss copy. + - Multi-iteration chunk scan (hot_buffer_size past one sub-group window). + - Fast-path (seq_len <= hot_buffer) short-circuit leaves cache state untouched + while still writing -1 to every output slot it cannot resolve. + - Graph-capture padding (num_real_reqs) leaves padded requests' cache state + untouched and writes -1 across their output rows. + - DSv4 page-padded C4 addressing on both the transfer and swap-in paths. +""" + +import pytest +import torch +from sgl_kernel import ( + load_cache_to_device_buffer_dsv4_mla, + load_cache_to_device_buffer_mla, + transfer_cache_dsv4_mla, +) + +HAS_XPU = torch.xpu.is_available() + +pytestmark = pytest.mark.skipif(not HAS_XPU, reason="Requires XPU device") + +DEVICE = "xpu" +DTYPE = torch.float32 +KV_DIM = 8 +HOT_BUFFER_SIZE = 4 +PADDED_BUFFER_SIZE = HOT_BUFFER_SIZE + 1 +HOST_CACHE_SIZE = 16 +DEVICE_CACHE_SIZE = 16 +ITEM_SIZE_BYTES = KV_DIM * torch.empty((), dtype=DTYPE).element_size() +DSV4_PAGE_SIZE = 64 +DSV4_VALUE_BYTES = 576 +DSV4_SCALE_BYTES = 8 +DSV4_ITEM_BYTES = DSV4_VALUE_BYTES + DSV4_SCALE_BYTES +DSV4_PAGE_BYTES = ((DSV4_ITEM_BYTES * DSV4_PAGE_SIZE + 575) // 576) * 576 +DSV4_SCALE_OFFSET = DSV4_VALUE_BYTES * DSV4_PAGE_SIZE + + +def _pinned(shape, dtype): + """Host tensor pinned for the current XPU device.""" + return torch.empty(shape, dtype=dtype, device="cpu").pin_memory() + + +def _host_cache() -> torch.Tensor: + host_cache = _pinned((HOST_CACHE_SIZE, 1, KV_DIM), DTYPE) + host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache)) + return host_cache + + +def _dsv4_token_pattern(seed: int) -> tuple[torch.Tensor, torch.Tensor]: + value = ( + (torch.arange(DSV4_VALUE_BYTES, dtype=torch.int16) + seed) + .remainder(256) + .to(torch.uint8) + ) + scale = ( + (torch.arange(DSV4_SCALE_BYTES, dtype=torch.int16) + seed + 17) + .remainder(256) + .to(torch.uint8) + ) + return value, scale + + +def _write_dsv4_token(cache: torch.Tensor, loc: int, seed: int) -> None: + page = loc // DSV4_PAGE_SIZE + offset = loc % DSV4_PAGE_SIZE + value, scale = _dsv4_token_pattern(seed) + cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES].copy_( + value.to(cache.device) + ) + scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES + cache[page, scale_start : scale_start + DSV4_SCALE_BYTES].copy_( + scale.to(cache.device) + ) + + +def _read_dsv4_token(cache: torch.Tensor, loc: int) -> torch.Tensor: + page = loc // DSV4_PAGE_SIZE + offset = loc % DSV4_PAGE_SIZE + value = cache[page, offset * DSV4_VALUE_BYTES : (offset + 1) * DSV4_VALUE_BYTES] + scale_start = DSV4_SCALE_OFFSET + offset * DSV4_SCALE_BYTES + scale = cache[page, scale_start : scale_start + DSV4_SCALE_BYTES] + return torch.cat([value, scale]) + + +def _dsv4_ptrs(cache: torch.Tensor) -> torch.Tensor: + return torch.tensor([cache.data_ptr()], dtype=torch.uint64, device=DEVICE) + + +def _run_kernel( + *, + top_k_tokens: torch.Tensor, + device_buffer_tokens: torch.Tensor, + host_cache_locs: torch.Tensor, + device_buffer_locs: torch.Tensor, + host_cache: torch.Tensor, + device_buffer: torch.Tensor, + lru_slots: torch.Tensor, + seq_len: int | None = None, + seq_lens: torch.Tensor | None = None, + seq_lens_dtype: torch.dtype = torch.int32, + req_pool_indices: torch.Tensor | None = None, + num_real_reqs: int | None = None, + output_fill_value: int = -1, +) -> torch.Tensor: + batch_size = top_k_tokens.shape[0] + if req_pool_indices is None: + req_pool_indices = torch.arange(batch_size, dtype=torch.int64, device=DEVICE) + if seq_lens is None: + seq_lens = torch.full( + (batch_size,), seq_len, dtype=seq_lens_dtype, device=DEVICE + ) + if num_real_reqs is None: + num_real_reqs = batch_size + + # Cases that assert the kernel *writes* -1 (rather than merely leaving a slot + # alone) pass output_fill_value, so the default -1 cannot mask a skipped write. + out = torch.full_like(top_k_tokens, output_fill_value) + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + lru_slots=lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=top_k_tokens.shape[1], + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([num_real_reqs], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + return out + + +def _make_state( + device_buffer_locs_rows: list[list[int]], + device_buffer_tokens_rows: list[list[int]], + newest_tokens: list[int], +): + host_cache = _host_cache() + device_buffer = torch.full( + (DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE + ) + device_buffer_locs = torch.tensor( + device_buffer_locs_rows, dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + device_buffer_tokens_rows, dtype=torch.int32, device=DEVICE + ) + lru_slots = ( + torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + host_cache_locs = ( + torch.arange(HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE) + .view(1, -1) + .repeat(device_buffer_locs.shape[0], 1) + ) + + # Slots 0..3 participate in LRU; slot 4 is the reserved newest slot. + for rid, newest_token in enumerate(newest_tokens): + for slot, token in enumerate(device_buffer_tokens_rows[rid][:HOT_BUFFER_SIZE]): + if token >= 0: + device_buffer[device_buffer_locs[rid, slot]].copy_( + host_cache[token].to(DEVICE, non_blocking=True) + ) + device_buffer[device_buffer_locs[rid, HOT_BUFFER_SIZE]].copy_( + host_cache[newest_token].to(DEVICE, non_blocking=True) + ) + torch.xpu.synchronize() + + return { + "host_cache": host_cache, + "device_buffer": device_buffer, + "device_buffer_locs": device_buffer_locs, + "device_buffer_tokens": device_buffer_tokens, + "lru_slots": lru_slots, + "host_cache_locs": host_cache_locs, + } + + +def test_transfer_cache_dsv4_mla_copies_paged_token() -> None: + src_cache = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + dst_cache = _pinned((2, DSV4_PAGE_BYTES), torch.uint8) + dst_cache.zero_() + src_loc = DSV4_PAGE_SIZE + 6 + dst_loc = DSV4_PAGE_SIZE + 1 + _write_dsv4_token(src_cache, src_loc, seed=41) + + transfer_cache_dsv4_mla( + src_ptrs=_dsv4_ptrs(src_cache), + dst_ptrs=_dsv4_ptrs(dst_cache), + src_indices=torch.tensor([src_loc], dtype=torch.int64, device=DEVICE), + dst_indices=torch.tensor([dst_loc], dtype=torch.int64, device=DEVICE), + ) + torch.xpu.synchronize() + + assert torch.equal( + _read_dsv4_token(dst_cache, dst_loc).to(DEVICE), + _read_dsv4_token(src_cache, src_loc), + ) + + +def test_dsv4_swap_in_reads_paged_host_layout() -> None: + host_cache = _pinned((2, DSV4_PAGE_BYTES), torch.uint8) + host_cache.zero_() + device_buffer = torch.zeros((2, DSV4_PAGE_BYTES), dtype=torch.uint8, device=DEVICE) + host_loc = DSV4_PAGE_SIZE + 1 + swap_loc = DSV4_PAGE_SIZE + 12 + _write_dsv4_token(host_cache, host_loc, seed=41) + + top_k_tokens = torch.tensor([[3]], dtype=torch.int32, device=DEVICE) + device_buffer_tokens = torch.full( + (1, PADDED_BUFFER_SIZE), -1, dtype=torch.int32, device=DEVICE + ) + host_cache_locs = torch.zeros((1, 8), dtype=torch.int64, device=DEVICE) + host_cache_locs[0, 3] = host_loc + device_buffer_locs = torch.tensor( + [[swap_loc, swap_loc + 1, swap_loc + 2, swap_loc + 3, swap_loc + 4]], + dtype=torch.int32, + device=DEVICE, + ) + lru_slots = torch.arange(HOT_BUFFER_SIZE, dtype=torch.int16, device=DEVICE).view( + 1, -1 + ) + out = torch.full_like(top_k_tokens, -1) + + load_cache_to_device_buffer_dsv4_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE), + lru_slots=lru_slots, + item_size_bytes=DSV4_ITEM_BYTES, + num_top_k=1, + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + assert out.item() == swap_loc + assert torch.equal( + _read_dsv4_token(device_buffer, swap_loc), + _read_dsv4_token(host_cache, host_loc).to(DEVICE), + ) + + +def _long_case(): + # One-request baseline used by the stateful cases below: + # req 0 LRU slots : [0, 1, 2, 3] + # req 0 cached tokens : slot0->1, slot1->4, slot2->2, slot3->5 + # req 0 physical locs : slot0->9, slot1->7, slot2->3, slot3->5 + # req 0 newest slot : slot4/newest -> token 7 at physical loc 11 + return _make_state([[9, 7, 3, 5, 11]], [[1, 4, 2, 5, -1]], [7]) + + +@pytest.mark.parametrize("seq_lens_dtype", [torch.int32, torch.int64]) +def test_load_cache_to_device_buffer_fast_path(seq_lens_dtype: torch.dtype) -> None: + host_cache = _host_cache() + device_buffer = torch.arange( + DEVICE_CACHE_SIZE * KV_DIM, dtype=DTYPE, device=DEVICE + ).view(DEVICE_CACHE_SIZE, 1, KV_DIM) + device_buffer_before = device_buffer.clone() + device_buffer_locs = torch.tensor( + [[13, 9, 5, 1, 15]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens = torch.tensor( + [[10, 11, 12, 13, -1]], dtype=torch.int32, device=DEVICE + ) + device_buffer_tokens_before = device_buffer_tokens.clone() + lru_slots = torch.tensor([[0, 1, 2, 3]], dtype=torch.int16, device=DEVICE) + lru_slots_before = lru_slots.clone() + + # seq_len <= HOT_BUFFER_SIZE should skip host loads and LRU mutations, + # so top_k_tokens acts like direct indexing into device_buffer_locs. + out = _run_kernel( + top_k_tokens=torch.tensor([[2, 0, 1]], dtype=torch.int32, device=DEVICE), + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=torch.arange( + HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE + ).view(1, -1), + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + lru_slots=lru_slots, + seq_len=3, + seq_lens_dtype=seq_lens_dtype, + ) + + assert torch.equal(out.cpu(), torch.tensor([[5, 13, 9]], dtype=torch.int32)) + assert torch.equal(device_buffer_tokens.cpu(), device_buffer_tokens_before.cpu()) + assert torch.equal(lru_slots.cpu(), lru_slots_before.cpu()) + assert torch.equal(device_buffer.cpu(), device_buffer_before.cpu()) + + +def test_load_cache_to_device_buffer_fast_path_overwrites_stale_output() -> None: + # The fast path must write every one of the num_top_k output slots, not just + # the first `count = min(seq_len, num_top_k)`: slots past seq_len, and slots + # whose token position is negative, have to be set to -1 rather than left + # holding whatever the caller's buffer contained. + state = _make_state([[9, 7, 3, 5, 11]], [[0, 1, 2, 3, -1]], [4]) + + out = _run_kernel( + top_k_tokens=torch.tensor([[1, -1, 0, 0]], dtype=torch.int32, device=DEVICE), + seq_len=2, + output_fill_value=123456, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[7, -1, -1, -1]], dtype=torch.int32)) + + +def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None: + state = _long_case() + + # Query [4, 2, 7]: + # 4 hits slot1 -> loc 7 + # 2 hits slot2 -> loc 3 + # 7 is the newest token -> reserved newest loc 11 + # + # Hits move to the MRU tail, so [0, 1, 2, 3] becomes [0, 3, 1, 2]. + out = _run_kernel( + top_k_tokens=torch.tensor([[4, 2, 7]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[7, 3, 11]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[1, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[0, 3, 1, 2]], dtype=torch.int16) + ) + + +def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None: + state = _long_case() + + # Step 1: touch tokens [4, 2], so LRU becomes [0, 3, 1, 2]. + # Step 2: query token 6, which is a miss. + # The kernel should reuse the new LRU head slot0, whose physical loc is 9. + _run_kernel( + top_k_tokens=torch.tensor([[4, 2]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + out = _run_kernel( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + seq_len=8, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[6, 4, 2, 5, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[3, 1, 2, 0]], dtype=torch.int16) + ) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +def test_load_cache_to_device_buffer_multiple_misses_copy_all_slots() -> None: + state = _make_state( + [[9, 7, 3, 5, 11]], + [[0, 1, 2, 3, -1]], + [8], + ) + + out = _run_kernel( + top_k_tokens=torch.tensor([[4, 5, 6, 7]], dtype=torch.int32, device=DEVICE), + seq_len=9, + **state, + ) + + assert torch.equal(out.cpu(), torch.tensor([[9, 7, 3, 5]], dtype=torch.int32)) + assert torch.equal( + state["device_buffer_tokens"].cpu(), + torch.tensor([[4, 5, 6, 7, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"].cpu(), torch.tensor([[0, 1, 2, 3]], dtype=torch.int16) + ) + for token, loc in zip([4, 5, 6, 7], [9, 7, 3, 5]): + assert torch.equal( + state["device_buffer"][loc].cpu(), state["host_cache"][token] + ) + + +def test_load_cache_to_device_buffer_multi_iteration_scan_compacts_evictables() -> None: + # Regression for the multi-iteration prefix scan over the hot-buffer chunks. + # + # If the classification loop scans s_evict_chunk_offset with the full + # num_buffer_chunks + 1 element count, all 32 lanes participate even though + # only num_sub_groups entries were written this iteration: the lanes past that + # window re-read what an earlier iteration's scan left behind and fold it into + # the accumulator, so every later iteration compacts evictable slots to the + # wrong positions in s_lru_slots_out. + # + # Reproducing it needs a stale region inside the second iteration's window: + # num_buffer_chunks > 2 * num_sub_groups -> hot_buffer_size > 512 at block + # 256. At block_size 512 or 1024 num_sub_groups is >= 16 and it cannot + # trigger. 8192 additionally keeps the corrupted offsets inside + # s_lru_slots_out, so the failure is a deterministic wrong LRU order rather + # than an out-of-bounds local-memory write. + hot_buffer_size = 8192 + num_top_k = 32 + padded_size = hot_buffer_size + 1 + seq_len = hot_buffer_size + 1 # > hot_buffer_size, so the fast path is skipped + + host_cache = _pinned((num_top_k, 1, KV_DIM), DTYPE) + host_cache.copy_(torch.arange(host_cache.numel(), dtype=DTYPE).view_as(host_cache)) + device_buffer = torch.full( + (padded_size, 1, KV_DIM), -1.0, dtype=DTYPE, device=DEVICE + ) + # Identity slot -> device loc mapping keeps the expected values readable. + device_buffer_locs = torch.arange( + padded_size, dtype=torch.int32, device=DEVICE + ).view(1, -1) + # Nothing is resident, so every buffer slot is evictable (the path whose + # chunk counts are non-zero, hence the one that carries the corruption) and + # every query is a miss. + device_buffer_tokens = torch.full( + (1, padded_size), -1, dtype=torch.int32, device=DEVICE + ) + lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view( + 1, -1 + ) + host_cache_locs = torch.arange(seq_len, dtype=torch.int64, device=DEVICE).view( + 1, -1 + ) + top_k_tokens = torch.arange(num_top_k, dtype=torch.int32, device=DEVICE).view(1, -1) + out = torch.full_like(top_k_tokens, -1) + + load_cache_to_device_buffer_mla( + top_k_tokens=top_k_tokens, + device_buffer_tokens=device_buffer_tokens, + host_cache_locs=host_cache_locs, + device_buffer_locs=device_buffer_locs, + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([seq_len], dtype=torch.int32, device=DEVICE), + lru_slots=lru_slots, + item_size_bytes=ITEM_SIZE_BYTES, + num_top_k=num_top_k, + hot_buffer_size=hot_buffer_size, + page_size=1, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + # Evictables compact backwards from hot_buffer_size - 1 in LRU order, so + # miss j takes buffer slot j, and the LRU list rotates left by num_top_k as + # the just-filled slots move to the MRU tail. + assert torch.equal( + out.cpu(), torch.arange(num_top_k, dtype=torch.int32).view(1, -1) + ) + assert torch.equal( + device_buffer_tokens[0, :num_top_k].cpu(), + torch.arange(num_top_k, dtype=torch.int32), + ) + assert torch.equal( + device_buffer_tokens[0, num_top_k:].cpu(), + torch.full((padded_size - num_top_k,), -1, dtype=torch.int32), + ) + assert torch.equal( + lru_slots[0].cpu(), + torch.cat( + [ + torch.arange(num_top_k, hot_buffer_size, dtype=torch.int16), + torch.arange(num_top_k, dtype=torch.int16), + ] + ), + ) + for token in range(num_top_k): + assert torch.equal(device_buffer[token].cpu(), host_cache[token]) + + +def test_load_cache_to_device_buffer_batched_with_padding() -> None: + state = _make_state( + [ + [9, 7, 3, 5, 11], + [12, 10, 8, 6, 14], + [15, 4, 2, 1, 13], + ], + [ + [1, 4, 2, 5, -1], + [0, 1, 2, 3, -1], + [9, 8, 7, 6, -1], + ], + [7, 4, 5], + ) + padded_tokens_before = state["device_buffer_tokens"][2].clone() + padded_lru_before = state["lru_slots"][2].clone() + + # req 0: long path; req 1: fast path; req 2: padded block (must be ignored). + out = _run_kernel( + top_k_tokens=torch.tensor( + [[4, 6, 7], [2, 1, 0], [9, 8, 7]], dtype=torch.int32, device=DEVICE + ), + seq_lens=torch.tensor([8, 3, 8], dtype=torch.int32, device=DEVICE), + num_real_reqs=2, + # The padded row must be *written* as -1, not merely left alone. + output_fill_value=123456, + **state, + ) + + assert torch.equal( + out.cpu(), + torch.tensor([[7, 9, 11], [8, 10, 12], [-1, -1, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["device_buffer_tokens"][:2].cpu(), + torch.tensor([[6, 4, 2, 5, -1], [0, 1, 2, 3, -1]], dtype=torch.int32), + ) + assert torch.equal( + state["lru_slots"][:2].cpu(), + torch.tensor([[2, 3, 0, 1], [0, 1, 2, 3]], dtype=torch.int16), + ) + assert torch.equal( + state["device_buffer_tokens"][2].cpu(), padded_tokens_before.cpu() + ) + assert torch.equal(state["lru_slots"][2].cpu(), padded_lru_before.cpu()) + assert torch.equal(state["device_buffer"][9].cpu(), state["host_cache"][6]) + + +def test_load_cache_to_device_buffer_dsv4_mla_miss_copy_layout() -> None: + # Both the host cache and the device buffer use the page-padded C4 layout. + # The miss copy must read the host source with paged addressing + # (get_pointer_paged), not a linear per-item stride. + num_pages = (HOST_CACHE_SIZE + DSV4_PAGE_SIZE - 1) // DSV4_PAGE_SIZE + + state = _long_case() + host_cache = _pinned((num_pages, DSV4_PAGE_BYTES), torch.uint8) + host_cache.zero_() + for token in range(HOST_CACHE_SIZE): + _write_dsv4_token(host_cache, token, seed=token + 1) + + device_buffer = torch.full( + (num_pages, DSV4_PAGE_BYTES), + 0xFF, + dtype=torch.uint8, + device=DEVICE, + ) + out = torch.full((1, 1), -1, dtype=torch.int32, device=DEVICE) + + # Token 6 is a miss in _long_case(), so it should be copied into evict slot 0, + # whose physical device loc is 9. + load_cache_to_device_buffer_dsv4_mla( + top_k_tokens=torch.tensor([[6]], dtype=torch.int32, device=DEVICE), + device_buffer_tokens=state["device_buffer_tokens"], + host_cache_locs=state["host_cache_locs"], + device_buffer_locs=state["device_buffer_locs"], + host_cache=host_cache, + device_buffer=device_buffer, + top_k_device_locs=out, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([8], dtype=torch.int32, device=DEVICE), + lru_slots=state["lru_slots"], + item_size_bytes=DSV4_ITEM_BYTES, + num_top_k=1, + hot_buffer_size=HOT_BUFFER_SIZE, + page_size=DSV4_PAGE_SIZE, + block_size=256, + num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE), + ) + torch.xpu.synchronize() + + assert torch.equal(out.cpu(), torch.tensor([[9]], dtype=torch.int32)) + assert torch.equal( + _read_dsv4_token(device_buffer, 9).cpu(), + _read_dsv4_token(host_cache, 6), + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v", "-s"])