Adding different lx layout to micro-benchmarks

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

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
+  ]
+}

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
+  ]
+}

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
+  ]
+}

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
+  ]
+}

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}
+    }
+  ]
+}