SIGN IN SIGN UP

FlashInfer: Kernel Library for LLM Serving

0 0 188 Python

feat(trace): add check callbacks to trace templates (#3330)

## Summary
- add optional TraceTemplate.check and emit its source into fi_trace
JSON
- add public default_check/default_tolerances helpers for
reference-vs-actual output validation
- add representative checks for norm, GEMM, and single-request attention
templates

Closes #3266.

## Tests
- python -m ruff check flashinfer/trace/__init__.py
flashinfer/trace/template.py flashinfer/trace/templates/__init__.py
flashinfer/trace/templates/attention.py
flashinfer/trace/templates/gemm.py flashinfer/trace/templates/norm.py
tests/trace/test_fi_trace.py
- python -m ruff format --check flashinfer/trace/__init__.py
flashinfer/trace/template.py flashinfer/trace/templates/__init__.py
flashinfer/trace/templates/attention.py
flashinfer/trace/templates/gemm.py flashinfer/trace/templates/norm.py
tests/trace/test_fi_trace.py
- python -m pytest tests/trace/test_fi_trace.py -q
- python -m pytest tests/trace/test_fi_trace_template_consistency.py -q

<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

* **New Features**
* Added trace correctness validation with dtype-aware default tolerances
and a standard check API; many trace templates (attention, GEMM/BMM,
norm) now include validation checks and generated trace definitions may
embed check logic.

* **Documentation**
* Template docs updated to describe optional reference/check callables
and the expected check contract.

* **Tests**
* Expanded and consolidated tests to cover default tolerances,
per-template checks, and representative acceptance/rejection behaviors;
test utilities unified to use the new checker.

<!-- review_stack_entry_start -->

[![Review Change
Stack](https://storage.googleapis.com/coderabbit_public_assets/review-stack-in-coderabbit-ui.svg)](https://app.coderabbit.ai/change-stack/flashinfer-ai/flashinfer/pull/3330?utm_source=github_walkthrough&utm_medium=github&utm_campaign=change_stack)

<!-- review_stack_entry_end -->
<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Avery Huang <averyh@nvidia.com>
E
eigen committed
b6040ed2bf7b5bce387e552b52dc454673b803bc
Parent: c5a2b06
Committed by GitHub <noreply@github.com> on 5/29/2026, 6:01:53 PM