diff --git a/scripts/microbench/README.md b/scripts/microbench/README.md index b743337a5..a6ce0fc00 100644 --- a/scripts/microbench/README.md +++ b/scripts/microbench/README.md @@ -67,6 +67,52 @@ Both lower onto the same `(query_lens, seq_lens)` path. `block_sizes: [64, 128]` sweeps block size as an extra axis. +## LX-resident kernel + +`--kv-layout lx` measures `_lx_page_attn_kernel` instead of the default kernel. It sets +`SPYRE_LX_KV_LAYOUT=1` before the impl is constructed, allocates the +`(page, kv_head)`-folded cache with `head_major_kv_layout` exactly as +`TorchSpyreModelRunner.initialize_kv_cache_tensors` does, and writes both the history +and the step's KV through `do_kv_cache_update` with the per-kv-head slot mapping +`attn_layer.SlotMapping` builds. The folded cache is allocated zeroed and filled on +device rather than host-populated, so the per-kv-head store path is exercised too. The +CPU reference reads the pages back and unfolds them, so the gate still compares against +what the kernel actually read. + +Only this layout needs the `SPYRE_LX_KV_LAYOUT` feature to be present; the others run on +a checkout without it, which is what makes a baseline arm possible. + +```bash +SPYRE_ATTN_PROFILING=1 .venv/bin/python3 scripts/microbench/spyre_attn_microbench.py \ + --config scripts/microbench/configs/granite33_8b_bs128_decode.json --kv-layout lx +``` + +Compiled variants only — eager cannot lower the 2-D page gather. Note the LX path caps +the attention compile's core count for small decode shapes (`_attn_max_cores`); +`SPYRE_ATTN_MAX_CORES` overrides it. + +Shipped sweep configs for the LX-vs-baseline comparison, each run twice — once plain, +once with `--kv-layout lx`: + +| config | axis | +| --- | --- | +| `lx_ab_short` | context length, single sequence | +| `lx_multiseq` | batch size at fixed context | +| `lx_mixed` | mixed prefill + decode in one batch | +| `lx_study_chunked_prefill` | chunked prefill, on- and off-bucket kv extents | +| `lx_pin_probe` | one shape, for reading the LX pin report | + +Because `num_blocks` is itself a latency axis — steeply so for the default kernel — every +config pins it. + +`max_model_len` (config key, or `--max-model-len`) sizes the bucketer's `num_blocks` +buckets. The stub model defaults to 2048, so a sweep reaching longer contexts must raise +it or those shapes come back as `error` rows. + +`--staging-rows` (config key `staging_rows`) shrinks the impl's staging buffers without +changing the measured shape, to separate the kernel's own cost from the staging +gather/store cost. + ## Output Tab-separated, written after every measurement (a crash keeps what completed) plus a diff --git a/scripts/microbench/configs/lx_ab_short.json b/scripts/microbench/configs/lx_ab_short.json new file mode 100644 index 000000000..8823d3d3b --- /dev/null +++ b/scripts/microbench/configs/lx_ab_short.json @@ -0,0 +1,50 @@ +{ + "run_label": "lx_ab_short", + "model": "granite-3.3-8b-instruct", + "num_query_heads": 32, + "num_kv_heads": 8, + "head_size": 128, + "block_size": 128, + "num_blocks": 64, + "seed": 0, + "warmup": 2, + "iterations": 10, + "atol": 0.3, + "rtol": 0.2, + "variants": [ + "online_softmax_compiled" + ], + "capture_batches": [ + { + "name": "decode_ctx512", + "capture_type": "decode", + "query_lens": [ + 1 + ], + "seq_lens": [ + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 1.0 + } + }, + { + "name": "decode_ctx1024", + "capture_type": "decode", + "query_lens": [ + 1 + ], + "seq_lens": [ + 1024 + ], + "extra_cols": { + "seqlen": 1024, + "decode_share": 1.0 + } + } + ], + "block_sizes": [ + 128 + ] +} diff --git a/scripts/microbench/configs/lx_mixed.json b/scripts/microbench/configs/lx_mixed.json new file mode 100644 index 000000000..5edbf7262 --- /dev/null +++ b/scripts/microbench/configs/lx_mixed.json @@ -0,0 +1,58 @@ +{ + "run_label": "lx_mixed", + "model": "granite-3.3-8b-instruct", + "num_query_heads": 32, + "num_kv_heads": 8, + "head_size": 128, + "block_size": 128, + "num_blocks": 64, + "seed": 0, + "warmup": 2, + "iterations": 10, + "atol": 0.3, + "rtol": 0.2, + "variants": [ + "online_softmax_compiled" + ], + "capture_batches": [ + { + "name": "pure_decode_3seq", + "capture_type": "decode", + "query_lens": [ + 1, + 1, + 1 + ], + "seq_lens": [ + 512, + 512, + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 1.0 + } + }, + { + "name": "mixed_1prefill_2decode", + "capture_type": "decode", + "query_lens": [ + 256, + 1, + 1 + ], + "seq_lens": [ + 256, + 512, + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 0.66 + } + } + ], + "block_sizes": [ + 128 + ] +} diff --git a/scripts/microbench/configs/lx_multiseq.json b/scripts/microbench/configs/lx_multiseq.json new file mode 100644 index 000000000..c98b73f96 --- /dev/null +++ b/scripts/microbench/configs/lx_multiseq.json @@ -0,0 +1,64 @@ +{ + "run_label": "lx_multiseq", + "model": "granite-3.3-8b-instruct", + "num_query_heads": 32, + "num_kv_heads": 8, + "head_size": 128, + "block_size": 128, + "num_blocks": 64, + "seed": 0, + "warmup": 2, + "iterations": 10, + "atol": 0.3, + "rtol": 0.2, + "variants": [ + "online_softmax_compiled" + ], + "capture_batches": [ + { + "name": "decode_1seq_ctx512", + "capture_type": "decode", + "query_lens": [ + 1 + ], + "seq_lens": [ + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 1.0 + } + }, + { + "name": "decode_8seq_ctx512", + "capture_type": "decode", + "query_lens": [ + 1, + 1, + 1, + 1, + 1, + 1, + 1, + 1 + ], + "seq_lens": [ + 512, + 512, + 512, + 512, + 512, + 512, + 512, + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 1.0 + } + } + ], + "block_sizes": [ + 128 + ] +} diff --git a/scripts/microbench/configs/lx_pin_probe.json b/scripts/microbench/configs/lx_pin_probe.json new file mode 100644 index 000000000..9e79a3dfc --- /dev/null +++ b/scripts/microbench/configs/lx_pin_probe.json @@ -0,0 +1,36 @@ +{ + "run_label": "lx_pin_probe", + "model": "granite-3.3-8b-instruct", + "num_query_heads": 32, + "num_kv_heads": 8, + "head_size": 128, + "block_size": 128, + "num_blocks": 64, + "seed": 0, + "warmup": 2, + "iterations": 10, + "atol": 0.3, + "rtol": 0.2, + "variants": [ + "online_softmax_compiled" + ], + "capture_batches": [ + { + "name": "decode_ctx512", + "capture_type": "decode", + "query_lens": [ + 1 + ], + "seq_lens": [ + 512 + ], + "extra_cols": { + "seqlen": 512, + "decode_share": 1.0 + } + } + ], + "block_sizes": [ + 128 + ] +} diff --git a/scripts/microbench/configs/lx_study_chunked_prefill.json b/scripts/microbench/configs/lx_study_chunked_prefill.json new file mode 100644 index 000000000..768921a1d --- /dev/null +++ b/scripts/microbench/configs/lx_study_chunked_prefill.json @@ -0,0 +1,69 @@ +{ + "run_label": "lx_chunked_prefill", + "model": "granite-3.3-8b-instruct", + "num_query_heads": 32, + "num_kv_heads": 8, + "head_size": 128, + "block_size": 128, + "num_blocks": 64, + "max_model_len": 8192, + "seed": 0, + "warmup": 2, + "iterations": 10, + "atol": 0.3, + "rtol": 0.2, + "variants": [ + "online_softmax_compiled" + ], + "capture_batches": [ + { + "name": "chunk_ctx0_kv512", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [512], + "extra_cols": {"seqlen": 512, "context_len_cfg": 0, "on_bucket": true, "decode_share": 0.0} + }, + { + "name": "chunk_ctx512_kv1024", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [1024], + "extra_cols": {"seqlen": 1024, "context_len_cfg": 512, "on_bucket": true, "decode_share": 0.0} + }, + { + "name": "chunk_ctx1024_kv1536", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [1536], + "extra_cols": {"seqlen": 1536, "context_len_cfg": 1024, "on_bucket": false, "decode_share": 0.0} + }, + { + "name": "chunk_ctx1536_kv2048", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [2048], + "extra_cols": {"seqlen": 2048, "context_len_cfg": 1536, "on_bucket": true, "decode_share": 0.0} + }, + { + "name": "chunk_ctx2560_kv3072", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [3072], + "extra_cols": {"seqlen": 3072, "context_len_cfg": 2560, "on_bucket": false, "decode_share": 0.0} + }, + { + "name": "chunk_ctx3584_kv4096", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [4096], + "extra_cols": {"seqlen": 4096, "context_len_cfg": 3584, "on_bucket": true, "decode_share": 0.0} + }, + { + "name": "chunk_ctx7680_kv8192", + "capture_type": "prefill", + "query_lens": [512], + "seq_lens": [8192], + "extra_cols": {"seqlen": 8192, "context_len_cfg": 7680, "on_bucket": true, "decode_share": 0.0} + } + ] +} diff --git a/scripts/microbench/spyre_attn_microbench.py b/scripts/microbench/spyre_attn_microbench.py index a232cdd7e..3f0867592 100644 --- a/scripts/microbench/spyre_attn_microbench.py +++ b/scripts/microbench/spyre_attn_microbench.py @@ -78,7 +78,7 @@ def register_variant(name, impl_label, compiled=True, available_fn=None): @contextlib.contextmanager -def spyre_vllm_config(compiled: bool): +def spyre_vllm_config(compiled: bool, max_model_len=None, staging_rows=None): """Establish a Spyre vLLM config context for standalone (non-pytest) use.""" from vllm.config import DeviceConfig, ModelConfig, VllmConfig, set_current_vllm_config from vllm.config.compilation import CompilationConfig, CompilationMode @@ -90,11 +90,18 @@ def spyre_vllm_config(compiled: bool): current_platform._enum = PlatformEnum.OOT register_all() mode = CompilationMode.STOCK_TORCH_COMPILE if compiled else CompilationMode.NONE + # max_model_len sizes the bucketer's num_blocks buckets, which cap the kv + # extent a shape can dispatch to; the stub model's default is 2048. config = VllmConfig( device_config=DeviceConfig(device="cpu"), compilation_config=CompilationConfig(custom_ops=["all"], mode=mode), - model_config=ModelConfig(dtype=DTYPE), + model_config=ModelConfig(dtype=DTYPE, max_model_len=max_model_len), ) + # Diagnostic only: SpyreAttentionImpl sizes its staging buffers as + # max_num_batched_tokens + 1, so this shrinks them without touching the measured + # product. Separates the kernel's own cost from the staging gather/store cost. + if staging_rows: + config.scheduler_config.max_num_batched_tokens = staging_rows - 1 with set_current_vllm_config(config), set_forward_context(None, config): yield @@ -171,6 +178,47 @@ def build_metadata( return builder.build(common_prefix_len=0, common_attn_metadata=common) +def _folded_slot_mapping(slots, num_kv_heads, block_size, device): + """Per-kv-head store index in the folded frame, as attn_layer.SlotMapping builds it.""" + from spyre_inference.custom_ops.utils import convert + + pages = torch.div(slots, block_size, rounding_mode="floor") + offsets = slots - pages * block_size + return [ + convert((pages * num_kv_heads + h) * block_size + offsets, device=device) + for h in range(num_kv_heads) + ] + + +def _fill_folded_cache(impl, inputs, kv_cache): + """Write history and this step's KV through the production scatter. + + The folded cache is allocated zeroed on device (as the runner does) and never + host-populated: the fold is a device-layout property, so a host tensor's + transfer does not reproduce it. + """ + from spyre_inference.custom_ops.utils import convert + + device = inputs["cache_device"] + for k_src, v_src, slots in inputs["kv_writes"]: + impl.do_kv_cache_update( + None, + convert(k_src, device=device), + convert(v_src, device=device), + kv_cache, + _folded_slot_mapping(slots, inputs["num_kv_heads"], inputs["block_size"], device), + ) + + +def read_kv_pages(inputs): + """KV pages back on the host in the plain [num_blocks, block_size, H, D] frame.""" + k, v = inputs["k_pages"].to("cpu"), inputs["v_pages"].to("cpu") + if not inputs["folded"]: + return k, v + shape = (-1, inputs["num_kv_heads"], inputs["block_size"], inputs["head_size"]) + return tuple(t.reshape(shape).permute(0, 2, 1, 3) for t in (k, v)) + + def _fused_qkv_kv_views(query, key, value, device): """K/V as the backend receives them: strided views of a fused QKV on device.""" from spyre_inference.custom_ops.utils import convert @@ -185,6 +233,20 @@ def _fused_qkv_kv_views(query, key, value, device): ) +def _padded_block_width(num_blocks: int) -> int: + """Block-table width build() can slice: the count rounded onto its buckets. + + Mirrors SpyreAttentionMetadataBuilder._pad_num_blocks. Kept as a local power-of-two + round rather than reaching for the bucketer, which needs a full vllm_config. The + real buckets derive from SPYRE_ATTN_KV_BUCKETS, so a sparse override can round to a + wider table than this predicts; the default buckets are powers of two. + """ + width = 1 + while width < num_blocks: + width *= 2 + return width + + def build_inputs_from_requests( query_lens, seq_lens, @@ -213,7 +275,13 @@ def build_inputs_from_requests( num_seqs = len(query_lens) max_kv = max(seq_lens) blocks_per_seq = (max_kv + block_size - 1) // block_size - if num_blocks < num_seqs * blocks_per_seq: + # build() pads each sequence's block count onto the bucket lattice and then + # slices block_table[s, :padded], so the table must be at least that wide -- + # production allocates it at ceil(max_model_len / block_size) for the engine's + # lifetime, for exactly this reason. A table sized to the real count raises + # "expanded size of the tensor" for any off-bucket kv length. + table_width = _padded_block_width(blocks_per_seq) + if num_blocks < num_seqs * table_width: return None # cache too small to give every sequence its own pages scale = head_size**-0.5 @@ -227,8 +295,10 @@ def build_inputs_from_requests( # Sample without replacement: an aliased page would let one sequence # overwrite another's KV and shrink the set of pages actually gathered. - block_tables = torch.randperm(num_blocks, dtype=torch.int32)[: num_seqs * blocks_per_seq].view( - num_seqs, blocks_per_seq + # Padded columns get real, distinct pages too: build() puts them in the gather, + # and their tiles are fully masked, so their contents never reach the output. + block_tables = torch.randperm(num_blocks, dtype=torch.int32)[: num_seqs * table_width].view( + num_seqs, table_width ) slot_mapping = [] @@ -268,6 +338,7 @@ def build_inputs_from_requests( ) cache_device = torch.device(device) + folded = kv_layout == "lx" and cache_device.type == "spyre" def to_device(cache): # plain: host-populated cache, plain transfer. Matches @@ -276,9 +347,21 @@ def to_device(cache): # reproduce the finding (see README). _reshape_and_cache views pages as # [-1, H, D] and relies on the slot-outermost device layout, which # convert() does not reproduce for a host tensor. + # lx: the folded frame, allocated zeroed exactly as + # TorchSpyreModelRunner.initialize_kv_cache_tensors does and filled on + # device afterwards; the fold is a device-layout property a host + # tensor's transfer cannot reproduce. if cache_device.type != "spyre" or kv_layout == "plain": return cache.to(cache_device) nb, bsz, h, d = cache.shape + if folded: + # Imported here, not at the top: a baseline checkout predating the folded + # frame has no head_major_kv_layout, and only ever asks for a plain cache. + from spyre_inference.v1.attention.backends.spyre_attn import head_major_kv_layout + + return torch.zeros(nb * h, bsz, d, dtype=cache.dtype).to( + cache_device, device_layout=head_major_kv_layout(nb * h, bsz, d, cache.dtype) + ) layout = slot_major_kv_layout(nb * bsz, h, d, cache.dtype) if kv_layout == "slot_major": return cache.to(cache_device, device_layout=layout) @@ -296,6 +379,19 @@ def to_device(cache): v_pages.view(view).index_copy_(0, slots_dev, hv_dev) key_dev, value_dev = _fused_qkv_kv_views(query, key, value, cache_device) + # Replayed on device by _fill_folded_cache. The step's KV goes in as the + # strided fused-QKV views the backend receives, not the host copies. + kv_writes = [(key_dev, value_dev, slot_mapping)] + if hist_slots: + kv_writes.insert( + 0, + ( + torch.cat(hist_k), + torch.cat(hist_v), + torch.tensor(hist_slots, dtype=torch.int64), + ), + ) + return { "query_dev": convert(query, cache_device), "key_dev": key_dev, @@ -312,6 +408,11 @@ def to_device(cache): "query_lens": list(query_lens), "seq_lens": list(seq_lens), "total_query_tokens": total_q, + "num_kv_heads": num_kv_heads, + "head_size": head_size, + "block_size": block_size, + "folded": folded, + "kv_writes": kv_writes, } @@ -435,6 +536,8 @@ def make_forward(inputs, num_query_heads, num_kv_heads, head_size): # not match, and the readback is then garbage. output = torch.empty_like(inputs["query_cpu"]).to(inputs["cache_device"]) kv_cache = SpyrePagedKVCache(k_pages=inputs["k_pages"], v_pages=inputs["v_pages"]) + if inputs["folded"]: + _fill_folded_cache(impl, inputs, kv_cache) @torch.inference_mode() def run(): @@ -538,7 +641,7 @@ def run_config(entry, variant, cfg, records, csv_path, block_size=None): kv_layout=cfg.get("kv_layout", "plain"), ) if inputs is None: - needed = len(query_lens) * row["num_kv_blocks_iterated"] + needed = len(query_lens) * _padded_block_width(row["num_kv_blocks_iterated"]) row["error"] = f"insufficient blocks (need num_blocks >= {needed})" print(f" -> skipped ({row['error']})", flush=True) records.append(row) @@ -558,10 +661,11 @@ def run_config(entry, variant, cfg, records, csv_path, block_size=None): atol, rtol = cfg.get("atol", 0.3), cfg.get("rtol", 0.2) max_outliers = cfg.get("max_outliers", 5) got = output.to("cpu").float() + k_ref, v_ref = read_kv_pages(inputs) ref = ref_attn( inputs["query_cpu"], - inputs["k_pages"].to("cpu"), - inputs["v_pages"].to("cpu"), + k_ref, + v_ref, inputs["query_lens"], inputs["seq_lens"], inputs["block_tables"], @@ -692,15 +796,31 @@ def main(): ap.add_argument("--iterations", type=int, default=None) ap.add_argument("--warmup", type=int, default=None) ap.add_argument("--device", default=None) + ap.add_argument( + "--max-model-len", + type=int, + default=None, + help="Context the bucketer sizes its num_blocks buckets from. Raise it above " + "the stub model's 2048 to reach longer kv extents.", + ) + ap.add_argument( + "--staging-rows", + type=int, + default=None, + help="Diagnostic: shrink the impl's staging buffers to this many rows without " + "changing the measured shape, to separate kernel cost from staging cost.", + ) ap.add_argument( "--kv-layout", - choices=["plain", "slot_major", "slot_major_devfill"], + choices=["plain", "slot_major", "slot_major_devfill", "lx"], default=None, help="KV page device layout. 'plain' (default) is correct for a " "host-populated cache. 'slot_major_devfill' matches the " "worker: zeroed slot-major alloc, history written on device. " "'slot_major' pins the worker layout on a host-populated " - "cache and is numerically wrong; kept to reproduce that.", + "cache and is numerically wrong; kept to reproduce that. " + "'lx' benchmarks the LX-resident kernel: sets SPYRE_LX_KV_LAYOUT=1 and " + "allocates the (page, kv_head)-folded cache as the worker does.", ) ap.add_argument( "--span", @@ -721,6 +841,8 @@ def main(): ("device", args.device), ("span", args.span), ("kv_layout", args.kv_layout), + ("max_model_len", args.max_model_len), + ("staging_rows", args.staging_rows), ): if val is not None: cfg[key] = val @@ -729,6 +851,14 @@ def main(): cfg["stop_on_failure"] = args.stop_on_failure cfg.setdefault("device", "spyre") + if cfg.get("kv_layout") == "lx": + # Read in SpyreAttentionImpl.__init__ and _attn_max_cores, and cached on + # first access, so set it before either runs. + os.environ["SPYRE_LX_KV_LAYOUT"] = "1" + from spyre_inference import envs + + envs.clear_env_cache() + variants = [v for v in cfg["variants"] if VARIANT_REGISTRY[v]["available"]()] if not variants: raise SystemExit("no available variants") @@ -740,6 +870,11 @@ def main(): "compiled and eager variants need separate runs (the compilation mode " "is fixed for the process). Re-run with --variants one at a time." ) + if cfg.get("kv_layout") == "lx" and not next(iter(compiled_modes)): + raise SystemExit( + "--kv-layout lx needs a compiled variant: the folded cache is gathered " + "with a 2-D index, which eager cannot lower." + ) entries = entries_from_config(cfg) stamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S") @@ -756,13 +891,18 @@ def main(): f"block_size={cfg.get('block_sizes') or cfg['block_size']} dtype={DTYPE}" ) print(f" span : {SPANS[cfg.get('span', 'online_softmax')]}") + print(f" kv_layout : {cfg.get('kv_layout', 'plain')}") print(f" variants : {variants}") print(f" shapes : {len(entries)}") print(f" iterations : {cfg.get('iterations', 10)} (warmup {cfg.get('warmup', 2)})") print(f" output : {out_dir if not args.no_output else 'none'}\n", flush=True) records = [] - with spyre_vllm_config(compiled=next(iter(compiled_modes))): + with spyre_vllm_config( + compiled=next(iter(compiled_modes)), + max_model_len=cfg.get("max_model_len"), + staging_rows=cfg.get("staging_rows"), + ): import torch._dynamo # The kernel specializes per (num_blocks, aligned_max_query_len), so a @@ -786,6 +926,8 @@ def main(): cfg["block_size"], cfg["num_blocks"], cfg["device"], + seed=cfg.get("seed", 0), + kv_layout=cfg.get("kv_layout", "plain"), ) probe_run, _ = make_forward( probe_inputs, cfg["num_query_heads"], cfg["num_kv_heads"], cfg["head_size"]