From 1bd98f2a65a4beed46cde2189ac4a31ed801b9ed Mon Sep 17 00:00:00 2001 From: Hao Yu Date: Wed, 5 Aug 2026 10:05:19 -0400 Subject: [PATCH 1/4] Add RMSNorm kernel to latency-demo roofline notebook (#165) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Integrate a parameterized RMSNorm generator and 3-config scaling study (4-core baseline, 32-core strong, 32-core weak) into the multi-kernel roofline sections alongside matmul/softmax/SDPA/paged-attention. Design choices: - Unfused standalone kernel — models the HBM pass-through between matmul stages as it appears in prefill. - 1D grid [N,1] row-partition — prefill has abundant row parallelism; hidden-dim sharding adds allreduce for zero benefit. - W is 1D [hidden_dim] per PyTorch convention — avoids inflated HBM traffic (8 KB stays LX-resident after one cold miss). Signed-off-by: Hao Yu --- notebooks/demo_gen_mlir.py | 113 +++++++++++++++++++++++++++++++++++ notebooks/demo_helpers.py | 12 +++- notebooks/latency_demo.ipynb | 8 +-- 3 files changed, 128 insertions(+), 5 deletions(-) diff --git a/notebooks/demo_gen_mlir.py b/notebooks/demo_gen_mlir.py index b93fc5c..df57fbc 100644 --- a/notebooks/demo_gen_mlir.py +++ b/notebooks/demo_gen_mlir.py @@ -355,6 +355,119 @@ def gen_rope_mlir(num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=256): }} }}""" +def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): + """Generate RMSNorm kernel: y = x * rsqrt(mean(x^2) + eps) * w. + + Two-pass embarrassingly parallel over rows. Grid = [num_cores, 1]. + Weight W is 1D [hidden_dim] — broadcast to match tile shape. + + Design decisions: + - Unfused standalone kernel (no predecessor/successor fusion) — models the + HBM-resident pass-through as it appears between matmul stages in prefill. + - 1D grid [N,1] row-partition: prefill has abundant row parallelism + (seq_len >> num_cores), so hidden-dim sharding (2x2) adds allreduce + communication for zero benefit. Matches adjacent matmul/softmax/SDPA grids. + - W is 1D [hidden_dim] per PyTorch convention (nn.Parameter(torch.ones(H))). + Avoids inflating HBM traffic — W stays LX-resident after one cold miss (8 KB). + """ + hd = hidden_dim + bs = block_size + x_view = _mem_view("x_view", "x_ptr", [n_rows, hd], [hd, 1]) + y_view = _mem_view("y_view", "y_ptr", [n_rows, hd], [hd, 1]) + w_view = _mem_view("w_view", "w_ptr", [hd], [1]) + return f"""module {{ + func.func @rmsnorm_kernel( + %x_ptr: index, %y_ptr: index, %w_ptr: index, + %n_rows: index, %N: index, %eps: f16, %BLOCK_SIZE: index + ) attributes {{grid = [{num_cores}, 1]}} {{ + %core_id = ktdp.get_compute_tile_id : index + %step = arith.constant {num_cores} : index + %c0 = arith.constant 0 : index + %c_hd = arith.constant {hd} : index + +{x_view} + +{y_view} + +{w_view} + + scf.for %row = %core_id to %n_rows step %step : index {{ + + // === Pass 1: sum of squares over hidden dim === + %zero_block = arith.constant dense<0.0> : tensor<1x{bs}xf16> + + %sq_acc = scf.for %col = %c0 to %c_hd step %BLOCK_SIZE + iter_args(%acc = %zero_block) -> tensor<1x{bs}xf16> {{ + + %x_acc = ktdp.construct_access_tile %x_view[%row, %col] {{ + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + + %x_blk = ktdp.load %x_acc : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16> + + %x_sq = arith.mulf %x_blk, %x_blk : tensor<1x{bs}xf16> + %acc_next = arith.addf %acc, %x_sq : tensor<1x{bs}xf16> + + scf.yield %acc_next : tensor<1x{bs}xf16> + }} + + // Reduce accumulator to scalar + %zero_scalar = arith.constant 0.0 : f16 + %reduce_init = tensor.splat %zero_scalar : tensor<1xf16> + %sum_sq = linalg.reduce {{ arith.addf }} + ins(%sq_acc : tensor<1x{bs}xf16>) + outs(%reduce_init : tensor<1xf16>) + dimensions = [1] + + // === Compute rstd = rsqrt(sum_sq / N + eps) === + %c0_idx = arith.constant 0 : index + %sum_scalar = tensor.extract %sum_sq[%c0_idx] : tensor<1xf16> + + %N_i32 = arith.index_cast %N : index to i32 + %N_f16 = arith.sitofp %N_i32 : i32 to f16 + %mean_sq = arith.divf %sum_scalar, %N_f16 : f16 + %mean_sq_plus_eps = arith.addf %mean_sq, %eps : f16 + %rstd_scalar = math.rsqrt %mean_sq_plus_eps : f16 + %rstd_block = tensor.splat %rstd_scalar : tensor<1x{bs}xf16> + + // === Pass 2: normalize and scale === + scf.for %col = %c0 to %c_hd step %BLOCK_SIZE {{ + + %x_acc2 = ktdp.construct_access_tile %x_view[%row, %col] {{ + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + + %w_acc = ktdp.construct_access_tile %w_view[%col] {{ + access_tile_set = affine_set<(d0) : (d0 >= 0, -d0 + {bs - 1} >= 0)>, + access_tile_order = affine_map<(d0) -> (d0)> + }} : memref<{hd}xf16> -> !ktdp.access_tile<{bs}xindex> + + %x2 = ktdp.load %x_acc2 : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16> + %w_1d = ktdp.load %w_acc : !ktdp.access_tile<{bs}xindex> -> tensor<{bs}xf16> + + %w_init = arith.constant dense<0.0> : tensor<1x{bs}xf16> + %w = linalg.broadcast ins(%w_1d : tensor<{bs}xf16>) outs(%w_init : tensor<1x{bs}xf16>) dimensions = [0] + + %x_norm = arith.mulf %x2, %rstd_block : tensor<1x{bs}xf16> + %y_blk = arith.mulf %x_norm, %w : tensor<1x{bs}xf16> + + %y_acc = ktdp.construct_access_tile %y_view[%row, %col] {{ + access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, + access_tile_order = affine_map<(d0, d1) -> (d0, d1)> + }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + + ktdp.store %y_blk, %y_acc : tensor<1x{bs}xf16>, !ktdp.access_tile<1x{bs}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 2d3058c..6eb748e 100644 --- a/notebooks/demo_helpers.py +++ b/notebooks/demo_helpers.py @@ -18,6 +18,7 @@ gen_matmul_mlir, gen_softmax_mlir, gen_rope_mlir, + gen_rmsnorm_mlir, gen_sdpa_mlir, gen_sdpa_decode_pv_mlir, gen_paged_attention_mlir, @@ -248,7 +249,6 @@ def run_kernel_softmax(hw, n_rows, row_width, num_cores, rng=None): input_ptr=rng.standard_normal((n_rows, row_width)).astype(np.float16)), 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: @@ -266,6 +266,16 @@ def run_kernel_rope(hw, num_heads, seq_len, head_dim, grid_s, grid_h, tile_seq=2 sin_ptr=sin_table, out_ptr=np.zeros((num_heads * seq_len, head_dim), dtype=np.float16))) +def run_kernel_rmsnorm(hw, n_rows, hidden_dim, num_cores, rng=None, block_size=1024): + """Generate RMSNorm MLIR, create tensors, run, return LatencyReport.""" + if rng is None: + rng = np.random.default_rng(0) + mlir = gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores, block_size) + return run_kernel(hw, mlir, "rmsnorm_kernel", + dict(x_ptr=rng.standard_normal((n_rows, hidden_dim)).astype(np.float16), + y_ptr=np.zeros((n_rows, hidden_dim), dtype=np.float16), + w_ptr=rng.uniform(0.5, 1.5, (hidden_dim,)).astype(np.float16)), + dict(n_rows=n_rows, N=hidden_dim, eps=np.float16(1e-5), BLOCK_SIZE=block_size)) def run_kernel_sdpa(hw, seq_len, head_dim, block_m, rng=None): """Generate SDPA MLIR, create tensors, run, return LatencyReport.""" diff --git a/notebooks/latency_demo.ipynb b/notebooks/latency_demo.ipynb index 8ad15e5..8b47ca2 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, 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)" + "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, gen_rmsnorm_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, run_kernel_rmsnorm,\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**, **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." + "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**, **rmsnorm**, 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**RMSNorm** computes `y = x * rsqrt(mean(x²) + ε) * w` — the normalization primitive in\nLlama-3, Granite, Gemma, etc. (appears twice per transformer layer). It is embarrassingly\nparallel over rows (tokens), with a 1D weight vector broadcast across hidden-dim blocks.\nTwo passes over the hidden dimension: (1) accumulate sum-of-squares, (2) normalize and scale.\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# --- 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)" + "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# --- RMSNorm (Granite-8B prefill slice, 4 cores) ---\nRMS_ROWS, RMS_HIDDEN, RMS_CORES = 256, 4096, 4\nprint(f\"RMSNorm: {RMS_ROWS}×{RMS_HIDDEN}, grid=[{RMS_CORES},1] = {RMS_CORES} cores\")\nrmsnorm_report = run_kernel_rmsnorm(hw, RMS_ROWS, RMS_HIDDEN, RMS_CORES, 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), (\"rmsnorm\", rmsnorm_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)\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)", + "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)\nrms_1 = run_kernel_rmsnorm(hw, RMS_ROWS, RMS_HIDDEN, 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)\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)\nrms_2 = run_kernel_rmsnorm(hw, RMS_ROWS, RMS_HIDDEN, 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)\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)\nrms_3 = run_kernel_rmsnorm(hw, RMS_ROWS * 8, RMS_HIDDEN, 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 (\"rope (4c)\", rope_1, \"o\"),\n (\"rmsnorm (4c)\", rms_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 (\"rmsnorm (strong-32)\", rms_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 (\"rmsnorm (weak-32)\", rms_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": [] From 34ff0b5c5ea2e393e0d832fba765631c9d2bc4e0 Mon Sep 17 00:00:00 2001 From: Hao Yu Date: Thu, 6 Aug 2026 09:33:05 -0400 Subject: [PATCH 2/4] Address PR #197 review: use _access_tile helper, add divisibility assert Signed-off-by: Hao Yu --- notebooks/demo_gen_mlir.py | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/notebooks/demo_gen_mlir.py b/notebooks/demo_gen_mlir.py index df57fbc..622c0c0 100644 --- a/notebooks/demo_gen_mlir.py +++ b/notebooks/demo_gen_mlir.py @@ -370,11 +370,16 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): - W is 1D [hidden_dim] per PyTorch convention (nn.Parameter(torch.ones(H))). Avoids inflating HBM traffic — W stays LX-resident after one cold miss (8 KB). """ + assert hidden_dim % block_size == 0, "hidden_dim must be divisible by block_size" hd = hidden_dim bs = block_size x_view = _mem_view("x_view", "x_ptr", [n_rows, hd], [hd, 1]) y_view = _mem_view("y_view", "y_ptr", [n_rows, hd], [hd, 1]) w_view = _mem_view("w_view", "w_ptr", [hd], [1]) + x_acc = _access_tile("x_acc", "x_view", ["%row", "%col"], [1, bs], [n_rows, hd]) + x_acc2 = _access_tile("x_acc2", "x_view", ["%row", "%col"], [1, bs], [n_rows, hd]) + w_acc = _access_tile("w_acc", "w_view", ["%col"], [bs], [hd]) + y_acc = _access_tile("y_acc", "y_view", ["%row", "%col"], [1, bs], [n_rows, hd]) return f"""module {{ func.func @rmsnorm_kernel( %x_ptr: index, %y_ptr: index, %w_ptr: index, @@ -399,10 +404,7 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): %sq_acc = scf.for %col = %c0 to %c_hd step %BLOCK_SIZE iter_args(%acc = %zero_block) -> tensor<1x{bs}xf16> {{ - %x_acc = ktdp.construct_access_tile %x_view[%row, %col] {{ - access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, - access_tile_order = affine_map<(d0, d1) -> (d0, d1)> - }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + {x_acc} %x_blk = ktdp.load %x_acc : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16> @@ -434,15 +436,9 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): // === Pass 2: normalize and scale === scf.for %col = %c0 to %c_hd step %BLOCK_SIZE {{ - %x_acc2 = ktdp.construct_access_tile %x_view[%row, %col] {{ - access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, - access_tile_order = affine_map<(d0, d1) -> (d0, d1)> - }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + {x_acc2} - %w_acc = ktdp.construct_access_tile %w_view[%col] {{ - access_tile_set = affine_set<(d0) : (d0 >= 0, -d0 + {bs - 1} >= 0)>, - access_tile_order = affine_map<(d0) -> (d0)> - }} : memref<{hd}xf16> -> !ktdp.access_tile<{bs}xindex> + {w_acc} %x2 = ktdp.load %x_acc2 : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16> %w_1d = ktdp.load %w_acc : !ktdp.access_tile<{bs}xindex> -> tensor<{bs}xf16> @@ -453,10 +449,7 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): %x_norm = arith.mulf %x2, %rstd_block : tensor<1x{bs}xf16> %y_blk = arith.mulf %x_norm, %w : tensor<1x{bs}xf16> - %y_acc = ktdp.construct_access_tile %y_view[%row, %col] {{ - access_tile_set = affine_set<(d0, d1) : (d0 >= 0, -d0 >= 0, d1 >= 0, -d1 + {bs - 1} >= 0)>, - access_tile_order = affine_map<(d0, d1) -> (d0, d1)> - }} : memref<{n_rows}x{hd}xf16> -> !ktdp.access_tile<1x{bs}xindex> + {y_acc} ktdp.store %y_blk, %y_acc : tensor<1x{bs}xf16>, !ktdp.access_tile<1x{bs}xindex> From c87cd523467bc51a93ac3562ac9c66b988895844 Mon Sep 17 00:00:00 2001 From: Hao Yu Date: Wed, 19 Aug 2026 09:39:33 -0400 Subject: [PATCH 3/4] Added f16-to-f32 and back conversions to address overflow in acc-reduce phase of rmsnorm kernel 1. fix(rmsnorm mlir): promote accumulation-reduction block to f32-based 2. feature(latency): charge SIMD cost for extf, truncf, and splat (reflecting Spyre Rapid Core specs) 3. LX-residency docstring claim 4. explanation for bottleneck=compute at 4 cores Signed-off-by: Hao Yu --- ktir_cpu/dialects/arith_ops.py | 15 ++++++++--- ktir_cpu/dialects/tensor_ops.py | 3 ++- notebooks/demo_gen_mlir.py | 44 +++++++++++++++++++-------------- notebooks/latency_demo.ipynb | 2 +- 4 files changed, 39 insertions(+), 25 deletions(-) diff --git a/ktir_cpu/dialects/arith_ops.py b/ktir_cpu/dialects/arith_ops.py index fdae2ca..04dbe42 100644 --- a/ktir_cpu/dialects/arith_ops.py +++ b/ktir_cpu/dialects/arith_ops.py @@ -220,12 +220,19 @@ def arith__constant(op, context, env): return value -# Cast ops — no latency category (Pattern A.4 — cast cluster) -# extsi stays bespoke: uses a lambda that can't be expressed as an ArithOps method. - -_CAST_UNARY_OPS = { +# Float cast ops — SIMD pipeline cost (Pattern A.4) +_FLOAT_CAST_OPS = { "arith.extf": (ArithOps.extf, np.float32), "arith.truncf": (ArithOps.truncf, None), +} +for _name, (_fn, _sfn) in _FLOAT_CAST_OPS.items(): + @register(_name, latency_category=LC.COMPUTE_FLOAT) + def _(op, context, env, _fn=_fn, _sfn=_sfn): + return _unary(op, context, _fn, _sfn) + +# Integer/index cast ops — no latency category (Pattern A.4 — cast cluster) +# extsi stays bespoke: uses a lambda that can't be expressed as an ArithOps method. +_CAST_UNARY_OPS = { "arith.extui": (ArithOps.extui, int), "arith.trunci": (ArithOps.trunci, int), "arith.fptosi": (ArithOps.fptosi, int), diff --git a/ktir_cpu/dialects/tensor_ops.py b/ktir_cpu/dialects/tensor_ops.py index 5e9e27b..d5891a7 100644 --- a/ktir_cpu/dialects/tensor_ops.py +++ b/ktir_cpu/dialects/tensor_ops.py @@ -24,6 +24,7 @@ from ..ir_types import Operation, Tile from ..parser_utils import find_ssa_names, parse_tensor_or_memref_type from .registry import register, register_parser +from ..latency import LatencyCategory as LC def _infer_splat_shape(context: CoreContext) -> Optional[Tuple[int, ...]]: @@ -54,7 +55,7 @@ def tensor__empty(op, context, env): return Tile(data, dtype_str, shape) -@register("tensor.splat") +@register("tensor.splat", latency_category=LC.COMPUTE_FLOAT) def tensor__splat(op, context, env): scalar = context.get_value(op.operands[0]) diff --git a/notebooks/demo_gen_mlir.py b/notebooks/demo_gen_mlir.py index 622c0c0..ec61062 100644 --- a/notebooks/demo_gen_mlir.py +++ b/notebooks/demo_gen_mlir.py @@ -368,7 +368,8 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): (seq_len >> num_cores), so hidden-dim sharding (2x2) adds allreduce communication for zero benefit. Matches adjacent matmul/softmax/SDPA grids. - W is 1D [hidden_dim] per PyTorch convention (nn.Parameter(torch.ones(H))). - Avoids inflating HBM traffic — W stays LX-resident after one cold miss (8 KB). + Loaded per-block in pass 2; small enough (8 KB) to be LX-resident on real + hardware, though the charge-model estimator costs every load from HBM. """ assert hidden_dim % block_size == 0, "hidden_dim must be divisible by block_size" hd = hidden_dim @@ -398,39 +399,44 @@ def gen_rmsnorm_mlir(n_rows, hidden_dim, num_cores=4, block_size=1024): scf.for %row = %core_id to %n_rows step %step : index {{ - // === Pass 1: sum of squares over hidden dim === - %zero_block = arith.constant dense<0.0> : tensor<1x{bs}xf16> + // === Pass 1: sum of squares over hidden dim (f32 accumulator) === + %zero_f32 = arith.constant 0.0 : f32 + %zero_block = tensor.splat %zero_f32 : tensor<1x{bs}xf32> %sq_acc = scf.for %col = %c0 to %c_hd step %BLOCK_SIZE - iter_args(%acc = %zero_block) -> tensor<1x{bs}xf16> {{ + iter_args(%acc = %zero_block) -> tensor<1x{bs}xf32> {{ {x_acc} %x_blk = ktdp.load %x_acc : !ktdp.access_tile<1x{bs}xindex> -> tensor<1x{bs}xf16> %x_sq = arith.mulf %x_blk, %x_blk : tensor<1x{bs}xf16> - %acc_next = arith.addf %acc, %x_sq : tensor<1x{bs}xf16> + %x_sq_f32 = arith.extf %x_sq : tensor<1x{bs}xf16> to tensor<1x{bs}xf32> + %acc_next = arith.addf %acc, %x_sq_f32 : tensor<1x{bs}xf32> - scf.yield %acc_next : tensor<1x{bs}xf16> + scf.yield %acc_next : tensor<1x{bs}xf32> }} - // Reduce accumulator to scalar - %zero_scalar = arith.constant 0.0 : f16 - %reduce_init = tensor.splat %zero_scalar : tensor<1xf16> + // Reduce f32 accumulator + %reduce_init = tensor.splat %zero_f32 : tensor<1xf32> %sum_sq = linalg.reduce {{ arith.addf }} - ins(%sq_acc : tensor<1x{bs}xf16>) - outs(%reduce_init : tensor<1xf16>) + ins(%sq_acc : tensor<1x{bs}xf32>) + outs(%reduce_init : tensor<1xf32>) dimensions = [1] - // === Compute rstd = rsqrt(sum_sq / N + eps) === - %c0_idx = arith.constant 0 : index - %sum_scalar = tensor.extract %sum_sq[%c0_idx] : tensor<1xf16> - + // === Compute rstd = rsqrt(sum_sq / N + eps) on tensor<1xf32> === %N_i32 = arith.index_cast %N : index to i32 - %N_f16 = arith.sitofp %N_i32 : i32 to f16 - %mean_sq = arith.divf %sum_scalar, %N_f16 : f16 - %mean_sq_plus_eps = arith.addf %mean_sq, %eps : f16 - %rstd_scalar = math.rsqrt %mean_sq_plus_eps : f16 + %N_f32 = arith.sitofp %N_i32 : i32 to f32 + %N_t = tensor.splat %N_f32 : tensor<1xf32> + %eps_f32 = arith.extf %eps : f16 to f32 + %eps_t = tensor.splat %eps_f32 : tensor<1xf32> + + %mean_sq = arith.divf %sum_sq, %N_t : tensor<1xf32> + %mean_sq_plus_eps = arith.addf %mean_sq, %eps_t : tensor<1xf32> + %rstd_f32 = math.rsqrt %mean_sq_plus_eps : tensor<1xf32> + %rstd_f16 = arith.truncf %rstd_f32 : tensor<1xf32> to tensor<1xf16> + %c0_idx = arith.constant 0 : index + %rstd_scalar = tensor.extract %rstd_f16[%c0_idx] : tensor<1xf16> %rstd_block = tensor.splat %rstd_scalar : tensor<1x{bs}xf16> // === Pass 2: normalize and scale === diff --git a/notebooks/latency_demo.ipynb b/notebooks/latency_demo.ipynb index 8b47ca2..1074445 100644 --- a/notebooks/latency_demo.ipynb +++ b/notebooks/latency_demo.ipynb @@ -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**, **rope**, **rmsnorm**, 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**RMSNorm** computes `y = x * rsqrt(mean(x²) + ε) * w` — the normalization primitive in\nLlama-3, Granite, Gemma, etc. (appears twice per transformer layer). It is embarrassingly\nparallel over rows (tokens), with a 1D weight vector broadcast across hidden-dim blocks.\nTwo passes over the hidden dimension: (1) accumulate sum-of-squares, (2) normalize and scale.\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 six 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**, **rmsnorm**, 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**RMSNorm** computes `y = x * rsqrt(mean(x²) + ε) * w` — the normalization primitive in\nLlama-3, Granite, Gemma, etc. (appears twice per transformer layer). It is embarrassingly\nparallel over rows (tokens), with a 1D weight vector broadcast across hidden-dim blocks.\nTwo passes over the hidden dimension: (1) accumulate sum-of-squares, (2) normalize and scale.\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", From f8b0afcc38bd53783658a95f3a742e9df527f94f Mon Sep 17 00:00:00 2001 From: Hao Yu Date: Fri, 21 Aug 2026 13:27:51 -0400 Subject: [PATCH 4/4] Minor fixes requested from reviews - F3: rmsnorm roofline docstring in the notebook - F4: keep the zero-latency premise for type-casting and splat ops Signed-off-by: Hao Yu --- ktir_cpu/dialects/arith_ops.py | 15 ++++----------- ktir_cpu/dialects/tensor_ops.py | 3 +-- notebooks/demo_helpers.py | 2 +- notebooks/latency_demo.ipynb | 32 +++++++++++--------------------- 4 files changed, 17 insertions(+), 35 deletions(-) diff --git a/ktir_cpu/dialects/arith_ops.py b/ktir_cpu/dialects/arith_ops.py index 04dbe42..fdae2ca 100644 --- a/ktir_cpu/dialects/arith_ops.py +++ b/ktir_cpu/dialects/arith_ops.py @@ -220,19 +220,12 @@ def arith__constant(op, context, env): return value -# Float cast ops — SIMD pipeline cost (Pattern A.4) -_FLOAT_CAST_OPS = { - "arith.extf": (ArithOps.extf, np.float32), - "arith.truncf": (ArithOps.truncf, None), -} -for _name, (_fn, _sfn) in _FLOAT_CAST_OPS.items(): - @register(_name, latency_category=LC.COMPUTE_FLOAT) - def _(op, context, env, _fn=_fn, _sfn=_sfn): - return _unary(op, context, _fn, _sfn) - -# Integer/index cast ops — no latency category (Pattern A.4 — cast cluster) +# Cast ops — no latency category (Pattern A.4 — cast cluster) # extsi stays bespoke: uses a lambda that can't be expressed as an ArithOps method. + _CAST_UNARY_OPS = { + "arith.extf": (ArithOps.extf, np.float32), + "arith.truncf": (ArithOps.truncf, None), "arith.extui": (ArithOps.extui, int), "arith.trunci": (ArithOps.trunci, int), "arith.fptosi": (ArithOps.fptosi, int), diff --git a/ktir_cpu/dialects/tensor_ops.py b/ktir_cpu/dialects/tensor_ops.py index d5891a7..5e9e27b 100644 --- a/ktir_cpu/dialects/tensor_ops.py +++ b/ktir_cpu/dialects/tensor_ops.py @@ -24,7 +24,6 @@ from ..ir_types import Operation, Tile from ..parser_utils import find_ssa_names, parse_tensor_or_memref_type from .registry import register, register_parser -from ..latency import LatencyCategory as LC def _infer_splat_shape(context: CoreContext) -> Optional[Tuple[int, ...]]: @@ -55,7 +54,7 @@ def tensor__empty(op, context, env): return Tile(data, dtype_str, shape) -@register("tensor.splat", latency_category=LC.COMPUTE_FLOAT) +@register("tensor.splat") def tensor__splat(op, context, env): scalar = context.get_value(op.operands[0]) diff --git a/notebooks/demo_helpers.py b/notebooks/demo_helpers.py index 6eb748e..819bac3 100644 --- a/notebooks/demo_helpers.py +++ b/notebooks/demo_helpers.py @@ -82,7 +82,7 @@ def print_hw_config(hw: HardwareConfig): print(f" {'cores_active':<14} {'systolic':>10} {'SIMD':>10}") for nc in [4, 32]: bw = chip_bw_per_cycle / nc - print(f" {nc:<14} {hw.systolic_flops_per_cycle / bw:>7.0f} F/B {hw.simd_elements_per_cycle / bw:>7.0f} F/B") + print(f" {nc:<14} {hw.systolic_flops_per_cycle / bw:>7.2f} F/B {hw.simd_elements_per_cycle / bw:>7.2f} F/B") def print_core_roofline(report, hw: HardwareConfig): diff --git a/notebooks/latency_demo.ipynb b/notebooks/latency_demo.ipynb index 1074445..09304ad 100644 --- a/notebooks/latency_demo.ipynb +++ b/notebooks/latency_demo.ipynb @@ -3,7 +3,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "# KTIR Latency Estimation Demo\n\nThis notebook shows how to use `ktir_cpu` to get cycle-approximate latency estimates\nfor KTIR kernels — no Spyre hardware required.\n\n**What we cover:**\n1. [Matmul — load, run, report](#1.-Matmul-—-load,-run,-report)\n2. [Per-core breakdown](#2.-Per-core-breakdown)\n3. [Softmax — compare with matmul](#3.-Softmax-—-compare-with-matmul)\n4. [Multi-kernel roofline](#4.-Multi-kernel-roofline)\n5. [Per-op trace: paged attention](#5.-Per-op-trace:-paged-attention)\n6. [Chip-level analysis](#6.-Chip-level-analysis-—-the-whole-chip-as-one-unit)\n7. [Tuning HardwareConfig](#7.-Tuning-HardwareConfig)\n8. [MLIR frontend parser (optional)](#8.-MLIR-frontend-parser-(optional))\n9. [Cross-core communication (comm) view](#9.-Cross-core-communication-(comm)-view)\n10. [Multi-core decode SDPA (P@V)](#10.-Multi-core-decode-SDPA-(P@V)-—-splitting-a-contraction-across-cores)" + "source": "# KTIR Latency Estimation Demo\n\nThis notebook shows how to use `ktir_cpu` to get cycle-approximate latency estimates\nfor KTIR kernels — no Spyre hardware required.\n\n**What we cover:**\n1. [Matmul — load, run, report](#1.-Matmul-—-load,-run,-report)\n2. [Per-core breakdown](#2.-Per-core-breakdown)\n3. [Softmax — compare with matmul](#3.-Softmax-—-compare-with-matmul)\n4. [Multi-kernel roofline](#4.-Multi-kernel-roofline)\n5. [Per-op trace: paged attention](#5.-Per-op-trace:-paged-attention)\n6. [Chip-level analysis](#6.-Chip-level-analysis-—-the-whole-chip-as-one-unit)\n7. [Tuning HardwareConfig](#7.-Tuning-HardwareConfig)\n8. [MLIR frontend parser](#8.-MLIR-frontend-parser)\n9. [Cross-core communication view](#9.-Cross-core-communication-view)\n10. [Multi-core decode SDPA — splitting a contraction across cores](#10.-Multi-core-decode-SDPA-—-splitting-a-contraction-across-cores)" }, { "cell_type": "markdown", @@ -120,9 +120,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": [ - "---\n\n## 2. Per-core breakdown\n\nIn this notebook, serial execution model of compute cores of the accelerator is assumed. The elapsed time of core is split into 3 categories: compute, memroy (access), and (core-to-core) comm (unication). The 'Per-core' her is from the perspective of a single core, when communication is not modeled, the compute-capacity of a core is defined for a given HW configuration, the data-move capacity of a core is a fraction of the total HBM bandwidth that are shared among all active cores (NOT total cores in the accelerator)." - ] + "source": "---\n\n## 2. Per-core breakdown\n\nThis notebook assumes a serial execution model for the accelerator's compute cores. Each core's elapsed time is split into three categories: compute, memory (HBM access), and comm (core-to-core communication). The per-core view is from the perspective of a single core: compute capacity is defined by the HW configuration, and data-movement capacity is a fraction of total HBM bandwidth shared among all active cores (not total cores in the accelerator)." }, { "cell_type": "code", @@ -176,7 +174,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "---\n\n## 4. Multi-kernel roofline\n\nRun six 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**, **rmsnorm**, 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**RMSNorm** computes `y = x * rsqrt(mean(x²) + ε) * w` — the normalization primitive in\nLlama-3, Granite, Gemma, etc. (appears twice per transformer layer). It is embarrassingly\nparallel over rows (tokens), with a 1D weight vector broadcast across hidden-dim blocks.\nTwo passes over the hidden dimension: (1) accumulate sum-of-squares, (2) normalize and scale.\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 seven kernels through the latency estimator and plot them together:\n- **matmul**, **sdpa**, **sdpa_decode_pv**, **paged_attn** — systolic-dominant (matmul ops)\n- **softmax**, **rope**, **rmsnorm** — 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**RMSNorm** computes `y = x * rsqrt(mean(x²) + ε) * w` — the normalization primitive in\nLlama-3, Granite, Gemma, etc. (appears twice per transformer layer). It is embarrassingly\nparallel over rows (tokens), with a 1D weight vector broadcast across hidden-dim blocks.\nTwo passes over the hidden dimension: (1) accumulate sum-of-squares, (2) normalize and scale.\nNote that RMSNorm reports `bottleneck=compute` at 4 cores — this is physically correct, not\na bug. The per-core SIMD ridge shifts with core count: at 4 cores, per-core HBM bandwidth is\n`1.024 TB/s ÷ 4 = 256 GB/s`, giving a SIMD ridge of `64 elems/cycle ÷ 256 B/cycle = 0.25 F/B`.\nRMSNorm's AI (~0.53) exceeds this ridge, so it is compute-bound. At 32 cores the ridge rises\nto 2.0 F/B and the same kernel becomes memory-bound — the scaling table in Section 6 shows\nthis crossover.\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 splits work along the contraction axis — section 10 takes it apart.\n`sdpa` comes out memory-bound below too, but for it that is a `block_m` choice: raise\n`block_m` to 32 and it crosses the ridge. Decode P@V cannot be moved — its arithmetic\nintensity is capped by the two tile widths alone — and it is the only one whose `comm`\ncolumn 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", @@ -235,7 +233,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "---\n\n## 6. Chip-level analysis — the whole chip as one unit\n\nAbove covers the **per-core** analysis (the critical core vs its own ceiling). What follows is a\n**separate analysis at a different granularity**: it collapses the chip to a single unit (the way Nsight\nviews a device) and asks *how much of the **entire chip** did this kernel use, and what bounds it* — a\nquestion per-core cannot answer.\n\n**Chip-level analysis** is most common scenario for a relative weak AI accelerators to host the kernels of ever growing LLM models, thus complements the **per-core analysis**. A core can have **~99% attainment**\n(`core_attainment`) while the **chip is mostly idle**. For instance, chip-level analysis makes visible:\n\n- `grid_coverage` / `mean_core_active_frac` — how many cores actually did work / for how much of the time.\n- `dram_throughput` — whether the kernel saturates the **whole chip's** HBM bandwidth (device memory-bound).\n- `compute_throughput` — the whole chip's compute utilisation vs its flat peak.\n\nConvention: plain NCU-style names (`AI`, `compute_throughput`, …) are **chip-level**; a `core_` prefix is\n**per-core**. The two are never compared on the same axes — their denominators differ." + "source": "---\n\n## 6. Chip-level analysis — the whole chip as one unit\n\nAbove covers the **per-core** analysis (the critical core vs its own ceiling). What follows is a\n**separate analysis at a different granularity**: it collapses the chip to a single unit (the way Nsight\nviews a device) and asks *how much of the **entire chip** did this kernel use, and what bounds it* — a\nquestion per-core cannot answer.\n\n**Chip-level analysis** complements the per-core view and is the most relevant framing for relatively\nweak AI accelerators hosting ever-growing LLM models. A core can have **~99% attainment**\n(`core_attainment`) while the **chip is mostly idle**. For instance, chip-level analysis makes visible:\n\n- `grid_coverage` / `mean_core_active_frac` — how many cores actually did work / for how much of the time.\n- `dram_throughput` — whether the kernel saturates the **whole chip's** HBM bandwidth (device memory-bound).\n- `compute_throughput` — the whole chip's compute utilisation vs its flat peak.\n\nConvention: plain NCU-style names (`AI`, `compute_throughput`, …) are **chip-level**; a `core_` prefix is\n**per-core**. The two are never compared on the same axes — their denominators differ." }, { "cell_type": "code", @@ -272,7 +270,7 @@ }, { "cell_type": "code", - "source": "print(\"### Scaling experiment results\\n\")\nprint_scaling_table(all_kernels)\n\n# Ridge = per-core peak (F/cycle) / per-core BW (B/cycle)\n# = flops_per_cycle / (chip_bw_per_cycle / cores_active)\nchip_bw_per_cycle = hw.hbm_bw_chip / hw.clock_hz\nprint(f\"\\nRidge (systolic / SIMD) by cores_active:\")\nfor nc in [4, 32]:\n bw = chip_bw_per_cycle / nc\n sys_ridge = hw.systolic_flops_per_cycle / bw\n simd_ridge = hw.simd_elements_per_cycle / bw\n print(f\" {nc} cores: systolic = {sys_ridge:.0f} F/B, SIMD = {simd_ridge:.1f} F/B\")", + "source": "print(\"### Scaling experiment results\\n\")\nprint_scaling_table(all_kernels)\n\n# Ridge = per-core peak (F/cycle) / per-core BW (B/cycle)\n# = flops_per_cycle / (chip_bw_per_cycle / cores_active)\nchip_bw_per_cycle = hw.hbm_bw_chip / hw.clock_hz\nprint(f\"\\nRidge (systolic / SIMD) by cores_active:\")\nfor nc in [4, 32]:\n bw = chip_bw_per_cycle / nc\n sys_ridge = hw.systolic_flops_per_cycle / bw\n simd_ridge = hw.simd_elements_per_cycle / bw\n print(f\" {nc} cores: systolic = {sys_ridge:.2f} F/B, SIMD = {simd_ridge:.2f} F/B\")", "metadata": {}, "execution_count": null, "outputs": [] @@ -299,20 +297,12 @@ } }, "outputs": [], - "source": "# Contrast: same softmax kernel, different HBM bandwidth\nhw_lo = HardwareConfig(\n num_cores=32, clock_ghz=1.0, lx_size_mb=2,\n hbm_bandwidth_tb_s=0.256, # 256 GB/s\n ring_bandwidth_tb_s=0.064,\n simd_elements_per_cycle=64, systolic_rows=8, transcendental_penalty=4,\n)\nhw_hi = HardwareConfig(\n num_cores=32, clock_ghz=1.0, lx_size_mb=2,\n hbm_bandwidth_tb_s=1.024, # 1 TB/s\n ring_bandwidth_tb_s=0.064,\n simd_elements_per_cycle=64, systolic_rows=8, transcendental_penalty=4,\n)\n\nprint(\"=== Low BW (256 GB/s) ===\")\nprint_hw_config(hw_lo)\nprint(\"\\n=== High BW (1 TB/s) ===\")\nprint_hw_config(hw_hi)\n\n# Run softmax on both configs (4 cores active)\nsm_lo = run_kernel_softmax(hw_lo, SM_ROWS, SM_WIDTH, SM_CORES, rng)\nsm_hi = run_kernel_softmax(hw_hi, SM_ROWS, SM_WIDTH, SM_CORES, rng)\n\nprint(f\"\\n{'Config':<12} {'Cycles':>10} {'Bottleneck':>12} {'SIMD ridge (4c)':>16}\")\nprint(f\"{'─'*12} {'─'*10} {'─'*12} {'─'*16}\")\nfor label, r, h in [(\"256 GB/s\", sm_lo, hw_lo), (\"1 TB/s\", sm_hi, hw_hi)]:\n bw_per_core = h.hbm_bw_chip / h.clock_hz / SM_CORES\n ridge = h.simd_elements_per_cycle / bw_per_core\n print(f\"{label:<12} {r.kernel_cycles:>7.0f} cy {r.bottleneck:>12} {ridge:>13.1f} F/B\")\nspeedup = sm_lo.kernel_cycles / sm_hi.kernel_cycles\nprint(f\"\\nSpeedup: {speedup:.2f}× (4× BW → {speedup:.1f}× faster)\")" + "source": "# Contrast: same softmax kernel, different HBM bandwidth\nhw_lo = HardwareConfig(\n num_cores=32, clock_ghz=1.0, lx_size_mb=2,\n hbm_bandwidth_tb_s=0.256, # 256 GB/s\n ring_bandwidth_tb_s=0.064,\n simd_elements_per_cycle=64, systolic_rows=8, transcendental_penalty=4,\n)\nhw_hi = HardwareConfig(\n num_cores=32, clock_ghz=1.0, lx_size_mb=2,\n hbm_bandwidth_tb_s=1.024, # 1 TB/s\n ring_bandwidth_tb_s=0.064,\n simd_elements_per_cycle=64, systolic_rows=8, transcendental_penalty=4,\n)\n\nprint(\"=== Low BW (256 GB/s) ===\")\nprint_hw_config(hw_lo)\nprint(\"\\n=== High BW (1 TB/s) ===\")\nprint_hw_config(hw_hi)\n\n# Run softmax on both configs (4 cores active)\nsm_lo = run_kernel_softmax(hw_lo, SM_ROWS, SM_WIDTH, SM_CORES, rng)\nsm_hi = run_kernel_softmax(hw_hi, SM_ROWS, SM_WIDTH, SM_CORES, rng)\n\nprint(f\"\\n{'Config':<12} {'Cycles':>10} {'Bottleneck':>12} {'SIMD ridge (4c)':>16}\")\nprint(f\"{'─'*12} {'─'*10} {'─'*12} {'─'*16}\")\nfor label, r, h in [(\"256 GB/s\", sm_lo, hw_lo), (\"1 TB/s\", sm_hi, hw_hi)]:\n bw_per_core = h.hbm_bw_chip / h.clock_hz / SM_CORES\n ridge = h.simd_elements_per_cycle / bw_per_core\n print(f\"{label:<12} {r.kernel_cycles:>7.0f} cy {r.bottleneck:>12} {ridge:>13.2f} F/B\")\nspeedup = sm_lo.kernel_cycles / sm_hi.kernel_cycles\nprint(f\"\\nSpeedup: {speedup:.2f}× (4× BW → {speedup:.1f}× faster)\")" }, { "cell_type": "markdown", "metadata": {}, - "source": [ - "---\n\n## 8. MLIR frontend parser (optional)\n", - "\n", - "As a well formed intermediate representation built in the MLIR eco-system, on top of the standard parser from the KTIR-defining project, the KTIR MLIR-frontend, kernels written in KTIR can be parsed with a light-weight parser built out of regex (Python regular expression).\n", - "By default `KTIRInterpreter` uses the built-in regex parser.\n", - "The MLIR frontend uses the real MLIR C++ bindings for full-fidelity parsing.\n", - "\n", - "See the [README](../README.md#mlir-frontend-bindings-optional) for build and install instructions." - ] + "source": "---\n\n## 8. MLIR frontend parser\n\nKTIR kernels can be parsed two ways: a lightweight regex-based parser (the default in\n`KTIRInterpreter`) and a full-fidelity MLIR frontend parser backed by the real MLIR C++\nbindings (`mlir_ktdp`). Both produce identical IR for well-formed inputs; the frontend\nadditionally validates against the dialect's type system and rejects malformed MLIR that\nthe regex path would silently accept.\n\nSee the [README](../README.md#mlir-frontend-bindings-optional) for build and install instructions." }, { "cell_type": "code", @@ -384,7 +374,7 @@ }, { "cell_type": "markdown", - "source": "## 9. Cross-core communication (comm) view\n\nEvery kernel above is embarrassingly parallel except `sdpa_decode_pv`, which splits a contraction and\ntherefore folds partial sums across cores; for the rest the `comm` column reads `0`. Communication *is* costed by\nthe model (`ktdp.inter_tile_reduce` → `comm_cycles`, folded into `total_cycles`; `bottleneck` recognises\n`comm`) — it just isn't exercised or surfaced here.\n\nThis section loads the in-repo cross-core all-reduce examples (`examples/ktir/ring_reduce*.mlir`) and surfaces\ncommunication: a per-core `compute / memory / comm` breakdown, plus the **critical-path core**'s\n`comm_fraction`, `comm_time`, and `comm_bytes`. Metrics are taken on the critical core (`max(total_cycles)`)\n— the wall-clock core `kernel_time` comes from — so this is *exposed communication* (PyTorch **HTA** framing):\ncomm on a non-critical core is hidden behind that core's own work and adds no wall time.", + "source": "## 9. Cross-core communication view\n\nEvery kernel above is embarrassingly parallel except `sdpa_decode_pv`, which splits a contraction and\ntherefore folds partial sums across cores; for the rest the `comm` column reads `0`. Communication *is* costed by\nthe model (`ktdp.inter_tile_reduce` → `comm_cycles`, folded into `total_cycles`; `bottleneck` recognises\n`comm`) — it just isn't exercised or surfaced here.\n\nThis section loads the in-repo cross-core all-reduce examples (`examples/ktir/ring_reduce*.mlir`) and surfaces\ncommunication: a per-core `compute / memory / comm` breakdown, plus the **critical-path core**'s\n`comm_fraction`, `comm_time`, and `comm_bytes`. Metrics are taken on the critical core (`max(total_cycles)`)\n— the wall-clock core `kernel_time` comes from — so this is *exposed communication* (PyTorch **HTA** framing):\ncomm on a non-critical core is hidden behind that core's own work and adds no wall time.", "metadata": {} }, { @@ -410,7 +400,7 @@ }, { "cell_type": "markdown", - "source": "---\n\n## 10. Multi-core decode SDPA (P@V) — splitting a contraction across cores\n\nSections 1–8 all parallelise the same way: **partition the output**, so each core owns a disjoint\nslice of the result and writes it alone — which is why their `comm` column reads `0`. Section 9\nshowed the model *does* charge communication, using standalone all-reduce examples. This section is\nthe two joined up: a kernel shape whose only viable decomposition forces a cross-core fold.\n\n**The kernel.** `C[q_per_kv, head_dim] = A[q_per_kv, kv_len] @ B[kv_len, head_dim]` is the `P @ V`\nstage of one decode step — `A` is the attention probabilities of the query heads sharing a KV head,\n`B` is that head's cached V. Shapes follow Granite-8B's grouped-query ratio (32 query heads over 8\nKV heads at `head_dim` 128, hence `q_per_kv = 4`) with an 8192-token KV cache.\n\n**Why decode is the shape that has to cross cores.** The output is `4 × 128` = 512 elements;\npartitioning it 32 ways would leave 16 elements per core. The only extent big enough to split is\n`kv_len` — and that is the *contraction*, so each core ends up holding a partial sum that has to be\nadded to the others. The grid is `[out_split, k_split]` = `[2, 16]`: the 16 cores sharing an output\nslice fold through `ktdp.inter_tile_produce` / `ktdp.inter_tile_reduce`, and one core per group stores.\n\nThree things to read off below:\n\n- **Memory-bound on the systolic array, and structurally so.** Arithmetic intensity is\n `m·n / (m + n)`, the *harmonic mean* of the two tile widths, because the contraction extent\n cancels out of `2mnk / (2(mk + kn))`. With `m = q_per_kv = 4` that caps at ~3.8 flop/byte.\n `sdpa` in section 4 is memory-bound as plotted as well, but that is a tiling choice — its AI\n rises with `block_m` and crosses the ridge at `block_m = 32`. Nothing moves decode P@V's.\n- **Splitting harder cannot fix it.** Since `k` cancels, the three configurations land on one\n x-coordinate, while the per-core systolic ridge *rises* with the core count (bandwidth is shared).\n At 4 cores the kernel sits just under the ridge; at 32 it is nowhere near it. Strong scaling\n shrinks only the compute term — chip bytes and chip bandwidth are both fixed — so it buys ~1.5×.\n- **`comm` is visible but never the bottleneck.** The fold is charged\n `(cores − 1) · payload / ring_bw`: priced across every *active core*, not across the 16-core reduce\n group, so it tracks the grid rather than the fan-in.", + "source": "---\n\n## 10. Multi-core decode SDPA (P@V) — splitting a contraction across cores\n\nMost kernels in this notebook parallelise by **partitioning the output**, so each core owns a\ndisjoint slice and writes it alone — which is why their `comm` column reads `0`. The exception\nis `sdpa_decode_pv`, which splits work along the contraction axis and requires a cross-core fold.\nSection 9 showed the model charges communication using standalone all-reduce examples. This section\njoins the two: a real kernel shape whose only viable decomposition forces a cross-core fold.\n\n**The kernel.** `C[q_per_kv, head_dim] = A[q_per_kv, kv_len] @ B[kv_len, head_dim]` is the `P @ V`\nstage of one decode step — `A` is the attention probabilities of the query heads sharing a KV head,\n`B` is that head's cached V. Shapes follow Granite-8B's grouped-query ratio (32 query heads over 8\nKV heads at `head_dim` 128, hence `q_per_kv = 4`) with an 8192-token KV cache.\n\n**Why decode is the shape that has to cross cores.** The output is `4 × 128` = 512 elements;\npartitioning it 32 ways would leave 16 elements per core. The only extent big enough to split is\n`kv_len` — and that is the *contraction*, so each core ends up holding a partial sum that has to be\nadded to the others. The grid is `[out_split, k_split]` = `[2, 16]`: the 16 cores sharing an output\nslice fold through `ktdp.inter_tile_produce` / `ktdp.inter_tile_reduce`, and one core per group stores.\n\nThree things to read off below:\n\n- **Memory-bound on the systolic array, and structurally so.** Arithmetic intensity is\n `m·n / (m + n)`, the *harmonic mean* of the two tile widths, because the contraction extent\n cancels out of `2mnk / (2(mk + kn))`. With `m = q_per_kv = 4` that caps at ~3.8 flop/byte.\n `sdpa` in section 4 is memory-bound as plotted as well, but that is a tiling choice — its AI\n rises with `block_m` and crosses the ridge at `block_m = 32`. Nothing moves decode P@V's.\n- **Splitting harder cannot fix it.** Since `k` cancels, the three configurations land on one\n x-coordinate, while the per-core systolic ridge *rises* with the core count (bandwidth is shared).\n At 4 cores the kernel sits just under the ridge; at 32 it is nowhere near it. Strong scaling\n shrinks only the compute term — chip bytes and chip bandwidth are both fixed — so it buys ~1.5×.\n- **`comm` is visible but never the bottleneck.** The fold is charged\n `(cores − 1) · payload / ring_bw`: priced across every *active core*, not across the 16-core reduce\n group, so it tracks the grid rather than the fan-in.", "metadata": {} }, { @@ -430,7 +420,7 @@ { "cell_type": "markdown", "metadata": {}, - "source": "---\n\n## Summary\n\nThroughout 10 curated exercises, the notebook demonstrates the usage interface of `ktir-cpu` for verify the correctiness of LLM kernels written in ktir mlir intermediate representation, and for estimating the latency and throughput performances in the framework of roofline models. Below is a complete list of the demonstrated API usages.\n\n```python\nfrom ktir_cpu.interpreter import KTIRInterpreter\nfrom ktir_cpu.latency import HardwareConfig\n\nhw = HardwareConfig() # or HardwareConfig(clock_ghz=1.5, ...)\ninterp = KTIRInterpreter(latency_config=hw) # add parser=MLIRFrontendParser() for MLIR\ninterp.load(\"path/to/kernel.mlir\") # or inline MLIR text\n\ninterp.execute_function(\"kernel_name\", arg=tensor, ...)\n\nreport = interp.get_latency_report()\nprint(report) # human-readable table + roofline\nreport.kernel_cycles # float\nreport.kernel_time_us # float\nreport.bottleneck # 'compute' | 'memory' | 'comm'\nreport.per_core_summary() # list[dict] — per-core breakdown\nreport.chip_roofline() # dict — AI, compute/dram_throughput, attainment, grid_coverage\nreport.core_roofline() # dict — core_AI, core_attainment, units, ...\nreport.roofline() # dict — chip + core merged\n```" + "source": "---\n\n## Summary\n\nThroughout 10 curated exercises, this notebook demonstrates the `ktir-cpu` interface for verifying the correctness of LLM kernels written in the KTIR MLIR intermediate representation, and for estimating latency and throughput in the framework of roofline models. Below is a complete list of the demonstrated API usages.\n\n```python\nfrom ktir_cpu.interpreter import KTIRInterpreter\nfrom ktir_cpu.latency import HardwareConfig\n\nhw = HardwareConfig() # or HardwareConfig(clock_ghz=1.5, ...)\ninterp = KTIRInterpreter(latency_config=hw) # add parser=MLIRFrontendParser() for MLIR\ninterp.load(\"path/to/kernel.mlir\") # or inline MLIR text\n\ninterp.execute_function(\"kernel_name\", arg=tensor, ...)\n\nreport = interp.get_latency_report()\nprint(report) # human-readable table + roofline\nreport.kernel_cycles # float\nreport.kernel_time_us # float\nreport.bottleneck # 'compute' | 'memory' | 'comm'\nreport.per_core_summary() # list[dict] — per-core breakdown\nreport.chip_roofline() # dict — AI, compute/dram_throughput, attainment, grid_coverage\nreport.core_roofline() # dict — core_AI, core_attainment, units, ...\nreport.roofline() # dict — chip + core merged\n```" } ], "metadata": { @@ -454,4 +444,4 @@ }, "nbformat": 4, "nbformat_minor": 4 -} +} \ No newline at end of file