Conversation
vgene
force-pushed
the
feat/kernel-agent
branch
from
822467b to
588500e
Compare
vgene
changed the title
Add new nkipy.tools.kernel_agent tool providing: - CLI for discovering NumPy op support status on NKI hardware - Multi-stage kernel execution (numpy, compile, hardware) - LLM-based kernel generation via AWS Bedrock - Op testing framework with dtype and shape validation
vgene
force-pushed
the
feat/kernel-agent
branch
from
588500e to
0aff2b0
Compare
…te functions Extract `_compile_kernel` and `_execute_neff` helper functions from `baremetal_run_traced_kernel`.
- Implement sweep functionality for continuous generation-and-test loops - Include 50+ kernel prompt templates for common operations - Add CLI commands: generate, sweep, list-prompts - Configure package data to include prompt text files - Update documentation with new commands and architecture overview
| 3. No loops or conditionals based on array values | ||
| 4. The function must be self-contained — only use numpy (imported as np) | ||
|
|
||
| Output ONLY a JSON object such as: |
There was a problem hiding this comment.
Having this output and then some manual parser script may still be unstable: there might be edge cases like improperly escaped characters within the embedded source code, incorrect indentations, etc. We can consider using anthropic SDK which has bedrock and structured output support.
There was a problem hiding this comment.
change the output to structured. But the numpy code is still hard to enforce. The code will be executed as numpy first in the test so if it's invalid, it will show in that stage.
| @@ -0,0 +1,62 @@ | |||
| # ML activations / norms | |||
| softmax along the last axis | |||
There was a problem hiding this comment.
Normally, I would also include things like input/output shapes/datatypes when asking the agent to write kernels, or the LLM can produce anything that are mostly reasonable but sometimes unexpected.
There was a problem hiding this comment.
I want to let the prompt as underspecified as possible so the LLM can generate "random" numpy code to test the lowering. The goal is to find corner cases that fail in NKIPy.
| input_specs = data.get("inputs", {}) | ||
|
|
||
| inputs = {} | ||
| for inp_name, spec in input_specs.items(): |
There was a problem hiding this comment.
This logic feels fragile. Maybe it's ok as a starting point.
There was a problem hiding this comment.
changed to tool use to make sure no fallthrough shape and dtype
| from nkipy.tools.kernel_agent.executor import ExecutionResult, run_kernel | ||
|
|
||
| # Target operations to test | ||
| TARGET_OPS = { |
vgene
force-pushed
the
feat/kernel-agent
branch
2 times, most recently
from
716a1fe to
81816ae
Compare
vgene
force-pushed
the
feat/kernel-agent
branch
from
81816ae to
fbb62c9
Compare
3 participants
Issue #, if available:
Description of changes:
Add new nkipy.tools.kernel_agent tool providing:
By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.