Tutorial: Grouped Query Attention with KV Caching Custom Operator

NOTE: The Jupyter Notebook below is included in the Chimera SDK and can be run interactively by running the following CLI command:

$ quadric sdk notebook

From the Jupyter Notebook window in your browser, select the notebook named /quadric/sdk-cli/examples/models/qwen/components/qwen_attention.ipynb.

Grouped Query Attention with KV Caching Custom Operator for Quadric CCL

1. Overview

This notebook demonstrates the implementation and validation of a custom operator for Grouped Query Attention (GQA) incorporating Key-Value (KV) Caching using Quadric CCL.

Note: groupedQueryAttention is a wrapper function around multiheadAttention. In this notebook, we are using an onnx that produces multiple tokens at a time. That is why there is a loop around multiheadAttention in groupedQueryAttention. In our final LLM lowering process, the graph will only produce one token at a time, and we will directly call multiheadAttention. We also currently focus on correctness/validation. We plan to explore performance optimizations for this op in the future.

2. Key Concepts Highlighted

2.1. Grouped Query Attention (GQA)

• What it is: Grouped Query Attention is an attention mechanism that serves as a middle ground between Multi-Head Attention (MHA) and Multi-Query Attention (MQA). In MHA, each query head has its own key and value head. In MQA, all query heads share a single key and value head. GQA groups multiple query heads to share a single key and value head.
• Why it's used: GQA significantly reduces the computational cost and memory bandwidth requirements compared to MHA, especially concerning the size of the KV Cache, while often retaining much of the model quality lost when using MQA. This makes inference faster and more memory-efficient, particularly on resource-constrained hardware.

2.2. KV Caching

• What it is: In autoregressive language models (like Qwen), generating the next token requires attending to all previous tokens. KV Caching is an optimization technique where the Key (K) and Value (V) states computed for previous tokens are stored (cached) and reused for subsequent token generation steps. Without it, these states would need to be recomputed every time, making inference prohibitively slow.
• How it's handled here: The custom_op_replacer function explicitly initializes zero-filled tensors named k_cache and v_cache as parameters for the custom operator. The custom groupedQueryAttention operator is designed to manage and utilize this KV cache internally during inference.

3. Work In Progress

3.1. Custom Operator Replacement

The notebook utilizes TVM's EPU utilities (graphutils.CustomOpReplacer) to swap out a sequence of standard ONNX operations (MatMuls, Adds, Transposes, Softmax, etc., corresponding to the attention calculation) with a single, opaque custom operator node within the ONNX graph. This allows leveraging hardware-specific acceleration available on the Quadric CCL platform via a custom C++ implementation (gqa.hpp).

Note that currently, we are using replacer.replace_subgraph_by_edge which is more basic and hardcodes the edges to replace on. There is also a more advanced custom op replacer in the works that utilizes pattern matching. This will be complete in the D5 deliverable (5/9).

3.2. Creating the ONNX graph

The provided ONNX graph comes from our huggingface to Quadric ONNX exporter script. This script will do the following steps:

1. Import QWEN torch decoder blocks from huggingface transformers repo.
2. Export to onnx with static shapes and kv caching based on the config.
3. Statically quantize the model with PTQ.

This script is also under development and a preliminary version will be given in the D5 deliverable.

4. Next Steps & Project Context:

This GQA+KV Cache operator is the only custom op required to lower the QWEN 1.5B model. The rest of the onnx graph will be lowered to CCL automatically with CGC.

Moving forward, we will be focusing on building general infrastructure for supporting LLMs on Quadric hardware. This will make supporting other LLMs more seamless.

With this, we will have the following generalizable pattern to lower QWEN and other future LLMs:

1. The onnx exporter script allows us to use huggingface config.json files to produce statically shaped and statically quantized ONNX graphs that we can support with today's software stack. (5/9)
2. We perform custom op replacement via pattern matching for the attention blocks. (5/9)
3. CGC lowers the rest of the model to CCL automatically. (5/9 for single decoder, 5/31 for multiple decoders).
4. Chipy defines autoregressive behavior to run the graph in a loop. (5/31)

Basic custom op replacement and CGC compilation

import onnx
import tvm.contrib.epu.graphutils as gutils
import numpy as np

MODEL_PATH = "../gqa.onnx"
quantized_model = onnx.load(MODEL_PATH)

input_edges = [
    "hidden_states_quantized",
]
output_edges = ["output_hidden_states_quantized"]

replacer = gutils.CustomOpReplacer(quantized_model)

param_name_dict = {
    "q_matmul_a_scale": "hidden_states_scale",
    "q_matmul_b": "onnx::MatMul_158_quantized",
    "q_matmul_b_scale": "onnx::MatMul_158_scale",
    "q_matmul_y_scale": "/decoder_layer/self_attn/q_proj/MatMul_output_0_scale",
    "q_add_a": "decoder_layer.self_attn.q_proj.bias_quantized",
    "q_add_a_scale": "decoder_layer.self_attn.q_proj.bias_scale",
    "q_add_b_scale": "/decoder_layer/self_attn/q_proj/MatMul_output_0_scale",
    "q_add_c_scale": "/decoder_layer/self_attn/q_proj/Add_output_0_scale",
    "k_matmul_a_scale": "hidden_states_scale",
    "k_matmul_b": "onnx::MatMul_174_quantized",
    "k_matmul_b_scale": "onnx::MatMul_174_scale",
    "k_matmul_y_scale": "/decoder_layer/self_attn/k_proj/MatMul_output_0_scale",
    "k_add_a": "decoder_layer.self_attn.k_proj.bias_quantized",
    "k_add_a_scale": "decoder_layer.self_attn.k_proj.bias_scale",
    "k_add_b_scale": "/decoder_layer/self_attn/k_proj/MatMul_output_0_scale",
    "k_add_c_scale": "/decoder_layer/self_attn/k_proj/Add_output_0_scale",
    "v_matmul_a_scale": "hidden_states_scale",
    "v_matmul_b": "onnx::MatMul_175_quantized",
    "v_matmul_b_scale": "onnx::MatMul_175_scale",
    "v_matmul_y_scale": "/decoder_layer/self_attn/v_proj/MatMul_output_0_scale",
    "v_add_a": "decoder_layer.self_attn.v_proj.bias_quantized",
    "v_add_a_scale": "decoder_layer.self_attn.v_proj.bias_scale",
    "v_add_b_scale": "/decoder_layer/self_attn/v_proj/MatMul_output_0_scale",
    "v_add_c_scale": "/decoder_layer/self_attn/v_proj/Add_output_0_scale",
    "qk_matmul_a_scale": "/decoder_layer/self_attn/Transpose_output_0_scale",
    "qk_matmul_b_scale": "/decoder_layer/self_attn/Transpose_3_output_0_scale",
    "qk_matmul_y_scale": "/decoder_layer/self_attn/MatMul_output_0_scale",
    "qk_mul_a_scale": "/decoder_layer/self_attn/MatMul_output_0_scale",
    "qk_mul_b": "/decoder_layer/self_attn/Constant_13_output_0_quantized",
    "qk_mul_b_scale": "/decoder_layer/self_attn/Constant_13_output_0_scale",
    "qk_mul_c_scale": "/decoder_layer/self_attn/Mul_2_output_0_scale",
    "qk_softmax_x_scale": "/decoder_layer/self_attn/Mul_2_output_0_scale",
    "qk_softmax_y_scale": "/decoder_layer/self_attn/Softmax_output_0_scale",
    "qkv_matmul_a_scale": "/decoder_layer/self_attn/Softmax_output_0_scale",
    "qkv_matmul_b_scale": "/decoder_layer/self_attn/Reshape_4_output_0_scale",
    "qkv_matmul_y_scale": "/decoder_layer/self_attn/MatMul_1_output_0_scale",
    "proj_matmul_a_scale": "/decoder_layer/self_attn/Reshape_5_output_0_scale",
    "proj_matmul_b": "onnx::MatMul_206_quantized",
    "proj_matmul_b_scale": "onnx::MatMul_206_scale",
    "proj_matmul_y_scale": "output_hidden_states_scale",
}


def custom_op_replacer(dict, idx):
    reversed_param_dict = {value: key for key, value in param_name_dict.items()}

    reversed_param_dict = {}
    for key, value in param_name_dict.items():
        reversed_param_dict.setdefault(value, []).append(key)
    params = {}
    results_list = []
    for k, v in dict.items():
        if v[1].size > 1:  # tensors
            results_list.append((reversed_param_dict[v[0]][0], v[1]))
        for name in reversed_param_dict[v[0]]:
            params[name] = v[1]

    results_list.append(("k_cache", np.zeros((1, 1, 64, 16), dtype=np.int32)))
    results_list.append(("v_cache", np.zeros((1, 1, 16, 256), dtype=np.int8)))
    results_list.extend(
        [
            (
                "q_matmul_scale",
                (
                    params["q_matmul_a_scale"]
                    * params["q_matmul_b_scale"]
                    / params["q_matmul_y_scale"]
                ),
            ),
            ("q_add_a_scale", (params["q_add_a_scale"] / params["q_add_c_scale"])),
            ("q_add_b_scale", (params["q_add_b_scale"] / params["q_add_c_scale"])),
            (
                "k_matmul_scale",
                (
                    params["k_matmul_a_scale"]
                    * params["k_matmul_b_scale"]
                    / params["k_matmul_y_scale"]
                ),
            ),
            ("k_add_a_scale", (params["k_add_a_scale"] / params["k_add_c_scale"])),
            ("k_add_b_scale", (params["k_add_b_scale"] / params["k_add_c_scale"])),
            (
                "v_matmul_scale",
                (
                    params["v_matmul_a_scale"]
                    * params["v_matmul_b_scale"]
                    / params["v_matmul_y_scale"]
                ),
            ),
            ("v_add_a_scale", (params["v_add_a_scale"] / params["v_add_c_scale"])),
            ("v_add_b_scale", (params["v_add_b_scale"] / params["v_add_c_scale"])),
            (
                "qk_matmul_scale",
                (
                    params["qk_matmul_a_scale"]
                    * params["qk_matmul_b_scale"]
                    / params["qk_matmul_y_scale"]
                ),
            ),
            (
                "qk_mul_scale",
                (params["qk_mul_a_scale"] * params["qk_mul_b_scale"] / params["qk_mul_c_scale"]),
            ),
            ("qk_mul_b", params["qk_mul_b"]),
            ("qk_softmax_x_scale", params["qk_softmax_x_scale"]),
            (
                "qk_softmax_y_inv_scale",
                (1.0 / params["qk_softmax_y_scale"]),
            ),  # Renamed
            (
                "qkv_matmul_scale",
                (
                    params["qkv_matmul_a_scale"]
                    * params["qkv_matmul_b_scale"]
                    / params["qkv_matmul_y_scale"]
                ),
            ),
            (
                "proj_matmul_scale",
                (
                    params["proj_matmul_a_scale"]
                    * params["proj_matmul_b_scale"]
                    / params["proj_matmul_y_scale"]
                ),
            ),
        ]
    )
    result_dict = {index: (item[0], np.array(item[1])) for index, item in enumerate(results_list)}
    return result_dict


ccl_func_name = "groupedQueryAttention<1536,2,16>"
_, new_model = replacer.replace_subgraph_by_edges(
    output_edges,
    input_edges,
    ccl_func_name,
    element_wise=False,
    keep_constants=list(param_name_dict.values()),
    process_const_callback=custom_op_replacer,
)
quantized_onnx = "gqa_custom_op.onnx"
onnx.save(new_model, quantized_onnx)

Run on EPU and validate against ORT

Functional correctness is validated within the notebook by the ChimeraJob.run_inference_harness method with the compare_ort=True flag. This compares the output of the custom operator run on EPU against a reference execution using ONNX Runtime, ensuring numerical equivalence within acceptable tolerances.

from tvm.contrib.epu.chimera_job.chimera_job import ChimeraJob

cgc_job = ChimeraJob(model_p=quantized_onnx, custom_ops="gqa.hpp")
cgc_job.compile()
cgc_job.run_inference_harness(compare_ort=True)
print(cgc_job)

2026-07-18 12:14 - INFO - epu - chimera_job - START==================================onnx_ingest
2026-07-18 12:14 - INFO - epu - codegen - START===============================optimize_relay
2026-07-18 12:14 - INFO - epu - codegen - START====================quantize_to_cpu_runnable_fx
2026-07-18 12:14 - INFO - epu - codegen - START====================build_cpu_runnable_fx_relay
2026-07-18 12:14 - INFO - epu - codegen - START=======================quantize_to_chimera_fx
2026-07-18 12:14 - INFO - epu - codegen - START=================================relay_to_tir
2026-07-18 12:14 - INFO - epu - codegen - START===========================relay_to_epu_relay
2026-07-18 12:14 - INFO - epu - codegen - START==============================adapt_and_order
2026-07-18 12:14 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-07-18 12:14 - INFO - epu - codegen - START=============================plan_lrm_virtual
2026-07-18 12:14 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-07-18 12:14 - INFO - epu - codegen - START===============================lrm_alloc_loop
2026-07-18 12:14 - INFO - epu - codegen - START==============================amend_ctrl_flow
2026-07-18 12:14 - INFO - epu - codegen - START================================lrm_splitting
2026-07-18 12:15 - INFO - epu - codegen - START==============================ext_split_relay
2026-07-18 12:15 - INFO - epu - codegen - START====================================build_tir
2026-07-18 12:15 - INFO - epu - chimera_job - Compilation of gqa_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 successful
2026-07-18 12:15 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7b8e3cc3f820>
FILM 1/1: 100%|██████████████████████████████████████████████████████| 1/1 [00:00<00:00, 641.43it/s]
2026-07-18 12:16 - INFO - epu - iss_testing - No tranges found for input, use default float range: <tvm.contrib.epu.interval.Interval object at 0x7b8e320eba00>
2026-07-18 12:16 - INFO - epu - iss_testing - Started Executing Onnxruntime...
2026-07-18 12:16 - INFO - epu - iss_testing - Done 0:00:00.051364



╒═════════════════════╤══════════════════════════════════════════════════════════════╕
│ Module Name         │ gqa_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1 │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ ONNX File           │ gqa_custom_op.onnx                                           │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Custom Ops          │ /quadric/sdk-cli/examples/models/qwen/components/gqa.hpp     │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Product Target      │ QC-U                                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Number of Cores     │ 1                                                            │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ ISS Clock Frequency │ 1.700                                                        │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ L2M Size            │ 16MB                                                         │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ LRM Size            │ 4kB                                                          │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ External Read BW    │ 128GBps                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ External Write BW   │ 128GBps                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ MACS per PE         │ 16                                                           │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max L2M             │ 0.000MB                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max LRM             │ 0.000kB                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Max Temp Ext Bytes  │ 0.023MB                                                      │
├─────────────────────┼──────────────────────────────────────────────────────────────┤
│ Network GMACs       │ 0.089                                                        │
╘═════════════════════╧══════════════════════════════════════════════════════════════╛

╒════╤════════╤════════════════════════════════╤═══════════════╤══════════════╤═══════╕
│    │ Type   │ Name                           │ shape         │ type         │ mse   │
╞════╪════════╪════════════════════════════════╪═══════════════╪══════════════╪═══════╡
│  0 │ Input  │ hidden_states_quantized        │ [1, 16, 1536] │ tensor[int8] │ n/a   │
├────┼────────┼────────────────────────────────┼───────────────┼──────────────┼───────┤
│  1 │ Output │ output_hidden_states_quantized │ [1, 16, 1536] │ tensor[int8] │ 0.159 │
╘════╧════════╧════════════════════════════════╧═══════════════╧══════════════╧═══════╛

Post-ISS Report 1.7 GHz ***
Fully placed-and-routed gate simulation: 
╒══════════════════════════════════╤════════╕
│ Latency (ms)                     │ 3.62   │
├──────────────────────────────────┼────────┤
│ FPS                              │ 276.07 │
├──────────────────────────────────┼────────┤
│ Average Power @ 3nm SSGNP (mW)   │ 943.22 │
├──────────────────────────────────┼────────┤
│ FPS per Watt @ 3nm SSGNP (FPS/W) │ 292.69 │
├──────────────────────────────────┼────────┤
│ Ext Rd Bytes (MB)                │ 5.34   │
├──────────────────────────────────┼────────┤
│ Ext Wr Bytes (MB)                │ 0.03   │
├──────────────────────────────────┼────────┤
│ Avg Ext Rd BW (GBps)             │ 1.44   │
├──────────────────────────────────┼────────┤
│ Avg Ext Wr BW (GBps)             │ 0.01   │
├──────────────────────────────────┼────────┤
│ MAC Utilization                  │ 0.09%  │
╘══════════════════════════════════╧════════╛
*** Data generated using 7nm SSGNP gatesim and scaled to 3nm

[SDK-CLI] : TotalCycles: 6,157,781
[SDK-CLI] : Executions/second: 276.07

compute      : ▇ 210.295K
data_array   : ▏ 38.138K
mac          : ▏ 13.248K
data_external: ▏ 77.469K
data_ocm     : ▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇▇ 5.819M

for more information check run directory: /quadric/sdk-cli/examples/models/qwen/components/ccl_build/gqa_custom_op_QC_U_1d7_16MB_4kB_128GBps_128GBps_16_OFF_x1_x1/run/20260718_121517_6c68bb