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 --> [](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