diff --git a/notebooks/demo_gen_mlir.py b/notebooks/demo_gen_mlir.py index 676ec92..b93fc5c 100644 --- a/notebooks/demo_gen_mlir.py +++ b/notebooks/demo_gen_mlir.py @@ -218,6 +218,144 @@ def gen_softmax_mlir(n_rows, row_width, num_cores=32): }}""" +def gen_rope_mlir(num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=256): + """Generate RoPE forward kernel (half-layout, LLaMA/Granite convention). + + y[:, 0:D/2] = x[:, 0:D/2] * cos - x[:, D/2:D] * sin + y[:, D/2:D] = x[:, 0:D/2] * sin + x[:, D/2:D] * cos + + Design decisions: + - 2D grid [seq_parts, head_groups] aligns with in-projection matmul and SDPA + in the prefill pipeline — no tensor redistribution at kernel boundaries. + This trades a tiny (~3%) intra-kernel AI loss vs [N,1] for zero + inter-kernel communication in a fused pipeline. + - cos/sin precomputed (standard Llama/Granite convention), loaded once per + seq-tile and reused across all heads in the inner loop. + - Single pass, no reduction — AI = 3h/(4h+2) ≈ 0.71–0.74 for typical + head counts; memory-bound above ~12 active cores. + (Assumes head_dim >= 128 so each half-row is a full HBM stick.) + """ + H = num_heads + S = seq_len + D = head_dim + half_d = D // 2 + rows = H * S + + if H % grid_h != 0: + raise ValueError( + f"gen_rope_mlir: num_heads={H} not divisible by grid_h={grid_h}") + if S % (grid_s * tile_seq) != 0: + raise ValueError( + f"gen_rope_mlir: seq_len={S} not divisible by " + f"grid_s*tile_seq={grid_s * tile_seq}") + + seq_per_core = S // grid_s + heads_per_core = H // grid_h + num_seq_tiles = seq_per_core // tile_seq + + x_view = _mem_view("x_view", "x_ptr", [rows, D], [D, 1]) + cos_view = _mem_view("cos_view", "cos_ptr", [S, half_d], [half_d, 1]) + sin_view = _mem_view("sin_view", "sin_ptr", [S, half_d], [half_d, 1]) + out_view = _mem_view("out_view", "out_ptr", [rows, D], [D, 1]) + + tile_shape = [tile_seq, half_d] + # Outer-loop access tiles (cos/sin): offset is [%cos_row, %c0] + cos_acc = _access_tile("cos_acc", "cos_view", ["%cos_row", "%c0"], + tile_shape, [S, half_d]) + sin_acc = _access_tile("sin_acc", "sin_view", ["%cos_row", "%c0"], + tile_shape, [S, half_d]) + # Inner-loop access tiles (x load, y store): offset uses %row + x_first_acc = _access_tile("x_first_acc", "x_view", ["%row", "%c0"], + tile_shape, [rows, D]) + x_second_acc = _access_tile("x_second_acc", "x_view", ["%row", "%c_half_d"], + tile_shape, [rows, D]) + y_first_acc = _access_tile("y_first_acc", "out_view", ["%row", "%c0"], + tile_shape, [rows, D]) + y_second_acc = _access_tile("y_second_acc", "out_view", ["%row", "%c_half_d"], + tile_shape, [rows, D]) + + return f"""module {{ + func.func @rope_fwd_kernel( + %x_ptr: index, %cos_ptr: index, %sin_ptr: index, %out_ptr: index + ) attributes {{grid = [{grid_s}, {grid_h}]}} {{ + %pid_s, %pid_h = ktdp.get_compute_tile_id : index, index + + %c0 = arith.constant 0 : index + %c1 = arith.constant 1 : index + %c_half_d = arith.constant {half_d} : index + %c_tile_seq = arith.constant {tile_seq} : index + %c_seq_per_core = arith.constant {seq_per_core} : index + %c_heads_per_core = arith.constant {heads_per_core} : index + %c_num_seq_tiles = arith.constant {num_seq_tiles} : index + %c_S = arith.constant {S} : index + + %seq_offset = arith.muli %pid_s, %c_seq_per_core : index + %head_offset = arith.muli %pid_h, %c_heads_per_core : index + +{x_view} + +{cos_view} + +{sin_view} + +{out_view} + + // Outer loop: seq-tiles (cos/sin loaded once, reused across heads) + scf.for %t = %c0 to %c_num_seq_tiles step %c1 {{ + + %tile_offset = arith.muli %t, %c_tile_seq : index + %cos_row = arith.addi %seq_offset, %tile_offset : index + + {cos_acc} + + %cos_tile = ktdp.load %cos_acc : !ktdp.access_tile<{tile_seq}x{half_d}xindex> -> tensor<{tile_seq}x{half_d}xf16> + + {sin_acc} + + %sin_tile = ktdp.load %sin_acc : !ktdp.access_tile<{tile_seq}x{half_d}xindex> -> tensor<{tile_seq}x{half_d}xf16> + + // Inner loop: heads (reuses cos/sin from outer loop) + scf.for %h = %c0 to %c_heads_per_core step %c1 {{ + + %h_abs = arith.addi %head_offset, %h : index + %row_base = arith.muli %h_abs, %c_S : index + %row = arith.addi %row_base, %cos_row : index + + {x_first_acc} + + %x_first = ktdp.load %x_first_acc : !ktdp.access_tile<{tile_seq}x{half_d}xindex> -> tensor<{tile_seq}x{half_d}xf16> + + {x_second_acc} + + %x_second = ktdp.load %x_second_acc : !ktdp.access_tile<{tile_seq}x{half_d}xindex> -> tensor<{tile_seq}x{half_d}xf16> + + // y_first = x_first * cos - x_second * sin + %tmp1 = arith.mulf %x_first, %cos_tile : tensor<{tile_seq}x{half_d}xf16> + %tmp2 = arith.mulf %x_second, %sin_tile : tensor<{tile_seq}x{half_d}xf16> + %y_first = arith.subf %tmp1, %tmp2 : tensor<{tile_seq}x{half_d}xf16> + + {y_first_acc} + + ktdp.store %y_first, %y_first_acc : tensor<{tile_seq}x{half_d}xf16>, !ktdp.access_tile<{tile_seq}x{half_d}xindex> + + // y_second = x_first * sin + x_second * cos + %tmp3 = arith.mulf %x_first, %sin_tile : tensor<{tile_seq}x{half_d}xf16> + %tmp4 = arith.mulf %x_second, %cos_tile : tensor<{tile_seq}x{half_d}xf16> + %y_second = arith.addf %tmp3, %tmp4 : tensor<{tile_seq}x{half_d}xf16> + + {y_second_acc} + + ktdp.store %y_second, %y_second_acc : tensor<{tile_seq}x{half_d}xf16>, !ktdp.access_tile<{tile_seq}x{half_d}xindex> + + scf.yield + }} + scf.yield + }} + return + }} +}}""" + + def gen_sdpa_mlir(seq_len, head_dim, block_m): """Generate a naive fused SDPA kernel (single-pass, no K-tiling). diff --git a/notebooks/demo_helpers.py b/notebooks/demo_helpers.py index 86b4305..2d3058c 100644 --- a/notebooks/demo_helpers.py +++ b/notebooks/demo_helpers.py @@ -17,6 +17,7 @@ from demo_gen_mlir import ( # noqa: F401 — re-exported for notebook compat gen_matmul_mlir, gen_softmax_mlir, + gen_rope_mlir, gen_sdpa_mlir, gen_sdpa_decode_pv_mlir, gen_paged_attention_mlir, @@ -248,6 +249,24 @@ def run_kernel_softmax(hw, n_rows, row_width, num_cores, rng=None): dict(n_rows=n_rows)) +def run_kernel_rope(hw, num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=256, rng=None): + """Generate RoPE MLIR, create tensors, run, return LatencyReport.""" + if rng is None: + rng = np.random.default_rng(0) + mlir = gen_rope_mlir(num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq) + half_dim = head_dim // 2 + freqs = 10000.0 ** (-np.arange(half_dim, dtype=np.float64) * 2.0 / head_dim) + positions = np.arange(seq_len, dtype=np.float64) + angles = np.outer(positions, freqs) + cos_table = np.cos(angles).astype(np.float16) + sin_table = np.sin(angles).astype(np.float16) + return run_kernel(hw, mlir, "rope_fwd_kernel", + dict(x_ptr=rng.standard_normal((num_heads * seq_len, head_dim)).astype(np.float16), + cos_ptr=cos_table, + sin_ptr=sin_table, + out_ptr=np.zeros((num_heads * seq_len, head_dim), dtype=np.float16))) + + def run_kernel_sdpa(hw, seq_len, head_dim, block_m, rng=None): """Generate SDPA MLIR, create tensors, run, return LatencyReport.""" if rng is None: diff --git a/notebooks/latency_demo.ipynb b/notebooks/latency_demo.ipynb index 4b29284..8ad15e5 100644 --- a/notebooks/latency_demo.ipynb +++ b/notebooks/latency_demo.ipynb @@ -34,7 +34,7 @@ } }, "outputs": [], - "source": "import numpy as np\nfrom ktir_cpu.interpreter import KTIRInterpreter\nfrom ktir_cpu.latency import HardwareConfig\nfrom demo_helpers import (\n gen_matmul_mlir, gen_softmax_mlir, gen_sdpa_mlir, gen_sdpa_decode_pv_mlir,\n gen_paged_attention_mlir,\n run_kernel, run_kernel_matmul, run_kernel_softmax, run_kernel_sdpa,\n run_kernel_sdpa_decode_pv, run_kernel_pa,\n show_mlir, plot_roofline,\n print_hw_config, print_core_roofline, print_per_core_table,\n print_kernel_comparison, print_chip_comparison, print_scaling_table, print_trace_summary,\n make_sdpa_tensors, make_pa_tensors, make_pa_scalars,\n)" + "source": "import numpy as np\nfrom ktir_cpu.interpreter import KTIRInterpreter\nfrom ktir_cpu.latency import HardwareConfig\nfrom demo_helpers import (\n gen_matmul_mlir, gen_softmax_mlir, gen_sdpa_mlir, gen_sdpa_decode_pv_mlir, gen_rope_mlir,\n gen_paged_attention_mlir,\n run_kernel, run_kernel_matmul, run_kernel_softmax, run_kernel_sdpa,\n run_kernel_sdpa_decode_pv, run_kernel_pa, run_kernel_rope,\n show_mlir, plot_roofline,\n print_hw_config, print_core_roofline, print_per_core_table,\n print_kernel_comparison, print_chip_comparison, print_scaling_table, print_trace_summary,\n make_sdpa_tensors, make_pa_tensors, make_pa_scalars,\n)" }, { "cell_type": "markdown", @@ -176,7 +176,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "---\n\n## 4. Multi-kernel roofline\n\nRun five kernels through the latency estimator and plot them together:\n- **matmul**, **sdpa**, **sdpa_decode_pv** and **paged_attn** — systolic-dominant (matmul ops)\n- **softmax** — simd-dominant (elementwise / transcendental ops)\n\n**SDPA** (Scaled Dot-Product Attention) computes `softmax(Q @ K^T / √d) @ V`. This is the\ncore attention primitive in every Transformer — our version is a naive single-pass fused\nkernel (no K-tiling). Flash Attention tiles over K blocks and uses online softmax to avoid\nmaterializing the full attention matrix; SDPA here shows what the un-tiled baseline looks like.\n\n**Multi-core decode SDPA (`sdpa_decode_pv`)** is the P@V stage of a single decode step, and\nthe only shape here that has to split work across cores — section 10 takes it apart. `sdpa`\ncomes out memory-bound below too, but for it that is a `block_m` choice: raise `block_m` to 32\nand it crosses the ridge. Decode P@V cannot be moved — its arithmetic intensity is capped by\nthe two tile widths alone — and it is the only one whose `comm` column is not zero.\n\n**Paged attention** extends SDPA with (1) a KV cache stored in fixed-size pages, (2) an\nindirection table (`block_tables`) so pages can be allocated non-contiguously, and (3) online\nsoftmax (the Flash Attention trick) streaming over pages. This is the decode-time attention\npattern in vLLM / TGI / SGLang — the indirect loads via `construct_indirect_access_tile` are\nthe key KTIR modelling challenge." + "source": "---\n\n## 4. Multi-kernel roofline\n\nRun five kernels through the latency estimator and plot them together:\n- **matmul**, **sdpa**, **sdpa_decode_pv** and **paged_attn** — systolic-dominant (matmul ops)\n- **softmax**, **rope**, and **paged_attn** — simd-dominant (elementwise / transcendental ops)\n\n**RoPE** (Rotary Position Embedding) applies position-dependent rotations to Q and K after\nthe in-projection matmul: `y[0:D/2] = x[0:D/2]*cos - x[D/2:]*sin; y[D/2:] = x[0:D/2]*sin + x[D/2:]*cos`. It is embarrassingly\nparallel over both sequence positions and attention heads, using a 2D grid `[seq, heads]`\nthat aligns with the adjacent matmul's tiling. cos/sin tables are precomputed and reused\nacross all heads within each sequence tile (the key locality optimization).\n\n**SDPA** (Scaled Dot-Product Attention) computes `softmax(Q @ K^T / √d) @ V`. This is the\ncore attention primitive in every Transformer — our version is a naive single-pass fused\nkernel (no K-tiling). Flash Attention tiles over K blocks and uses online softmax to avoid\nmaterializing the full attention matrix; SDPA here shows what the un-tiled baseline looks like.\n\n**Multi-core decode SDPA (`sdpa_decode_pv`)** is the P@V stage of a single decode step, and\nthe only shape here that has to split work across cores — section 10 takes it apart. `sdpa`\ncomes out memory-bound below too, but for it that is a `block_m` choice: raise `block_m` to 32\nand it crosses the ridge. Decode P@V cannot be moved — its arithmetic intensity is capped by\nthe two tile widths alone — and it is the only one whose `comm` column is not zero.\n\n**Paged attention** extends SDPA with (1) a KV cache stored in fixed-size pages, (2) an\nindirection table (`block_tables`) so pages can be allocated non-contiguously, and (3) online\nsoftmax (the Flash Attention trick) streaming over pages. This is the decode-time attention\npattern in vLLM / TGI / SGLang — the indirect loads via `construct_indirect_access_tile` are\nthe key KTIR modelling challenge." }, { "cell_type": "code", @@ -190,11 +190,11 @@ } }, "outputs": [], - "source": "rng = np.random.default_rng(0)\n\n# --- Matmul (4 cores, high K for AI) ---\nMM_M, MM_N, MM_K = 256, 256, 128\nMM_BM, MM_BN, MM_BK = 128, 128, 128\nmm_grid = (MM_M // MM_BM, MM_N // MM_BN)\nprint(f\"Matmul: {MM_M}×{MM_K}×{MM_N}, tiles {MM_BM}×{MM_BN}×{MM_BK}, grid={mm_grid} = {mm_grid[0]*mm_grid[1]} cores\")\nmatmul_report = run_kernel_matmul(hw, MM_M, MM_N, MM_K, MM_BM, MM_BN, MM_BK, rng)\n\n# --- Softmax (small, 4 cores) ---\nSM_ROWS, SM_WIDTH, SM_CORES = 64, 1024, 4\nprint(f\"Softmax: {SM_ROWS}×{SM_WIDTH}, grid=[{SM_CORES},1] = {SM_CORES} cores\")\nsoftmax_full_report = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, SM_CORES, rng)\n\n# --- SDPA (4 cores, smaller hd for speed, same AI structure) ---\nSDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M = 64, 128, 16\nsdpa_grid = SDPA_SEQ_LEN // SDPA_BLOCK_M\nprint(f\"SDPA: seq_len={SDPA_SEQ_LEN}, head_dim={SDPA_HEAD_DIM}, block_m={SDPA_BLOCK_M}, grid=[{sdpa_grid}] = {sdpa_grid} cores\")\nsdpa_report = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\n\n# --- Paged Attention (4 cores, shorter context, smaller hd) ---\nPA_NUM_TOKENS, PA_CONTEXT_LEN = 64, 512\nPA_NUM_QUERY_HEADS, PA_NUM_KV_HEADS = 4, 1\nPA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q = 64, 64, 16\n\npa_grid = (PA_NUM_TOKENS // PA_BLOCK_Q, PA_NUM_KV_HEADS)\npa_num_tiles = (PA_CONTEXT_LEN + PA_BLOCK_SIZE - 1) // PA_BLOCK_SIZE\nprint(f\"Paged Attention: {PA_NUM_TOKENS} tokens, ctx={PA_CONTEXT_LEN}, grid={pa_grid} = {pa_grid[0]*pa_grid[1]} cores\")\nprint(f\" num_tiles={pa_num_tiles} iterations\")\npaged_attn_report = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n\n# --- Multi-core decode SDPA, P@V stage (output x2, KV contraction x16 = 32 cores) ---\n# Granite-8B grouped-query ratio: 32 query heads over 8 key/value heads at head_dim 128,\n# so one KV head is shared by 4 query heads. Own generator, so the draws above stay put.\nPV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV = 8192, 128, 4\nPV_OUT_SPLIT, PV_K_SPLIT = 2, 16\npv_grid = (PV_OUT_SPLIT, PV_K_SPLIT)\nprint(f\"Decode SDPA P@V: kv_len={PV_KV_LEN}, head_dim={PV_HEAD_DIM}, \"\n f\"q_per_kv={PV_Q_PER_KV}, grid={pv_grid} = {PV_OUT_SPLIT*PV_K_SPLIT} cores, \"\n f\"cross-core fan-in {PV_K_SPLIT}\")\nsdpa_pv_report, sdpa_pv_err = run_kernel_sdpa_decode_pv(\n hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, PV_OUT_SPLIT, PV_K_SPLIT,\n rng=np.random.default_rng(169))\nprint(f\" Max abs error vs fp32 numpy: {sdpa_pv_err['max_abs']:.4f} \"\n f\"(relative {sdpa_pv_err['max_rel']:.2%}, \"\n f\"{sdpa_pv_err['zero_rows']} all-zero output rows)\")\n\nkernels = [(\"matmul\", matmul_report), (\"softmax\", softmax_full_report),\n (\"sdpa\", sdpa_report), (\"sdpa_decode_pv\", sdpa_pv_report),\n (\"paged_attn\", paged_attn_report)]\nprint_kernel_comparison(kernels)" + "source": "rng = np.random.default_rng(0)\n\n# --- Matmul (4 cores, high K for AI) ---\nMM_M, MM_N, MM_K = 256, 256, 128\nMM_BM, MM_BN, MM_BK = 128, 128, 128\nmm_grid = (MM_M // MM_BM, MM_N // MM_BN)\nprint(f\"Matmul: {MM_M}×{MM_K}×{MM_N}, tiles {MM_BM}×{MM_BN}×{MM_BK}, grid={mm_grid} = {mm_grid[0]*mm_grid[1]} cores\")\nmatmul_report = run_kernel_matmul(hw, MM_M, MM_N, MM_K, MM_BM, MM_BN, MM_BK, rng)\n\n# --- Softmax (small, 4 cores) ---\nSM_ROWS, SM_WIDTH, SM_CORES = 64, 1024, 4\nprint(f\"Softmax: {SM_ROWS}×{SM_WIDTH}, grid=[{SM_CORES},1] = {SM_CORES} cores\")\nsoftmax_full_report = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, SM_CORES, rng)\n\n# --- SDPA (4 cores, smaller hd for speed, same AI structure) ---\nSDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M = 64, 128, 16\nsdpa_grid = SDPA_SEQ_LEN // SDPA_BLOCK_M\nprint(f\"SDPA: seq_len={SDPA_SEQ_LEN}, head_dim={SDPA_HEAD_DIM}, block_m={SDPA_BLOCK_M}, grid=[{sdpa_grid}] = {sdpa_grid} cores\")\nsdpa_report = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\n\n# --- RoPE (Granite-8B: 32Q+8K heads, head_dim=128, 4 cores [2,2]) ---\nROPE_HEADS, ROPE_SEQ, ROPE_HEAD_DIM, ROPE_TILE_SEQ = 40, 4096, 128, 256\nprint(f\"RoPE: {ROPE_HEADS}h × {ROPE_SEQ}s × {ROPE_HEAD_DIM}d, grid=[2,2] = 4 cores\")\nrope_report = run_kernel_rope(hw, ROPE_HEADS, ROPE_SEQ, ROPE_HEAD_DIM, grid_s=2, grid_h=2, tile_seq=ROPE_TILE_SEQ, rng=rng)\n\n# --- Paged Attention (4 cores, shorter context, smaller hd) ---\nPA_NUM_TOKENS, PA_CONTEXT_LEN = 64, 512\nPA_NUM_QUERY_HEADS, PA_NUM_KV_HEADS = 4, 1\nPA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q = 64, 64, 16\n\npa_grid = (PA_NUM_TOKENS // PA_BLOCK_Q, PA_NUM_KV_HEADS)\npa_num_tiles = (PA_CONTEXT_LEN + PA_BLOCK_SIZE - 1) // PA_BLOCK_SIZE\nprint(f\"Paged Attention: {PA_NUM_TOKENS} tokens, ctx={PA_CONTEXT_LEN}, grid={pa_grid} = {pa_grid[0]*pa_grid[1]} cores\")\nprint(f\" num_tiles={pa_num_tiles} iterations\")\npaged_attn_report = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n\n# --- Multi-core decode SDPA, P@V stage (output x2, KV contraction x16 = 32 cores) ---\n# Granite-8B grouped-query ratio: 32 query heads over 8 key/value heads at head_dim 128,\n# so one KV head is shared by 4 query heads. Own generator, so the draws above stay put.\nPV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV = 8192, 128, 4\nPV_OUT_SPLIT, PV_K_SPLIT = 2, 16\npv_grid = (PV_OUT_SPLIT, PV_K_SPLIT)\nprint(f\"Decode SDPA P@V: kv_len={PV_KV_LEN}, head_dim={PV_HEAD_DIM}, \"\n f\"q_per_kv={PV_Q_PER_KV}, grid={pv_grid} = {PV_OUT_SPLIT*PV_K_SPLIT} cores, \"\n f\"cross-core fan-in {PV_K_SPLIT}\")\nsdpa_pv_report, sdpa_pv_err = run_kernel_sdpa_decode_pv(\n hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, PV_OUT_SPLIT, PV_K_SPLIT,\n rng=np.random.default_rng(169))\nprint(f\" Max abs error vs fp32 numpy: {sdpa_pv_err['max_abs']:.4f} \"\n f\"(relative {sdpa_pv_err['max_rel']:.2%}, \"\n f\"{sdpa_pv_err['zero_rows']} all-zero output rows)\")\n\nkernels = [(\"matmul\", matmul_report), (\"softmax\", softmax_full_report),\n (\"rope\", rope_report), (\"sdpa\", sdpa_report), (\"sdpa_decode_pv\", sdpa_pv_report),\n (\"paged_attn\", paged_attn_report)]\nprint_kernel_comparison(kernels)" }, { "cell_type": "code", - "source": "# --- Set 1: 4-core, tiles 256³ ---\nM, N, K = 512, 512, 256\nBM, BN, BK = 256, 256, 256\n\nmm_1 = run_kernel_matmul(hw, M, N, K, BM, BN, BK, rng)\nsm_1 = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, num_cores=4, rng=rng)\nsdpa_1 = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\npa_1 = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n# 4 cores: the KV contraction is split 2 ways instead of 16, so the fan-in drops with it.\npv_1, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, 2, 2,\n rng=np.random.default_rng(169))\n\n# --- Set 2: strong-scaling, 32-core (same problem, smaller tiles) ---\nmm_2 = run_kernel_matmul(hw, M, N, K, BM // 4, BN // 2, BK, rng)\nsm_2 = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, num_cores=32, rng=rng)\nsdpa_2 = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M // 8, rng)\npa_2 = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q // 8, rng)\n# Same problem, 8x the cores: the extra cores come from splitting the contraction harder.\npv_2, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, 2, 16,\n rng=np.random.default_rng(169))\n\n# --- Set 3: weak-scaling, 32-core (8x problem, same tiles) ---\nmm_3 = run_kernel_matmul(hw, M * 4, N * 2, K, BM, BN, BK, rng)\nsm_3 = run_kernel_softmax(hw, SM_ROWS * 8, SM_WIDTH, num_cores=32, rng=rng)\nsdpa_3 = run_kernel_sdpa(hw, SDPA_SEQ_LEN * 8, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\npa_3 = run_kernel_pa(hw, PA_NUM_TOKENS * 8, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n# 8x the KV length on 8x the cores, so each core keeps the 4096-token slice it had at 4 cores.\npv_3, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN * 8, PV_HEAD_DIM, PV_Q_PER_KV, 2, 16,\n rng=np.random.default_rng(169))\n\n# All 15 experiments: ● set 1 (4-core), ■ set 2 (strong 32), ▲ set 3 (weak 32)\nall_kernels = [\n (\"matmul (4c)\", mm_1, \"o\"), (\"softmax (4c)\", sm_1, \"o\"),\n (\"sdpa (4c)\", sdpa_1, \"o\"), (\"paged_attn (4c)\", pa_1, \"o\"),\n (\"sdpa_decode_pv (4c)\", pv_1, \"o\"),\n (\"matmul (strong-32)\", mm_2, \"s\"), (\"softmax (strong-32)\", sm_2, \"s\"),\n (\"sdpa (strong-32)\", sdpa_2, \"s\"), (\"paged_attn (strong-32)\", pa_2, \"s\"),\n (\"sdpa_decode_pv (strong-32)\", pv_2, \"s\"),\n (\"matmul (weak-32)\", mm_3, \"^\"), (\"softmax (weak-32)\", sm_3, \"^\"),\n (\"sdpa (weak-32)\", sdpa_3, \"^\"), (\"paged_attn (weak-32)\", pa_3, \"^\"),\n (\"sdpa_decode_pv (weak-32)\", pv_3, \"^\"),\n]\nprint_chip_comparison(all_kernels)", + "source": "# --- Set 1: 4-core, tiles 256³ ---\nM, N, K = 512, 512, 256\nBM, BN, BK = 256, 256, 256\n\nmm_1 = run_kernel_matmul(hw, M, N, K, BM, BN, BK, rng)\nsm_1 = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, num_cores=4, rng=rng)\nrope_1 = run_kernel_rope(hw, ROPE_HEADS, ROPE_SEQ, ROPE_HEAD_DIM, grid_s=2, grid_h=2, tile_seq=ROPE_TILE_SEQ, rng=rng)\nsdpa_1 = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\npa_1 = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n# 4 cores: the KV contraction is split 2 ways instead of 16, so the fan-in drops with it.\npv_1, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, 2, 2,\n rng=np.random.default_rng(169))\n\n# --- Set 2: strong-scaling, 32-core (same problem, smaller tiles) ---\nmm_2 = run_kernel_matmul(hw, M, N, K, BM // 4, BN // 2, BK, rng)\nsm_2 = run_kernel_softmax(hw, SM_ROWS, SM_WIDTH, num_cores=32, rng=rng)\nrope_2 = run_kernel_rope(hw, ROPE_HEADS, ROPE_SEQ, ROPE_HEAD_DIM, grid_s=8, grid_h=4, tile_seq=ROPE_TILE_SEQ, rng=rng)\nsdpa_2 = run_kernel_sdpa(hw, SDPA_SEQ_LEN, SDPA_HEAD_DIM, SDPA_BLOCK_M // 8, rng)\npa_2 = run_kernel_pa(hw, PA_NUM_TOKENS, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q // 8, rng)\n# Same problem, 8x the cores: the extra cores come from splitting the contraction harder.\npv_2, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN, PV_HEAD_DIM, PV_Q_PER_KV, 2, 16,\n rng=np.random.default_rng(169))\n\n# --- Set 3: weak-scaling, 32-core (8x problem, same tiles) ---\nmm_3 = run_kernel_matmul(hw, M * 4, N * 2, K, BM, BN, BK, rng)\nsm_3 = run_kernel_softmax(hw, SM_ROWS * 8, SM_WIDTH, num_cores=32, rng=rng)\nrope_3 = run_kernel_rope(hw, ROPE_HEADS, ROPE_SEQ * 8, ROPE_HEAD_DIM, grid_s=8, grid_h=4, tile_seq=ROPE_TILE_SEQ, rng=rng)\nsdpa_3 = run_kernel_sdpa(hw, SDPA_SEQ_LEN * 8, SDPA_HEAD_DIM, SDPA_BLOCK_M, rng)\npa_3 = run_kernel_pa(hw, PA_NUM_TOKENS * 8, PA_CONTEXT_LEN, PA_NUM_QUERY_HEADS,\n PA_NUM_KV_HEADS, PA_HEAD_DIM, PA_BLOCK_SIZE, PA_BLOCK_Q, rng)\n# 8x the KV length on 8x the cores, so each core keeps the 4096-token slice it had at 4 cores.\npv_3, _ = run_kernel_sdpa_decode_pv(hw, PV_KV_LEN * 8, PV_HEAD_DIM, PV_Q_PER_KV, 2, 16,\n rng=np.random.default_rng(169))\n\n# All 15 experiments: ● set 1 (4-core), ■ set 2 (strong 32), ▲ set 3 (weak 32)\nall_kernels = [\n (\"matmul (4c)\", mm_1, \"o\"), (\"softmax (4c)\", sm_1, \"o\"),\n (\"rope (4c)\", rope_1, \"o\"),\n (\"sdpa (4c)\", sdpa_1, \"o\"), (\"paged_attn (4c)\", pa_1, \"o\"),\n (\"sdpa_decode_pv (4c)\", pv_1, \"o\"),\n (\"matmul (strong-32)\", mm_2, \"s\"), (\"softmax (strong-32)\", sm_2, \"s\"),\n (\"rope (strong-32)\", rope_2, \"s\"),\n (\"sdpa (strong-32)\", sdpa_2, \"s\"), (\"paged_attn (strong-32)\", pa_2, \"s\"),\n (\"sdpa_decode_pv (strong-32)\", pv_2, \"s\"),\n (\"matmul (weak-32)\", mm_3, \"^\"), (\"softmax (weak-32)\", sm_3, \"^\"),\n (\"rope (weak-32)\", rope_3, \"^\"),\n (\"sdpa (weak-32)\", sdpa_3, \"^\"), (\"paged_attn (weak-32)\", pa_3, \"^\"),\n (\"sdpa_decode_pv (weak-32)\", pv_3, \"^\"),\n]\nprint_chip_comparison(all_kernels)", "metadata": {}, "execution_count": null, "outputs": [] diff --git a/tests/test_nb_run_kernels.py b/tests/test_nb_run_kernels.py new file mode 100644 index 0000000..c20c189 --- /dev/null +++ b/tests/test_nb_run_kernels.py @@ -0,0 +1,84 @@ +"""Tests for notebook kernel generators (demo_gen_mlir / demo_helpers). + +Covers the per-unit watermark checks and input guards that the notebook +itself does not exercise programmatically. +""" + +import sys +from pathlib import Path + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "notebooks")) + +from demo_gen_mlir import gen_rope_mlir # noqa: E402 +from demo_helpers import run_kernel_rope # noqa: E402 +from ktir_cpu.latency import HardwareConfig # noqa: E402 + + +@pytest.fixture +def hw(): + return HardwareConfig() + + +class TestRoPEAIWatermark: + """Verify AI = 3h/(4h+2) for various head/grid configs.""" + + @pytest.mark.parametrize("num_heads,grid_h,expected_ai", [ + (4, 2, 3 * 2 / (4 * 2 + 2)), # h=2, AI=0.6 + (40, 2, 3 * 20 / (4 * 20 + 2)), # h=20, AI=0.7317 + (40, 4, 3 * 10 / (4 * 10 + 2)), # h=10, AI=0.7143 + ]) + def test_ai_matches_formula(self, hw, num_heads, grid_h, expected_ai): + report = run_kernel_rope( + hw, num_heads=num_heads, seq_len=512, head_dim=128, + grid_s=2, grid_h=grid_h, tile_seq=256) + rf = report.core_roofline() + assert rf["core_AI"] == pytest.approx(expected_ai, rel=1e-3) + + +class TestRoPEPerCoreSummary: + """Verify per-core cycle structure for embarrassingly-parallel kernel.""" + + def test_memory_bound(self, hw): + report = run_kernel_rope( + hw, num_heads=4, seq_len=512, head_dim=128, + grid_s=2, grid_h=2, tile_seq=256) + summary = report.per_core_summary() + for core in summary: + assert core["memory_cycles"] > core["compute_cycles"] + + def test_all_cores_equal(self, hw): + report = run_kernel_rope( + hw, num_heads=4, seq_len=512, head_dim=128, + grid_s=2, grid_h=2, tile_seq=256) + summary = report.per_core_summary() + cycles = [c["total_cycles"] for c in summary] + assert all(c == cycles[0] for c in cycles) + + def test_no_communication(self, hw): + report = run_kernel_rope( + hw, num_heads=4, seq_len=512, head_dim=128, + grid_s=2, grid_h=2, tile_seq=256) + summary = report.per_core_summary() + for core in summary: + assert core["comm_cycles"] == 0.0 + + +class TestRoPEDivisibilityGuard: + """Verify ValueError on non-divisible inputs.""" + + def test_heads_not_divisible(self): + with pytest.raises(ValueError, match="num_heads=40 not divisible by grid_h=3"): + gen_rope_mlir(num_heads=40, seq_len=1024, head_dim=128, + grid_s=2, grid_h=3, tile_seq=256) + + def test_seq_not_divisible(self): + with pytest.raises(ValueError, match="seq_len=100 not divisible"): + gen_rope_mlir(num_heads=4, seq_len=100, head_dim=128, + grid_s=2, grid_h=2, tile_seq=256) + + def test_seq_tiles_not_divisible(self): + with pytest.raises(ValueError, match="not divisible"): + gen_rope_mlir(num_heads=4, seq_len=1024, head_dim=128, + grid_s=4, grid_h=2, tile_seq=512)