forked from ccf-ai-infra/TileOPs-Metax
2 Commits
| Author | SHA1 | Message | Date |
|---|---|---|---|
Cao Ying
|
fd77066565
|
[Perf][Elementwise] Generalize kernel caching and autotune to all kernel families (#560)
## Summary Extract BinaryKernel's `init_config()` caching pattern (`_compiled_fn` pre-compilation) to **UnaryKernel**, **FusedGatedKernel**, and all **11 custom kernels** (LeakyRelu, Elu, Hardtanh, Softplus, Prelu, Where, Clamp, MaskedFill, NanToNum, Alibi, Sinusoidal). Add `autotune_configs` property with serialization-fallback `autotune()` override to UnaryKernel and FusedGatedKernel. Closes #551 ## Changes - **UnaryKernel**: Added `autotune_configs` (9 configs for fp16/bf16/fp32, 6 for fp8), `autotune()` with serialization fallback, `init_config()` override that caches `_compiled_fn` - **FusedGatedKernel**: Same pattern as UnaryKernel — `autotune_configs`, `autotune()` fallback, `init_config()` caching - **11 custom kernels**: `init_config()` override to cache `_compiled_fn` (custom kernels use factory functions, so autotune is not applicable) - **docs/perf/elementwise.md**: Updated checklist to reflect all kernel families now have caching and autotune ## Test plan - [x] AC-1: Modified files pass unit tests — 432 tests pass (31 new + 401 existing), 0 failures - [x] AC-2: All UnaryKernel subclasses use cached compiled functions — `_compiled_fn` set in `init_config()`, verified by `TestUnaryCaching` (5 tests) - [x] AC-3: FusedGatedKernel subclasses use cached compiled functions — `_compiled_fn` set in `init_config()`, verified by `TestFusedGatedCaching` (4 tests) - [x] AC-4: `autotune_configs` defined for UnaryKernel and FusedGatedKernel with >= 3 configs — verified by `TestAutotuneConfigs` (8 tests) - [x] AC-5: Existing benchmark numbers show no regression — all 401 existing tests pass, BinaryKernel autotune still works - [x] AC-6: `docs/perf/elementwise.md` updated with new caching/autotune patterns ## Benchmark **Environment**: NVIDIA H200 · PyTorch 2.9.1+cu128 · CUDA 12.8 · Driver 575.57.08 173 benchmarks passed. No runtime regression. Caching eliminates per-forward JIT lookup overhead. ### Unary ops (TileOPs vs PyTorch baseline) | Op | N | dtype | TileOPs (ms) | Baseline (ms) | BW (TB/s) TileOPs | BW (TB/s) Baseline | | --- | --- | --- | --- | --- | --- | --- | | exp | 4000000 | fp16 | 0.01 | 0.01 | 2.40 | 2.42 | | exp | 4000000 | bf16 | 0.01 | 0.01 | 2.41 | 2.42 | | gelu | 4000000 | fp16 | 0.01 | 0.01 | 1.64 | 1.77 | | gelu | 4000000 | bf16 | 0.01 | 0.01 | 1.58 | 1.71 | | logical_not | 4000000 | fp16 | 0.01 | 0.01 | 0.82 | 2.07 | | bitwise_not | 4000000 | int32 | 0.02 | 0.01 | 1.94 | 3.06 | | isnan | 4000000 | fp16 | 0.01 | 0.01 | 0.82 | 2.09 | ### Binary ops (TileOPs vs PyTorch baseline, fp16) | Op | Shape | TileOPs (ms) | Baseline (ms) | BW (TB/s) TileOPs | BW (TB/s) Baseline | | --- | --- | --- | --- | --- | --- | | sub | 1024×20480 | 0.03 | 0.03 | 3.80 | 3.82 | | mul | 1024×20480 | 0.03 | 0.03 | 3.80 | 3.82 | | div | 1024×20480 | 0.03 | 0.03 | 3.73 | 3.68 | | remainder | 1024×10240 | 0.02 | 0.02 | 3.36 | 2.66 | | floor_divide | 1024×10240 | 0.02 | 0.06 | 3.36 | 1.03 | | maximum | 1024×20480 | 0.03 | 0.03 | 3.78 | 3.78 | | minimum | 1024×20480 | 0.03 | 0.03 | 3.79 | 3.78 | ### Fused gated ops (TileOPs vs PyTorch baseline, fp16) | Op | Shape | TileOPs (ms) | Baseline (ms) | Speedup | | --- | --- | --- | --- | --- | | gelu_and_mul | 1024×4096 | 0.01 | 0.04 | 4.0× | | gelu_and_mul | 1024×10240 | 0.02 | 0.08 | 4.0× | | gelu_and_mul | 1024×20480 | 0.04 | 0.15 | 3.8× | | gelu_tanh_and_mul | 1024×4096 | 0.01 | 0.03 | 3.0× | | gelu_tanh_and_mul | 1024×10240 | 0.02 | 0.07 | 3.5× | | gelu_tanh_and_mul | 1024×20480 | 0.04 | 0.14 | 3.5× | ### Broadcast ops (TileOPs vs PyTorch baseline, fp16) | Op | TileOPs (ms) | Baseline (ms) | Speedup | | --- | --- | --- | --- | | sub (bcast, largest) | 0.02 | 0.06 | 3.0× | | mul (bcast, largest) | 0.02 | 0.06 | 3.0× | | div (bcast, largest) | 0.03 | 0.07 | 2.3× | ### Strategy comparison (largest shape 1024×20480, fp16) | Op | Strategy | Latency (ms) | BW (TB/s) | | --- | --- | --- | --- | | relu (unary) | direct | 0.07 | 1.20 | | relu (unary) | explicit_parallel | 0.08 | 1.03 | | relu (unary) | register_copy | 0.02 | 3.71 | | add (binary) | direct | 0.08 | 1.65 | | add (binary) | explicit_parallel | 0.03 | 3.81 | | silu_and_mul (fused) | direct | 0.04 | 1.46 | | silu_and_mul (fused) | explicit_parallel | 0.02 | 2.84 | <details> <summary>Full benchmark report (profile_run.log)</summary> ``` 173 passed, 0 failed (516.61s) See profile_run.log for per-op latency/TFLOPS/bandwidth tables. ``` </details> 🤖 Generated with [Claude Code](https://claude.com/claude-code) --------- Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |
|
|
Cao Ying
|
d4852d2082
|
[Doc] Add elementwise kernel performance checklist and evidence (#556)
## Summary - Add `docs/perf/` directory with structured performance documentation for elementwise kernels - **Checklist** (`elementwise.md`, ~25 lines): heuristic rules for agent/reviewer audit covering strategy selection, vectorization, config tuning, caching, and code quality - **Evidence** (`elementwise-evidence.md`): measured data and reasoning from PRs #500, #537, #538, #539 — loaded on demand only when justification is needed - **Index** (`README.md`): test environment table (H200/CUDA 12.8/TileLang 0.1.8) scoping all conclusions, plus category index for future op families - Checklist framed as heuristics with explicit guidance: violations need justification, not mandatory fixes ## Test plan - [x] All markdown files pass mdformat and codespell pre-commit hooks - [x] Cross-links between checklist, evidence, and README verified - [x] No runtime code changes 🤖 Generated with [Claude Code](https://claude.com/claude-code) Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> |