Add Knowledge Database to Kernel optimization

[kaiming-cheng]

This PR Introduces a hierarchical optimization database that stores GPU kernel optimization techniques and code examples for the RAG-based optimization.

Key components:

• OptNode / OptHierarchy: Tree structure organizing optimizations by bottleneck type (latency, memory, utilization) → technique → code example
• docs/: Optimization technique documentation (TMA, PID swizzling, persistence)
• code_samples/: Reference Triton kernel implementations (matmul, matadd with various optimizations applied)

Optimization techniques covered:

• Host-side and device-side Tensor Memory Accelerator (TMA)
• PID swizzling for L2 cache locality
• Persistent kernel programming style

This database enables the agent to retrieve relevant optimization strategies and reference implementations based on diagnosed performance bottlenecks.

Test

query = "use TMA for memory optimization"
prescriber = RAGPrescriber()

opt_node, similarities = prescriber.retrieve(query)
Retrieved: 
============================= On-Device Tensor Memory Accel... (similarity: 0.573)

Generated context (4620 chars):
--------------------------------------------------------------------------------
## Optimization Technique
...

context = prescriber.build_context(opt_node, max_code_examples=1, max_chars=2000)
## Code Examples
...
 
add_kernel[grid](
        x,
        y,
        output,
        M,
        N,
        x.stride(0),
        x.stride(1),
        BLOCK_SIZE_M=BLOCK_SIZE_M,
        BLOCK_SIZE_N=BLOCK_SIZE_N,
    )
    return output

(Showing 1 of 2 examples)

[Jack-Khuu]

Looks good, is it hard to add the integration code using the RAG into this PR too?

Remember to cite for the code_samples/docs

• Drop [Optimization 7/n] from the title just to avoid confusion

[Jack-Khuu]

Do we need this for any reason?

[kaiming-cheng]

good call - I've removed it in following commit

[Jack-Khuu]

nit: For legibility can we add a helper like add_relation or something that updates the child+parent symmetrically

It's easy to parse here, but level3 is a harder to parse

[Laurawly]

database_path seems wrong, use Path(__file__).resolve().parents[...] until you hit the project root (where pyproject.toml is)

[Laurawly]

retrieve() calls embeddings.embed_query(node.opt_desc) for every node each time. Some nodes include full code examples which is slow and costly.
Precompute embeddings once at init and cache them per node.
Or at least cache in-memory dict {OptNode: embedding} after first compute.
Also consider embedding only L1/L2 text nodes for retrieval, then traverse down for code examples. Embedding code blobs is noisy and expensive.

[Laurawly]

It traverses from the selected node down and concatenates every descendant’s opt_desc, including entire code files. That will quickly blow context limits and drown signal.
Put a max character/token budget and stop after N leaf examples.
Add separators between nodes (right now it just concatenates).
Optionally include only: (a) technique description + (b) top-k leaf code examples.

[Laurawly]

I don’t see a kernel_perf_agent/kernel_opt/database/docs/__init__.py added in this PR.

[kaiming-cheng]

Thanks for the catch, updated this in the c57e13c

[Laurawly]

Adding langchain-openai is a big dependency. If you only need embeddings, consider using the project’s existing LLM client (if any) or a thinner dependency.

If you keep it, I’d suggest pinning compatible versions or adding it as an optional dependency for the RAG feature.

[kaiming-cheng]

Good point! We can use OpenAI's text-embedding model for simplicity