[Bench][Linear-Attn] Add Qwen GDN prefill benchmark surface#5
Draft
superAngGao wants to merge 35 commits into
Draft
[Bench][Linear-Attn] Add Qwen GDN prefill benchmark surface#5superAngGao wants to merge 35 commits into
superAngGao wants to merge 35 commits into
Conversation
superAngGao
force-pushed
the
impl/gdn-prefill-op
branch
3 times, most recently
from
9d21d88 to
9282b2b
Compare
…16, causal) (tile-ai#1614) ## Summary Adds `GQAPrefillFwdWsPersistentCausalKernel` — a faithful port of FlashInfer's `single_prefill_sm90` — and wires it into `GroupedQueryAttentionPrefillFwdOp`. On Hopper, the prefill op now dispatches to this kernel for the **fp16 / causal / dim==128 / default sm_scale / no-softcap** path, reaching FlashInfer-level performance; all other cases fall back to the existing `GQAPrefillFwdKernel`. ## Performance Kernel-level parity with FlashInfer. Per-kernel CUDA time from `benchmarks/ops/attention/bench_gqa.py` (H200, fp16, causal, L2 flushed per iteration; `single_prefill_sm90` is the kernel this ports): | batch · Sq · Skv | tileops kernel | flashinfer kernel | tileops/flashinfer | | --- | --- | --- | --- | | 4 · 512 · 512 | 42.0 µs | 40.4 µs | 1.04x | | 2 · 512 · 4096 | 122.6 µs | 126.2 µs | 0.97x | | 1 · 512 · 4096 | 122.0 µs | 125.3 µs | 0.97x | Within ~3–4% across the workloads: at parity / marginally faster on long context (Skv=4096, ~525 TFLOPS), marginally slower on the short square case (Skv=512). The bench now reports per-kernel CUPTI time directly — it wraps the timed call in a `record_function` scope and sums only the kernels Kineto attributes to it, so the per-iteration L2 flush is excluded from the measurement. ## Kernel 2-warpgroup ping-pong (1 TMA producer WG + 2 consumer WG, 384 threads); register-P (rs-wgmma) held in a fp16 fragment; `set_max_nreg` producer→24 / consumer→240; named-barrier WarpScheduler ping-pong; delayed PV + `wait_wgmma(1)` so softmax(k) overlaps in-flight PV(k-1); pipelined rescale; smem-staged coalesced output. Causal uses bottom-right alignment (`co = seq_len_kv - seq_len_q`). ## Notes - **No log-sum-exp output**, matching FlashInfer's forward kernel. The op's `(output, lse)` contract is satisfied with `lse=None`. Emitting lse from the epilogue perturbed the register-tuned schedule and cost ~35% on long context; a backward pass that needs lse must compute it separately. - `bench_gqa`: the flashinfer baseline helper is generalized to handle `Sq != Skv` and is now recorded for the prefill bench. - Manifest `kernel_map` / `source.kernel` are intentionally **not** touched here — adding a new-file kernel requires a `source.kernel` change, which is outside the implementation-PR carve-out and will follow as a separate manifest PR. ## Verification - 22/22 prefill-fwd correctness tests pass (fp16 → FA3 kernel; bf16 / non-causal / softcap / non-default scale → fallback). - 6/6 prefill bench cases pass. - `ruff` clean.
) ## Summary Fixes tile-ai#1617 - GLA decode benchmark timeout This PR resolves the kernel compilation issue that prevented GLA decode from working. The root cause was intermediate variable assignments inside `T.Parallel(k_tile, dim_v)` loops within `T.Pipelined()`, which created `tirx.Bind` IR nodes that TileLang's warp-specialization pass couldn't handle. ## Changes ### Kernel Compilation Fix (`tileops/kernels/gla_recurrence.py`) **Root cause**: Intermediate variables inside `T.Parallel(k_tile, dim_v)` loops within `T.Pipelined()` created `tirx.Bind` IR nodes that TileLang's warp-specialization pass couldn't handle. **Solution**: - Preload `k`, `v`, `gk` tensors into shared memory before the pipelined loop - Inline all computations in the pipelined loop (eliminate intermediate variables) - Applied to both `_gla_decode_tl()` (fp16/bf16) and `_gla_decode_fp32_tl()` **Before** (BROKEN): ```python for kt in T.Pipelined(dim_k // k_tile, num_stages=num_stages): T.copy(state[bid, hid, kt * k_tile, 0], h_tile) for kk, j in T.Parallel(k_tile, dim_v): gk_val = gk[bid, hid, kt * k_tile + kk] # ← Intermediate variable alpha_i = T.exp2(gk_val * _LOG2E) # ← Creates BindNode new_state[...] = alpha_i * h_tile[kk, j] + ... ``` **After** (FIXED): ```python # Preload to shared memory k_shared = T.alloc_shared([dim_k], dtype) gk_shared = T.alloc_shared([dim_k], dtype) v_shared = T.alloc_shared([dim_v], dtype) for i in T.Parallel(dim_k): k_shared[i] = k[bid, hid, i] gk_shared[i] = gk[bid, hid, i] for i in T.Parallel(dim_v): v_shared[i] = v[bid, hid, i] # Inline computations - no intermediate variables for kt in T.Pipelined(dim_k // k_tile, num_stages=num_stages): T.copy(state[bid, hid, kt * k_tile, 0], h_tile) for kk, j in T.Parallel(k_tile, dim_v): new_state[...] = ( T.exp2(gk_shared[kt * k_tile + kk] * _LOG2E) * h_tile[kk, j] + k_shared[kt * k_tile + kk] * v_shared[j] ) ``` ## Testing All GLA decode tests pass: ```bash tests/ops/test_gla_recurrence.py::test_gla_decode[1-4-64-64-dtype0-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[1-4-64-64-dtype1-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[1-4-64-64-dtype2-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[2-8-64-64-dtype3-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[2-4-128-128-dtype4-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[2-8-64-64-dtype5-False] PASSED tests/ops/test_gla_recurrence.py::test_gla_decode[2-8-64-64-dtype6-False] PASSED ``` ✅ Kernels compile successfully for all dtypes (fp32, fp16, bf16) ✅ No NaN/Inf in outputs ✅ Benchmark unblocked ## Impact - **Performance**: May improve due to reduced global memory traffic; shared memory usage increases by ~768 bytes per block (dim_k=64, dim_v=64, fp32) - **Pattern alignment**: Now matches DeltaNet/GatedDeltaNet coding patterns - **Correctness**: All tests pass with proper numerical results ## Environment Note The test reference implementation uses `torch.einsum`. This requires PyTorch built with matching CUDA version (e.g., PyTorch 2.10.0+cu129 for CUDA 12.9). If you encounter CUBLAS errors with batch operations, ensure your PyTorch CUDA variant matches your system CUDA version. ## Benchmark Results 🚀 Benchmarked on **NVIDIA H200** (CUDA 12.9, PyTorch 2.10.0+cu129) ### Performance Comparison: TileOPs vs FLA (Flash Linear Attention) | Config | dtype | TileOPs (ms) | FLA (ms) | Speedup | |--------|-------|--------------|----------|---------| | B=1, H=32, D=64 | fp32 | 0.0080 | 0.0071 | **0.89x**⚠️ | | B=1, H=32, D=128 | fp32 | 0.0143 | 0.0080 | **0.56x**⚠️ | | B=1, H=32, D=128 | fp16 | 0.0143 | 0.0089 | **0.62x**⚠️ | | B=1, H=32, D=128 | bf16 | 0.0139 | 0.0090 | **0.65x**⚠️ | | B=8, H=32, D=128 | fp32 | 0.0302 | 0.0215 | **0.71x**⚠️ | | B=8, H=32, D=128 | fp16 | 0.0313 | 0.0187 | **0.60x**⚠️ | | B=8, H=32, D=128 | bf16 | 0.0298 | 0.0188 | **0.63x**⚠️ | | B=16, H=32, D=128 | fp32 | 0.0831 | 0.0372 | **0.45x**⚠️ | | B=16, H=32, D=128 | fp16 | 0.0785 | 0.0284 | **0.36x**⚠️ | | B=16, H=32, D=128 | bf16 | 0.0761 | 0.0285 | **0.37x**⚠️ | | B=32, H=32, D=128 | fp32 | 0.1889 | 0.0670 | **0.35x**⚠️ | | B=32, H=32, D=128 | fp16 | 0.1828 | 0.0470 | **0.26x**⚠️ | | B=32, H=32, D=128 | bf16 | 0.1752 | 0.0472 | **0.27x**⚠️ | | B=64, H=32, D=128 | fp32 | 0.4045 | 0.1318 | **0.33x**⚠️ | | B=64, H=32, D=128 | fp16 | 0.3890 | 0.0841 | **0.22x**⚠️ | | B=64, H=32, D=128 | bf16 | 0.3772 | 0.0846 | **0.22x**⚠️ | ### Analysis **Current Status:** -⚠️ TileOPs is currently **slower** than FLA by 1.1x-4.6x across all configurations - FLA achieves **0.9-1.6 TFLOPS** vs TileOPs' **0.07-0.6 TFLOPS** - The gap widens with larger batch sizes (B=64: 4.5x slower) **Performance Gap:** - **Small batches (B=1)**: 11-78% slower - **Medium batches (B=8-16)**: 2.2-2.8x slower - **Large batches (B=32-64)**: 2.8-4.6x slower **Key Observations:** 1. ✅ **Correctness verified**: All tests pass against FLA reference 2. ✅ **Kernel compiles successfully**: Fixed BindNode compilation issue 3.⚠️ **Performance needs optimization**: Current implementation is functional but not yet competitive with FLA's highly optimized kernels ### Next Steps for Performance This PR successfully **unblocks GLA decode** by fixing the compilation issue. Performance optimization can be addressed in follow-up work: 1. **Tuning opportunities**: Adjust `num_stages`, `k_tile`, thread counts 2. **Memory access patterns**: Further optimize shared memory usage 3. **Warp-level optimizations**: Leverage H200's compute capabilities 4. **Benchmark-driven iteration**: Profile with NSight Compute to identify bottlenecks
## Summary Fixes tile-ai#1613. TileLang 0.1.11 validates/binds autotuned JIT signatures before candidate configs are applied. Several TileOPs kernels expose tunable JIT parameters such as `block_m`, `block_n`, and `bdim` as required arguments, so calling the autotuned wrapper with no kwargs can fail before autotuning starts. This seeds the autotuned wrapper call with bindable config values while preserving the existing autotune search space. ## Changes - Add shared `Kernel._call_autotuned_kernel()` / `_autotune_initial_kwargs()` helpers. - Seed base `Kernel.autotune()` with `default_config`, filtered to the JIT signature. - Support common JIT parameter aliases such as `threads_arg` and `npt_arg`. - Update custom autotune overrides that previously called autotuned wrappers with no kwargs: - `SoftmaxKernel` - `LogSumExpKernel` - `SparseMlaKernel` - `FP8LightingIndexerKernel` - DeltaNet / GatedDeltaNet sub-kernel autotune paths - Add unit coverage for required tunable binding, signature filtering, and alias mapping. ## Validation - `python3 -m pytest -q tests/kernels/test_kernel_autotune_binding.py` - `python3 -m ruff check tileops/kernels/kernel_base.py tileops/kernels/reduction/softmax.py tileops/kernels/reduction/logsumexp.py tileops/kernels/attention/deepseek_dsa_decode.py tileops/kernels/fp8_lighting_indexer.py tileops/kernels/deltanet/deltanet_fwd.py tileops/kernels/deltanet/deltanet_bwd.py tileops/kernels/gated_deltanet/gated_deltanet_fwd.py tests/kernels/test_kernel_autotune_binding.py` - `python3 -m compileall -q ...` on changed Python files - commit hooks / pre-commit checks Note: local environment has TileLang 0.1.9, so the regression test mocks the TileLang 0.1.11-style binding requirement instead of running the exact 0.1.11 autotuner.
…al handling (tile-ai#1626) Fixes tile-ai#1625. ## Summary Resolves the nightly Triton/Inductor `PermissionError` under `/ci-cache` by establishing a UID-agnostic ownership invariant for the shared CI build cache, plus a related runner stop-signal fix and the gpu-smoke cache de-namespacing. ### Cache ownership - `.github/runner/Dockerfile`: pin `ci-runner` to a deterministic UID/GID 1000 (`groupadd -g 1000 ci-runner && useradd -u 1000 -g 1000 ...`) so the account deterministically owns the shared `/ci-cache` bind-mount across image rebuilds. - The `chown -R ci-runner:ci-runner ... /ci-cache` line is retained but documented as serving only the **unmounted** image-verification path; in the mounted CI path the host bind-mount shadows those dirs (host ownership wins). - `.github/runner/entrypoint.sh`: set `umask 002` before the runner agent starts so compiled kernels land group-writable (`0775`/`2775`) for cross-UID cache reuse. ### Runner stop-signal handling - `.github/runner/entrypoint.sh`: `exec ./run.sh` + `export RUNNER_MANUALLY_TRAP_SIG=1`. - **Why:** `docker stop` (and the launcher's `systemctl stop`) delivers SIGTERM only to PID 1. With `./run.sh` as a foreground child, PID 1 was the entrypoint bash, which defers its trap while waiting and never relays the signal; an idle ephemeral runner then ignored SIGTERM and was SIGKILLed only after docker's full stop grace (~30 min via the launcher's `docker stop -t 1800`). - **Fix:** `exec` makes `run.sh` PID 1 so docker's SIGTERM reaches it; `RUNNER_MANUALLY_TRAP_SIG=1` makes `run.sh` trap SIGTERM/SIGINT and forward SIGINT to the listener process group for graceful shutdown. - **Trade-off:** a stop signal now cancels an in-progress job rather than draining it. ### gpu-smoke cache de-namespacing - `.github/workflows/gpu-smoke.yml`: drop the per-run `TRITON_CACHE_DIR=/ci-cache/triton/gpu-smoke-<run>-<attempt>` step so gpu-smoke shares the warm Triton cache like nightly. `TRITON_CACHE_DIR` falls back to the image ENV default `/ci-cache/triton`. This is safe now that ownership is UID-agnostic (above) and fork write isolation lives at the runner's overlay mount, not the workflow. It also stops the per-run dirs accumulating. - `nightly.yml` and `scripts/ci/prepare_cache_dirs.sh` need no change: the earlier run-local-cache hotfix (PR tile-ai#1618) was never merged, so nightly already uses the shared cache and `prepare_cache_dirs.sh` does not exist on `main`. ## Test plan - [x] Dockerfile pins `ci-runner` to UID/GID 1000; `entrypoint.sh` sets `umask 002` - [x] Image rebuilt (same tilelang commit / tag) and deployed on the runner host; verified `id ci-runner` = `1000:1000`, `umask` = `0002` - [x] Host cache normalized (group `1000` + group-writable + setgid); a container `ci-runner` can write into a previously-foreign cache dir (was `Permission denied`) - [x] Idle runner `docker stop` returns in ~0.3s (was ~30 min) after the signal fix - [x] gpu-smoke.yml parses; `TRITON_CACHE_DIR` resolves to the shared `/ci-cache/triton` via image ENV - [ ] Nightly op tests (GLA / Mamba / MoE) pass the Triton/Inductor cache step with no `PermissionError` (final end-to-end confirmation) - [x] pre-commit passed on changed files --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
…gsumexp kernels under TileLang 0.1.11 (tile-ai#1627) ## Summary Fixes autotune failures under TileLang 0.1.11 for three kernel families that were not covered by PR tile-ai#1619. ### Root cause TileLang 0.1.11 changed its autotuner: when `autotuned_kernel_fn(**initial_kwargs)` is called with the tunable JIT parameters already bound, it records those params in the cache key and — on a cache hit — treats them as "already tuned," skipping the benchmarking sweep entirely and returning `config=None`. PR tile-ai#1619 added `_call_autotuned_kernel()` to seed required JIT params (fixing the `TypeError: missing N required positional arguments`), but this inadvertently triggered the new 0.1.11 skip behavior, causing: - `SSDStatePassingFwdKernel` / `SSDChunkScanFwdKernel` / `SSDDecodeKernel`: `Best config: None` - `SoftmaxKernel` / `LogSumExpKernel`: `TypeError: '<' not supported between instances of 'NoneType' and 'float'` ### Fix Pass `do_not_specialize=tunable_params` to TileLang's `autotune()` decorator. This tells 0.1.11 to exclude those parameter names from the cache key, so the seeded initial values no longer trigger the skip path — the full benchmarking sweep runs as expected. ## Changes - **`tileops/kernels/kernel_base.py`** (`Kernel.autotune()`): extract tunable param names from the JIT signature and pass as `do_not_specialize`. Fixes `SSDStatePassingFwdKernel`, `SSDChunkScanFwdKernel`, `SSDDecodeKernel`, and any future kernel that uses the base `autotune()`. - **`tileops/kernels/reduction/softmax.py`** (`SoftmaxKernel.autotune()`): same pattern for the custom per-tile_n autotune loop. - **`tileops/kernels/reduction/logsumexp.py`** (`LogSumExpKernel.autotune()`): same pattern for the custom per-tile_n autotune loop. ## Validation Tested against real TileLang 0.1.11 (`0.1.11+cuda.gitcd37ed5f`) in `flashmlaenv`: **Mamba kernels** — all run full sweeps and return valid configs: - `SSDStatePassingFwdKernel` → `{'block_d': 64, 'threads': 128, 'vectorize': False}` - `SSDChunkScanFwdKernel` → `{'block_l': 64, 'block_p': 64, 'block_n': 64, 'block_s': 64, 'threads': 128}` - `SSDDecodeKernel` → `{'block_p': 4, 'block_n': 32, 'threads': 128}` **Softmax/logsumexp** — `tests/ops/test_softmax.py`: **135 passed** (previously 3 failed with `TypeError`) **Autotune binding unit tests** — `tests/kernels/test_kernel_autotune_binding.py`: **3 passed** **deepseek_dsa_decode** — `tests/ops/attention/test_deepseek_dsa_decode.py`: **1 passed** **Ruff** — all PR tile-ai#1619 validation files pass cleanly.
…try (tile-ai#1629) ## Summary Makes dense `GemmOp` **input-inferred** and brings it under the manifest, aligning with DeepGEMM's call-time-JIT model. Implements step 1 of tile-ai#1628. - `GemmOp` no longer commits shape/dtype at construction: `forward(a, b)` infers `m/n/k` and `dtype` from the inputs, builds & caches the dtype-specialized kernel per `(m, n, k, dtype)` on first use (mirroring DeepGEMM's compile-on-first-call + per-config cache), and returns a freshly-allocated `d`. - Default layout changed **NN → NT** (`trans_b=True`) to match DeepGEMM's primary path and `GroupedGemmOp`. `(trans_a, trans_b)` ≡ DeepGEMM `nt/nn/tn/tt`. - New `tileops/manifest/gemm.yaml` with the `GemmOp` entry (`status: implemented`). - GEMV fast-path dispatch (`m==1`/`n==1` on SM90) moved into `forward`. ## Verification (H200, container stack) - `scripts/validate_manifest.py` → exit 0, no GemmOp errors. - `pytest tests/ops/test_gemm.py` → **22 passed** (incl. NN/NT/TN/TT layouts, GEMV boundary, fp16/bf16, and the `tune=True` case). - `pytest benchmarks/ops/bench_gemm.py` → 4 passed (manifest-driven, `eval_roofline`-based, cuBLAS baseline). - `ruff check` clean on all changed Python files. ## Deviations from the issue plan (please review) 1. **Roofline is func mode, not inline.** A single static inline `vars` expression can't extract `M/N/K` correctly across all `trans_a/trans_b` layouts, and `M/N/K` can't be declared as manifest params (they're neither `__init__` nor `forward` args, so the signature check would reject them). Func mode (`tileops.ops.gemm._gemm_roofline`, reads `op.m/n/k/dtype` bound during `forward`) is correct for every layout. `eval_roofline()` is therefore valid only after the first `forward()`. 2. **Includes a small kernel autotune fix (`tileops/kernels/gemm.py`).** `tune=True` was broken for *both* `GemmKernel` and `GemvKernel` on the current tilelang stack: the `@tilelang.jit` builders had no parameter defaults, so the autotuner's validation `get_tir()` call (empty args) raised `missing required argument 'block_m'/'block_n'`. Fix: added default values to the tunable builder params. This is scoped to the GEMM/GEMV builders only — it does **not** touch the shared `kernel_base.autotune()` path, so other ops' autotune (also affected by the same upstream tilelang drift) is intentionally left for a separate tracking issue. (Authorized expansion beyond the issue's "no kernel" constraint.) 3. **Trust-model split.** The manifest entry is isolated in its own commit (`[Manifest][GEMM] …`) so it can be cherry-picked into a standalone human-reviewed manifest-only PR before merge, per `docs/design/trust-model.md`. This draft bundles it with the implementation only for review convenience. ## Benchmark (H200, torch 2.10.0+cu129) `pytest benchmarks/ops/bench_gemm.py` on an **idle GPU** (`CUDA_VISIBLE_DEVICES` pinned to a 0%-util device; SM clock locked at 1500 MHz) — TileOPs `GemmOp` (default cp.async `GemmKernel`) vs `torch.matmul` (cuBLAS), via the repo's CUPTI / cold-L2 profiler. Speedup = cuBLAS latency ÷ TileOPs latency (>1 = TileOPs faster). | m | n | k | layout | dtype | TileOPs (ms / TFLOPS) | cuBLAS (ms / TFLOPS) | speedup | |---|---|---|---|---|---|---|---| | 1024 | 1024 | 1024 | NN | fp16 | 0.0165 / 130 | 0.0075 / 288 | 0.45× | | 1024 | 1024 | 1024 | NN | bf16 | 0.0165 / 130 | 0.0075 / 288 | 0.45× | | 128 | 2112 | 7168 | NT | bf16 | 0.1146 / 34 | 0.0184 / 211 | 0.16× | | 128 | 7168 | 2048 | NT | bf16 | 0.0289 / 130 | 0.0135 / 278 | 0.47× | | 4096 | 2112 | 7168 | NT | bf16 | 0.4437 / 280 | 0.3062 / 405 | 0.69× | | 4096 | 7168 | 2048 | NT | bf16 | 0.3241 / 371 | 0.3132 / 384 | 0.97× | | 4096 | 4096 | 7168 | NT | fp16 | 0.6130 / 392 | 0.6242 / 385 | 1.02× | | 4096 | 4096 | 7168 | NT | bf16 | 0.6078 / 396 | 0.6078 / 396 | 1.00× | | 4096 | 7168 | 16384 | NT | bf16 | 2.5455 / 378 | 2.5278 / 381 | 0.99× | | 4096 | 24576 | 1536 | NT | bf16 | 0.9270 / 334 | 0.8552 / 362 | 0.92× | **Takeaway:** at cuBLAS parity (0.92–1.02×) on large compute-bound prefill shapes (m=4096, K-aligned), but meaningfully slower on small / decode-token shapes (m ≤ 1024 → 0.16–0.47×) and on non-256-aligned N (4096×2112×7168 → 0.69×), where the simple cp.async `GemmKernel` is launch/memory/tail-bound. This refactor preserves the existing kernel's numerics and perf — closing the small-m and tile-alignment gaps is what the follow-up persistent / split-k work targets. ## Out of scope (follow-ups) - `GemmAccumulateOp` (optional `c`/accumulate as a `variant_of` class). - persistent / split-k / stream-k kernels. - fp8/fp4 GEMM sibling op. - grouped / gemv manifest entries. - repo-wide `kernel_base.autotune()` fix for the tilelang API drift.
…tune kernels (tile-ai#1631) ## Summary Follow-up to tile-ai#1627, which fixed the TileLang 0.1.11 autotune-skip bug (seeded tunable params recorded in the cache key → treated as "already tuned" → sweep skipped, `config=None` → crash) in the base `Kernel.autotune()` and in `SoftmaxKernel` / `LogSumExpKernel`. These custom-autotune kernels build the TileLang `autotune(...)` decorator directly and were not covered: - `deltanet_fwd` / `deltanet_bwd` (per sub-kernel) - `gated_deltanet_fwd` (per sub-kernel) - `deepseek_dsa_decode` (`SparseMlaKernel`) ## Fix - At each autotuned call site, derive the seeded tunable parameter names via `_autotune_initial_kwargs` and pass them as `do_not_specialize` so TileLang excludes them from the cache key and runs the full sweep (guarded so it's only passed when non-empty). - `deepseek_dsa_decode` also passes `supply_prog`: TileLang calls it with the candidate JIT params, but `SparseMlaKernel.supply_prog` takes none (it builds inputs from instance attrs), so it's wrapped (`lambda *args, **kwargs: self.supply_prog()`) to discard the extra arg. ## Test coverage Added one `tune=True` (`full` tier) regression case to `test_deltanet_fwd.py`, `test_deltanet_chunkwise_bwd.py`, `test_gated_deltanet_fwd.py`, and `test_deepseek_dsa_decode.py` so these autotune paths are covered by CI. ## Not included — fp8_lighting_indexer `fp8_lighting_indexer` was in scope initially but is **reverted**: its `supply_prog(q, kv, ...)` needs tensors not available at tune time and requires an instance-attr fallback (a separate change). Deferred to a follow-up. ## Test plan - [x] H200: `DeltaNetFwdKernel(tune=True)` and `SparseMlaKernel(tune=True)` autotune and return a populated config (were `None` / crash). - [x] New tuned cases collect cleanly and satisfy the conftest tier validator. - [x] `ruff` + pre-commit pass. ### Test node delta 4 new nodes (one `tune=True` `full` case per touched test file). `scripts/test_node_delta.py` could not run in the runner image (host-owned workspace not writable by the container user); the growth is exactly those 4 cases.
) ## Problem `FP8LightingIndexerOp(..., tune=True)` crashed before ever producing a valid config. The autotune path had **three independent defects**, all in `FP8LightingIndexerKernel`: 1. **Skip-guard leaves `config=None`** — seeded tunable JIT params landed in the autotuner cache key, so TileLang 0.1.11 treated the kernel as *already tuned*, skipped the sweep, and left `self.config = None`. Downstream `self.config[...]` then raised `TypeError`. 2. **`supply_prog` signature incompatible with the autotuner** — TileLang calls `supply_prog(params)` positionally for each candidate, but the signature was `(self, q, kv, dtype=..., ...)`, so every candidate raised `TypeError` and the sweep found nothing. 3. **`supply_prog` built wrong inputs** even once callable: - `torch.randn(..., dtype=torch.float8_e4m3fn)` → `NotImplementedError: normal_kernel_cuda not implemented for Float8_e4m3fn`. Generate in fp16, then cast. - `IndexK` was cast to `self.dtype`; the kernel signature requires fp8 → `IndexK dtype mismatch, expected float8_e4m3fn`. - The `Logits` output buffer (kernel arg 4 of 7) was missing from the returned tuple. ## Fix - Pass `do_not_specialize=<tunable param names>` to `autotune(...)` so seeded params stay out of the cache key (mirrors the central fix in `kernel_base`). - Rewrite `supply_prog` to `(self, *args, **kwargs)`, building inputs from the kernel's shape attributes and returning all 7 kernel args (incl. `Logits`). ## Test delta `tests/ops/test_fp8_lighting_indexer.py`: +1 node — a `full`-tier `tune=True` case tracing the repaired autotune path. `full`-tier only (excluded from smoke PRs); the full 54-config sweep is the dominant cost there. --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
…y race (tile-ai#1633) ## Summary Fixes tile-ai#1616, plus a defensive fix for the same latent race in the plain MHA/GQA forward WGMMA kernels. The fixed-length `GQASlidingWindowFwdWgmmaPipelinedKernel` (selected on H200) miscomputed non-causal right-window cases — up to ~19% of elements wrong, max abs diff ~0.47 (tolerance 0.01), reproduced deterministically across recent `main` gpu-smoke runs. ## Root cause The kernel writes the output accumulator through a **swizzled** `o_shared` buffer (`acc_o -> o_shared -> output`) with **no barrier** around the round-trip. Across warps this is a shared-memory race: a warp reads `o_shared` before peer warps finish writing it, corrupting the output under WGMMA pipelining (`num_stages >= 2`). The varlen sibling already guards the identical round-trip with `T.sync_threads(3, threads)` (added in the closed tile-ai#1404 "wgmma varlen output shared-memory race"); that guard was never ported to the fixed-length kernel. This PR ports it. **Why a named barrier (slot 3, not bare `T.sync_threads()`):** verified on H200 — a bare `T.sync_threads()` (slot 0) aliases the implicit `__syncthreads()` the surrounding `T.copy` lowering already emits, gets elided, and the race survives. Any non-zero slot works. **Why non-causal right-window specifically:** some query rows have a fully-masked trailing KV tile, which shifts warp scheduling enough to expose the race. Causal / left-only / no-window paths and `num_stages=1` happened to avoid the timing window — which is why those passed and masked the latent bug. ## Commits 1. `gqa_sliding_window_fwd.py` — the tile-ai#1616 fix. 2. `gqa_fwd.py` — defensive fix: `_mha_fwd_wgmma_pipelined_main` / `_gqa_fwd_wgmma_pipelined_main` have the identical unguarded write-back. Plain attention has no masked-tile trigger so it is latent (tests pass), but the hazard is real; barriers only add synchronization. 3. comment clarifying the named-barrier requirement. ## Correctness verification (CI image `ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10-dev`, torch 2.10.0+cu129, H200) | suite | result | |---|---| | `tests/ops/attention/test_gqa_sliding_window.py` | **21/21 passed** (incl. the two previously-failing cases) | | `tests/ops/attention/test_mha.py` + `test_gqa.py` | **55 passed, 4 skipped** (no regression from commit 2) | ## Barrier perf impact — before/after latency (median over 100 iters, same image) The two `T.sync_threads` run once per block at the end (not in the KV loop), so the cost is below run-to-run noise. | kernel | shape | latency before (ms) | latency after (ms) | |---|---|---|---| | GQASlidingWindowFwdWgmmaPipelined | fp16 wl=-1 wr=64 | 0.0167 | 0.0164 | | GQASlidingWindowFwdWgmmaPipelined | bf16 wl=64 wr=64 | 0.0149 | 0.0146 | | MHAFwdWgmmaPipelined | causal fp16 | 0.0162 | 0.0159 | | GQAFwdWgmmaPipelined | causal fp16 | 0.0160 | 0.0157 | | MHAFwdWgmmaPipelined | causal bf16 | 0.0161 | 0.0159 | | GQAFwdWgmmaPipelined | causal bf16 | 0.0158 | 0.0158 | > "before" = the unguarded `origin/main` kernels (temporarily reverted, then restored). The sliding-window "before" output is numerically wrong (the bug); only latency is comparable. ## GQA forward vs reference baselines (`bench_gqa.py`, same image, H200) Metric is **throughput in TFLOPS (higher is better)** — not latency. TileOPs = `GroupedQueryAttentionFwdOp` (kernel variant in parens). Baselines = FlashAttention-3 and FlashInfer. Bold = fastest in the row. ### Compute-bound (S=2048, causal, autotuned) — the representative comparison | model | dtype | TileOPs (TFLOPS) | FA3 (TFLOPS) | FlashInfer (TFLOPS) | |---|---|---|---|---| | Llama-3.1-8B (B=2, H=32, Hkv=8, D=128) | fp16 | 412 (ws_causal) | **508** | 215 | | Llama-3.1-8B (B=2, H=32, Hkv=8, D=128) | bf16 | 107 (wgmma_pipelined) | **288** | 186 | | Llama-3.1-70B (B=1, H=64, Hkv=8, D=128) | fp16 | 212 (ws_causal) | **236** | 216 | | Llama-3.1-70B (B=1, H=64, Hkv=8, D=128) | bf16 | 146 (wgmma_pipelined) | **236** | 186 | ### Small / overhead-bound (S=512, causal) — latency-dominated, not throughput-representative | model | dtype | TileOPs (TFLOPS) | FA3 (TFLOPS) | FlashInfer (TFLOPS) | |---|---|---|---|---| | Llama-3.1-8B (B=4, H=32, Hkv=8, D=128) | fp16 | 215 (ws_causal) | **249** | 88 | | Llama-3.1-8B (B=4, H=32, Hkv=8, D=128) | bf16 | 81 (wgmma_pipelined) | **94** | 92 | | Llama-3.1-70B (B=2, H=64, Hkv=8, D=128) | fp16 | 92 (ws_causal) | 94 | **196** | | Llama-3.1-70B (B=2, H=64, Hkv=8, D=128) | bf16 | 86 (wgmma_pipelined) | 94 | **198** | **Reading the compute-bound numbers:** - **fp16 (ws_causal)** reaches **81–90% of FA3** and matches/beats FlashInfer — TileOPs' strong path. - **bf16 (wgmma_pipelined)** is markedly weaker (**37–62% of FA3**), even behind FlashInfer. For the same causal shape fp16 dispatches to `ws_causal` but bf16 falls back to `wgmma_pipelined`, a separate perf gap unrelated to this fix (the barrier change has no measurable latency impact, see above). Worth a follow-up (bf16 support in the persistent kernel, or better wgmma bf16 autotune). ## Acceptance criteria (tile-ai#1616) - [x] Both failing fixed-length GQA sliding-window cases pass at `atol=1e-2, rtol=1e-2`. - [x] Causal, left-window-only, no-window, and varlen cases remain green. - [x] Right-window masking covered by the existing (now-passing) parametrizations in `test_gqa_sliding_window.py`. --------- Co-authored-by: Ibuki 🍃 — a wind born from GPTs <Ibuki-wind@users.noreply.github.com>
## Summary - Canonicalize GQA inference around fixed manifest-backed OP contracts: - `GroupedQueryAttentionPrefillFwdOp` now owns packed dense/varlen/sliding-window/FP8 prefill dispatch. - `GroupedQueryAttentionPrefillPagedWithKVCacheFwdOp` now owns normal and FP8 paged prefill cache dispatch through fixed scale inputs and `cache_dtype`. - `GroupedQueryAttentionDecodePagedWithKVCacheFwdOp` keeps paged decode as the canonical serving decode path. - Keep compatibility wrappers where useful, including routing square `GroupedQueryAttentionFwdOp` through canonical prefill instead of adding logical Kernel wrappers. - Keep selector helpers local to `tileops/ops/attention/gqa.py` and expose concrete kernels through `default_kernel_map` / manifest `source.kernel_map`. - Thread `sm_scale` and `softcap` through contiguous and paged GQA decode, and widen the Hopper WS prefill path to cover bf16 plus non-default scale/softcap. - Keep canonical inference outputs fixed as `o`; no dynamic `return_lse` output arity is introduced. - Keep backward-only LSE setup out of the production forward op surface: `GroupedQueryAttentionFwdOp.forward()` returns only `o`, while the backward workload constructs the real base-2 LSE in its own setup helper. Closes tile-ai#1610 Closes tile-ai#1611 Closes tile-ai#1622 Closes tile-ai#1623 ## Validation Local validation: - `ruff check tileops/ops/attention/gqa.py workloads/attention/gqa.py tests/ops/attention/test_gqa.py` - `python -m compileall -q tileops/ops/attention/gqa.py workloads/attention/gqa.py tests/ops/attention/test_gqa.py` - `python -m pytest -q tests/ops/attention/test_gqa.py::test_gqa_fwd tests/ops/attention/test_gqa.py::test_gqa_bwd --tb=short` (`10 passed`) - `python -m pytest -q tests/ops/attention/test_gqa.py -m smoke --tb=short` (`58 passed`) - `python scripts/validate_manifest.py --strict --check-op GroupedQueryAttentionFwdOp` - `git diff --check` Earlier broader GQA validation on this PR branch: - targeted prefill/paged/decode/FP8 GQA suite passed before the final output-only forward cleanup; - GQA manifest strict checks for `GroupedQueryAttentionPrefillFwdOp`, `GroupedQueryAttentionPrefillPagedWithKVCacheFwdOp`, `GroupedQueryAttentionDecodeWithKVCacheFwdOp`, and `GroupedQueryAttentionDecodePagedWithKVCacheFwdOp` passed. GitHub CI on latest push: - `pre-commit`: passing - `validate-manifest`: passing - `benchmark-contract-tests`: passing - `gpu-smoke`: passing
Closes tile-ai#1521 ## Summary - Add native grouped Conv2d/Conv3d TileLang kernels for bias and no-bias variants. - Wire Conv2d/Conv3d op dispatch to grouped kernels while preserving existing dilation support from `main`. - Update the convolution manifest contract for grouped Conv2d/Conv3d status, kernel maps, and per-group roofline accounting. - Add grouped correctness coverage and model-derived grouped convolution benchmark cases. ## Benchmark ```bash CUDA_VISIBLE_DEVICES=1 TILELANG_CLEANUP_TEMP_FILES=1 python -m pytest \ 'benchmarks/ops/bench_convolution.py::test_conv2d_bench[mobilenetv2-depthwise-fp16]' \ 'benchmarks/ops/bench_convolution.py::test_conv2d_bench[resnext-grouped-3x3-fp16]' \ 'benchmarks/ops/bench_convolution.py::test_conv3d_bench[3d-resnext-grouped-k3-fp16]' \ 'benchmarks/ops/bench_convolution.py::test_conv3d_bench[3d-resnext-grouped-k3-b8-fp16]' \ -vvs ``` Result: `4 passed in 138.39s`. | Case | TileOps latency | TileOps TFLOP/s | TileOps bandwidth | Torch latency | Torch TFLOP/s | Torch bandwidth | | --- | ---: | ---: | ---: | ---: | ---: | ---: | | `conv2d mobilenetv2-depthwise-fp16` | 0.0064 ms | 0.2812 | 0.0626 TB/s | 0.0065 ms | 0.2772 | 0.0617 TB/s | | `conv2d resnext-grouped-3x3-fp16` | 0.0041 ms | 3.4882 | 0.1498 TB/s | 0.0200 ms | 0.7207 | 0.0309 TB/s | | `conv3d 3d-resnext-grouped-k3-fp16` | 0.0147 ms | 5.8872 | 0.1645 TB/s | 0.7442 ms | 0.1165 | 0.0033 TB/s | | `conv3d 3d-resnext-grouped-k3-b8-fp16` | 0.1269 ms | 5.4643 | 0.1519 TB/s | 5.7221 ms | 0.1212 | 0.0034 TB/s |
tile-ai#1634) > **Hotfix / stopgap.** Unblocks CI and closes a deadlock footgun. Does **not** fix the root design issue (`block_Q` is exposed as an independent autotune knob while it is coupled to `heads`); see Follow-up. ## Problem `FP8LightingIndexerOp(..., tune=True)` hangs until the per-test timeout and fails the job. It surfaced on push-to-main / nightly but not on the introducing PR (tile-ai#1632), because `gpu-smoke` runs `-m smoke` on PRs and `-m "smoke or full"` only on push — so the `full`-tier `tune=True` node never ran on that PR's gate. ## Root cause Mapped on H200 across heads ∈ {32, 64, 128}: **the kernel deadlocks iff `block_Q == 1` and `num_stages >= 1`** (some warps never reach the software-pipeline barrier; the launch spins forever). An abandoned hung launch wedges the GPU, so every later autotune candidate times out too — the sweep never finishes within budget. Also a non-autotune footgun: the default `block_Q = 128 // heads` is `1` for `heads >= 65`. ## Fix - Force `num_stages = 0` when `block_Q < 2` (guards autotune, manual config, and the default path). - Harden default to `block_Q = max(1, 128 // heads)`. - Drop `block_Q == 1` from `autotune_configs`. Verified on H200: previously-hanging configs now run; `tune=True` test passes in ~9s warm / ~54s cold (was: timeout). ## Follow-up Clean fix is to tune the q-tile and *derive* `block_Q` so degenerate configs can't be expressed (removes the guard). The silent deadlock is also a tilelang pipeline-codegen robustness bug.
Remove the `cache-warmup` job from `nightly.yml`. - Delete the `cache-warmup` job. - `benchmark` becomes the entry job: drop `needs: [cache-warmup]` and the now-redundant `always() &&`. - Renumber phase comments (benchmark→1, op_test→2, publish→2b, packaging→3). - Keep `scripts/warmup_kernel_cache.py` + `scripts/conftest_warmup.py` for manual one-off warmup. Redundant under the persistent `/ci-cache`: the benchmark phase autotunes at op construction (outside the timed region) and writes results back, so it self-warms across runs; benchmark numbers are unchanged. The job never gated anything (warmup step `continue-on-error`, benchmark `if: always()`) but turned the nightly red on runner network blips.
…anup (tile-ai#1637) ## Summary - nightly `op_test` runs `-m "full or nightly"` instead of the whole suite; the `smoke` tier is already covered by every PR and every push-to-main gpu-smoke. - `reclaim-runner-disk` drops the venv-prune pass (venvs retired) and the `RUNNER_TEMP` orphan-purge pass (ephemeral `--rm` jobs leave no cross-run orphans), plus their inputs `venv-keep`, `temp-orphan-minutes`, `current-run-dir`, `prune-tool-cache`. - `reclaim-runner-disk` keeps only the `/ci-cache` file-trim, atomic autotuner sentinel-repair/age-trim, and post-reclaim disk floor. - `gpu-smoke.yml` drops the removed `prune-tool-cache` input; `runner-maintenance.yml` and `reclaim_cache.sh` are unchanged. ## Test plan - [x] all workflow/action YAML parses (`yaml.safe_load`) - [x] `tests/test_reclaim_action.py` 13/13 pass (`reclaim_cache.sh` untouched) - [x] pre-commit passed
…ne trace tool (tile-ai#1638) ## Summary - **Warp-specialized GEMM kernel.** Rewrites the dense `GemmKernel` as a hand-written two-warpgroup warp-specialized implementation for Hopper (SM90): a producer warpgroup issues TMA loads into a double-buffered SMEM ring, a consumer warpgroup runs WGMMA over K. Covers all four `(trans_a, trans_b)` layouts via the WGMMA transpose flags. fp16/bf16 in, fp32 accumulate. The `GemvKernel` fast path (m==1 / n==1) is unchanged. - **In-kernel timeline trace tool (`tileops/trace/`).** Annotate a kernel body with `trace.group` / `range` / `record` / `dag`; `trace.out_idx` + `trace.finalize` wire the markers into a vanilla `@tilelang.jit` builder (lowered when tracing is on, stripped to zero cost when off); `trace.run` executes the kernel and dumps a self-contained Plotly HTML timeline. Process-local switch via `trace.enable()` — no monkeypatch, no env var. `GemmKernel` is annotated; opt in per pytest run with `--trace-kernel`. ## Motivation Small-m dense NT GEMM (e.g. 128×2112×7168) was running well under cuBLAS. The WS kernel targets that regime, and the trace tool provides a timeline view (kernel execution / gaps / producer–consumer overlap) to diagnose it. ## Notes / tradeoffs - **SM90-only.** WS relies on TMA + WGMMA, so `GemmKernel.supported_archs = [90]`. - **TMA alignment.** Each input's innermost (contiguous) dimension must be a multiple of 8 fp16 elements (16-byte). For NT (`A @ Bᵀ`) that is `K` only; other layouts additionally require the relevant `M` / `N` to be aligned. Documented on the kernel. - `_gemm_wrapped_kernel` (torch custom op) is kept for `torch.compile` compatibility; the eager forward calls the compiled JIT directly (cf. `GemvKernel`). ## Tests - All four layouts (NN/NT/TN/TT) verified numerically against torch (max abs err 0.0), including the `transpose_A` paths. - Added an NT dense numerics case to `tests/ops/test_gemm.py`; suite passes (23 passed). ## Follow-ups (not in this PR) - Unit tests for the `tileops/trace/` package. - Promote the trace-tool design doc into `docs/design/`.
…tile-ai#1608) ## Summary - Update default tuning configs for SSD chunk_scan and chunk_state kernels based on NCU profiling. ## Test plan - [ ] Run benchmarks/ops/bench_mamba.py for chunk_scan and chunk_state ### ssd_chunk_scan: threads=64 (before) vs threads=128 (after, current default) > Config: `{block_l: 64, block_p: 64, block_n: 64, block_s: 64, threads: 64/128}` | Shape | before (t=64) ms | after (t=128) ms | Speedup | |---|---|---|---| | 130m-b1-4k | 0.1137 | 0.0992 | **1.146x** | | 130m-b8-4k | 0.5349 | 0.4107 | **1.303x** | | 130m-b4-32k | 1.9326 | 1.4292 | **1.352x** | | 370m-b1-4k | 0.1337 | 0.1122 | **1.192x** | | 370m-b8-4k | 0.6749 | 0.5033 | **1.341x** | | 370m-b4-32k | 2.5624 | 1.8263 | **1.403x** | | 370m-b32-2k | 1.2850 | 0.9423 | **1.364x** | | 780m-b1-4k | 0.1778 | 0.1452 | **1.225x** | | 780m-b8-4k | 0.9909 | 0.7403 | **1.339x** | | 780m-b4-32k | 3.7741 | 2.7922 | **1.352x** | | 780m-b16-2k | 0.9901 | 0.7428 | **1.333x** | | 1p3b-b1-4k | 0.2205 | 0.1733 | **1.273x** | | 1p3b-b8-4k | 1.2873 | 0.9422 | **1.366x** | | 1p3b-b2-32k | 2.5459 | 1.8293 | **1.392x** | | 1p3b-b8-2k | 0.6737 | 0.5040 | **1.337x** | | 2p7b-b1-4k | 0.2596 | 0.2067 | **1.256x** | | 2p7b-b4-4k | 0.8391 | 0.6356 | **1.320x** | | 2p7b-b2-32k | 3.1504 | 2.3596 | **1.335x** | | 2p7b-b4-2k | 0.4594 | 0.3500 | **1.313x** | **threads=128 wins on every single production shape, by 15–40%.** The pre-tile-ai#1608 default of threads=64 is consistently slower. The threads=128 change in this PR is correct and necessary. --- ### ssd_chunk_state: old config (block_n=128, block_p=64, block_l=64, threads=256) vs new (block_n=min(128,N), block_p=64, block_l=32, threads=128) | Shape | before ms | after ms | Speedup | |---|---|---|---| | 370m-b1-L2k | 0.0706 | 0.0704 | 1.003x | | 370m-b4-L2k | 0.0887 | 0.0867 | 1.023x | | 1p3b-b1-L2k | 0.0707 | 0.0707 | 1.000x | | 1p3b-b4-L2k | 0.0896 | 0.0852 | 1.051x | | 2p7b-b1-L2k | 0.0810 | 0.0834 | 0.972x | | 2p7b-b4-L2k | 0.0837 | 0.0818 | 1.023x | | 1p3b-b1-L8k | 0.0702 | 0.0675 | 1.039x | The chunk_state config change is largely neutral (±3%), with a small regression at 2.7b-b1 (-2.8%). The main benefit from the chunk_state changes in this PR is the `threads=256 → 128` reduction, which reduces register pressure. The `block_l=64→32` change has minimal latency impact at these shapes.
…-ai#1640) ## What Makes the `tileops.trace` module documentation-ready and adds a usage tutorial, so the trace tool renders cleanly on the docs site (mkdocstrings + Material). - **Docstrings: Sphinx roles → plain code spans.** Across the whole `tileops/trace/` package (`__init__`, `api`, `decode`, `markers`, `passes`, `record`, `state`, `ui`), replace `:func:` / `:class:` / `:mod:` / `:data:` cross-reference roles with double-backtick code spans. The project is Google-docstring-only; mkdocstrings does not resolve Sphinx roles, so they rendered as literal `:func:`...`` text on the site. - **`ui.py` timeline legend fix.** The horizontal legend overlapped the x-axis title. Give the x-axis title a `standoff`, anchor the legend below it (`y: -0.30`, `yanchor: top`), and grow the bottom margin (120 → 170). - **New tutorial: `docs/perf/trace-timeline.md`.** End-to-end walkthrough — write a traced warp-specialized GEMM, run it, enable tracing, read the timeline. Each code block uses Material code annotations (the `+` markers) linked to the auto-generated API reference. - **`docs/perf/README.md`**: add a "Tools & Guides" entry pointing at the tutorial. ## Why The trace tool shipped in tile-ai#1638 with no rendered API docs or guide; this makes its public surface (`tileops.trace.api`) render correctly and gives users a single walkthrough. ## Notes - No behavior change: docstring text and one Plotly layout dict only. - The site-side wiring (mkdocstrings page, nav, embedded timeline) is a separate PR against `tile-ai/TileOPs.github.io`. Closes tile-ai#1639
) ## Summary - Add explicit dense-backend GQA prefill manifest workloads for custom sm_scale and softcap. - Thread backend, sm_scale, and softcap from manifest workloads into the dense prefill benchmark. - Apply the same sm_scale/softcap settings to the torch reference and FlashInfer baseline so benchmark comparisons use matching semantics. - Pass 0.0 instead of None for FlashInfer logits_soft_cap when softcap is unset. This PR intentionally updates only manifest/benchmark coverage. The high-performance kernel scheduling updates will be submitted separately after this lands. ## Testing - python -m py_compile benchmarks/ops/attention/bench_gqa.py benchmarks/ops/attention/manifest_params.py - python manifest workload expansion check for GroupedQueryAttentionPrefillFwdOp - git diff --check - docker/CI image: pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py --collect-only -q - docker/CI image: pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench and softcap50' - docker/CI image: pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench and sm and scale' - docker/CI image: pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench and llama-3.1-8b-prefill-dense-float16'
## Summary - Dispatch dense GQA prefill to Hopper warp-specialized kernels for H200 fp16/bf16 causal dim-128 paths. - Add a square dense prefill fast path that reuses the high-performance square GQA WS kernel and supports custom `sm_scale` / `softcap`. - Avoid timed CUDA `torch.equal` checks for explicit dense backend inputs by validating fixed packed lengths directly. - Make paged prefill forward kernels output-only for inference paths and keep RoPE / FP8-cache / softcap variants on the canonical paged OP. - Add focused tests for dense dispatch, feature variants, and dense backend packed-input validation. This follows tile-ai#1641, which added the manifest/benchmark coverage for scale and softcap cases. ## Performance Notes H200, torch 2.10.0+cu129, CUDA 12.9, driver 575.57.08. Dense prefill benchmark (`test_gqa_prefill_fwd_bench`): - B4 S512 H32/Hkv8 D128 fp16: TileOps 0.0394 ms vs FlashInfer 0.0397 ms - B4 S512 H32/Hkv8 D128 bf16: TileOps 0.0393 ms vs FlashInfer 0.0398 ms - B2 Sq512 Skv4096 H32/Hkv8 D128 fp16: TileOps 0.1247 ms vs FlashInfer 0.1266 ms - B2 Sq512 Skv4096 H32/Hkv8 D128 bf16: TileOps 0.1240 ms vs FlashInfer 0.1258 ms - B1 Sq512 Skv4096 H64/Hkv8 D128 fp16: TileOps 0.1233 ms vs FlashInfer 0.1245 ms - B1 Sq512 Skv4096 H64/Hkv8 D128 bf16: TileOps 0.1223 ms vs FlashInfer 0.1239 ms Focused feature benchmarks: - sm_scale=0.125: TileOps 0.0391 ms vs FlashInfer 0.0395 ms - softcap=50.0: TileOps 0.0455 ms vs FlashInfer 0.0459 ms Paged representative benchmarks: - Llama3.1 paged RoPE: 1.9863 ms - paged softcap: 0.1759 ms - Llama3.1 FP8 KV cache: 2.0418 ms - FP8 KV cache softcap: 0.2440 ms ## Testing - `python -m py_compile tileops/kernels/attention/gqa_prefill_fwd_ws.py tileops/kernels/attention/gqa_fwd_ws.py tileops/kernels/attention/gqa_fwd.py tileops/ops/attention/gqa.py tests/ops/attention/test_gqa.py` - `git diff --check` - docker/CI image: `pytest /workspace/TileOPs/tests/ops/attention/test_gqa.py -q -k 'prefill_fwd'` -> 19 passed - docker/CI image: `pytest /workspace/TileOPs/tests/ops/attention/test_gqa_prefill_paged.py -q` -> 27 passed - docker/CI image: `pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py --collect-only -q` -> 34 collected - docker/CI image: `python scripts/validate_manifest.py --levels schema,signature,shape,dtype,bench --strict` - docker/CI image: `python -m pytest -q benchmarks/tests` -> 8 passed - docker/CI image: `pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench'` -> 8 passed - docker/CI image: `pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench and sm and scale'` -> 1 passed - docker/CI image: `pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_fwd_bench and softcap50'` -> 1 passed - docker/CI image: `pytest /workspace/TileOPs/benchmarks/ops/attention/bench_gqa.py -q -k 'test_gqa_prefill_paged_with_kv_cache_fwd_bench and (llama31 or softcap50)'` -> 4 passed
## Summary - fix `MHCPostKernel` shared-memory indexing to use `block_C` as the row stride for `x_out_shared` - avoid overlapping writes across `n_expand` rows in the default `n_expand=4, block_C=64` configuration Fixes tile-ai#1584. ## Tests - `python -m py_compile tileops/kernels/mhc/mhc_post.py tileops/ops/mhc.py tests/ops/test_mhc_post.py` - `python -m ruff check tileops/kernels/mhc/mhc_post.py tests/ops/test_mhc_post.py` - Docker runner: `python -m pytest tests/ops/test_mhc_post.py -vvs -k smoke` - Docker runner: `compute-sanitizer --tool racecheck python -m pytest tests/ops/test_mhc_post.py -q -k smoke` - `RACECHECK SUMMARY: 0 hazards displayed (0 errors, 0 warnings)` - Docker runner: `python -m pytest tests/ops/test_mhc_post.py -q` Host pytest was not used for validation because the host environment is missing `tvm.tirx`; the Docker runner image includes the TileOps CI dependencies.
…ile-ai#1645) Closes tile-ai#1644 ## Summary This PR addresses the poor `conv2d` performance reported in tile-ai#1644 (CI run `28398916051`) by introducing `Conv2dSymmetricKernel`, an NHWC-implicit-GEMM path for symmetric conv2d cases (`kernel_h == kernel_w`, `groups == 1`, `c_in % 32 == 0`). Key changes: - Added `Conv2dSymmetricKernel` and dispatched symmetric conv2d cases to it. - Unified the three layout transpose macros (`nchw_to_nhwc_input`, `kcrs_to_krsc_weight`, `nhwc_to_nchw_output`) into a single `transpose_spatial_channel` macro. - Parameterized tile size and fastest-varying dimension so the transpose kernels can be tuned per shape in follow-up work. - Reduced the weight transpose `spatial_block` from `32` to `16`, which increases the spatial grid for small kernels (`3x3`, `5x5`, `7x7`). - Added compile-time asserts for tile/thread constraints. No public API is modified. ## Test plan - [x] `python -m py_compile tileops/kernels/convolution.py` passed - [x] `tests/ops/test_convolution.py -k "test_conv2d"` — 20 passed - [x] `tests/ops/test_convolution.py -k "dispatches"` — 6 passed ## Benchmark Comparison of this branch (`Conv2dSymmetricKernel`) against `upstream/main` (`Conv2dKernel`) on H200, clock-locked 1500 MHz: | case | main Conv2dKernel (ms) | this branch Conv2dSymmetricKernel (ms) | speedup | |------|------------------------|----------------------------------------|---------| | deeplabv3-aspp-3x3-rate12-fp16 | 0.4278 | 0.1152 | **3.71x** | | stage-transition-5x5-s2-fp16 | 0.0748 | 0.0242 | **3.09x** | | midres-5x5-s1-fp16 | 0.0392 | 0.0183 | **2.14x** | | stride2-bf16 | 0.0279 | 0.0136 | **2.05x** | | stage-transition-3x3-s2-fp16 | 0.0279 | 0.0163 | **1.71x** | | resnet-3x3-fp16 | 0.0169 | 0.0150 | **1.13x** | `Conv2dSymmetricKernel` is faster on all representative symmetric cases, with larger gains on bigger/dilated kernels. Further improvements are expected once the new transpose tile parameters are exposed to autotune. ## Additional context Addressed Gemini Code Assist review feedback: - Renamed local `h`/`w` to `h_idx`/`w_idx` in `transpose_spatial_channel` to avoid shadowing outer parameters. - Switched to interleaved channel mapping for coalesced NHWC writes in the `channel_fastest` branch. - Added `c_in % block_k == 0` filtering in `Conv2dSymmetricKernel.autotune_configs`. - Made bias addition conditional on `has_bias` in `conv_nhwc_implicit_gemm_bias`. This change is the first step toward a larger performance fix for tile-ai#1644. The unified macro makes it straightforward to add `input_spatial_block`, `output_spatial_block`, etc. to `Conv2dSymmetricKernel.autotune_configs` in a follow-up PR.
…1647) ## Summary - keep public packed `backend="dense"` and `backend="fp8"` GQA prefill safe by validating uniform `cu_seqlens` by default before reshaping to BSHD - add an explicit public `validate_uniform_cu_seqlens` switch so high-performance callers with known-uniform packed metadata can skip the CUDA value check without using a separate benchmark-only path - keep `backend="auto"` requiring uniform validation because auto dispatch needs the `cu_seqlens` values to choose dense vs varlen - declare `validate_uniform_cu_seqlens` in the op manifest and mark known-uniform dense/fp8 manifest workloads with `validate_uniform_cu_seqlens=False` - update dense and fp8 prefill benchmarks to consume the manifest setting and call the same public op path users call - add tests covering default dense/fp8 ragged rejection, explicit dense validation skip, auto validation requirements, varlen-style skip behavior, and the BSHD wrapper path This addresses the dense packed-input contract concern from review while keeping benchmark and user calls on the same public API. ## Testing - `python -m ruff check tileops/ops/attention/gqa.py tests/ops/attention/test_gqa.py benchmarks/ops/attention/bench_gqa.py benchmarks/ops/attention/bench_gqa_fp8.py benchmarks/ops/attention/manifest_params.py` - `python -m py_compile tileops/ops/attention/gqa.py tests/ops/attention/test_gqa.py benchmarks/ops/attention/bench_gqa.py benchmarks/ops/attention/bench_gqa_fp8.py benchmarks/ops/attention/manifest_params.py` - `git diff --check` - Docker (`ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10`): `python scripts/validate_manifest.py --check-op GroupedQueryAttentionPrefillFwdOp --levels schema,signature,shape,dtype,bench --strict` - Docker (`ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10`): `python -m pytest tests/ops/attention/test_gqa.py -q -k "explicit_varlen_backends or explicit_dense_backends or explicit_dense_can_skip or auto_backend_requires or bshd_wrapper or auto_backend_checks"` - Docker as root for writable TileLang cache/benchmark log: `python -m pytest benchmarks/ops/attention/bench_gqa.py::test_gqa_prefill_fwd_bench[llama-3.1-8b-prefill-dense-float16] -q -s` Host pytest was not used for validation because the host environment is missing `tvm.tirx`; the Docker runner image includes the TileOps CI dependencies.
## Summary Closes tile-ai#1648. This updates the initial tile-ai#1648 scope so derivable shape metadata is inferred from op inputs instead of being required as constructor arguments: - Norm family: `FusedAddLayerNormFwdOp`, `FusedAddRMSNormFwdOp`, `AdaLayerNormFwdOp`, `AdaLayerNormZeroFwdOp` - Scan family: `CumsumFwdOp`, `CumprodFwdOp` - MoE leaf ops: `FusedTopKOp`, `MoePermuteNopadFwdOp` The old explicit constructor arguments remain supported as compatibility/strict-validation commitments, while the preferred API derives `M`, `N`, token counts, top-k shape, hidden size, and dtype from the input tensors at `forward()` time. Kernels are built lazily and cached by the inferred specialization. ## Manifest / callers - Removed input-derivable `static_dims` from the affected norm and scan manifest entries. - Updated `MoePermuteNopadFwdOp` manifest params/workloads/roofline to treat `hidden_states` and `topk_ids` shapes as the source of truth, keeping `num_experts` explicit. - Updated direct benchmarks and composite call sites to use the preferred input-inferred APIs. - Added compatibility and cache-reuse/changing-shape test coverage. ## Testing Using the TileOpsGov runner image `ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10` on H200 GPU: - `python -m pytest -q tests/ops/test_cumulative.py tests/ops/test_fused_add_layer_norm.py tests/ops/test_fused_add_rms_norm.py tests/ops/test_ada_layer_norm.py tests/ops/test_ada_layer_norm_zero.py tests/ops/test_moe_fused_topk.py tests/ops/test_moe_permute_nopad.py -m smoke --tb=short --timeout=900 --timeout-method=thread` - `67 passed, 52 deselected, 14 warnings` - `python -m pytest -q tests/ops/test_cumulative.py tests/ops/test_fused_add_layer_norm.py tests/ops/test_fused_add_rms_norm.py tests/ops/test_ada_layer_norm.py tests/ops/test_ada_layer_norm_zero.py tests/ops/test_moe_fused_topk.py tests/ops/test_moe_permute_nopad.py --tb=short --timeout=900 --timeout-method=thread` - `119 passed, 14 warnings` - `python scripts/validate_manifest.py` - `All manifest checks passed` Note: the validator reports an advisory warning that `MoePermuteNopadFwdOp.num_experts` has an `__init__` default of `None` while the manifest has no default. This is intentional to preserve old positional construction compatibility; `forward()` still requires `num_experts` to be provided.
…ads (tile-ai#1652) ## Summary Adds a warp-specialized (Hopper) batch=1 GQA decode kernel and routes the `GroupedQueryAttentionDecodeWithKVCacheFwdOp` to it when the fast-path preconditions hold. Migrated from a validated PoC and hardened for production (runtime `real_seqlen_kv` masking + split redistribution). Includes the batch=1 workloads and a proper FlashInfer baseline needed to benchmark it. ## What's in it **Workloads + benchmark baseline** (`manifest`, `bench_gqa_decode.py`) - Eight batch=1 Qwen3-30B-A3B decode workloads (KV 1K–256K). - FlashInfer baseline now uses `single_decode_with_kv_cache` (contiguous KV, no paging overhead) for batch=1 — the apples-to-apples comparison; the paged batched wrapper is kept for batch>1. **Kernel + dispatch** (`gqa_decode_bs1.py`, `gqa.py`, tests, `kernel_map`) - `GQADecodeBs1Kernel` (`supported_archs=[90]`): TMA + mbarrier + wgmma warp-specialized producer/consumer, online softmax in exp2 domain, fp32 accumulators, ctx-split + LSE-combine. - Op-level ctor gate selects the fast-path kernel only for batch=1, Hopper, fp16, D=128, softcap=0, group in [1,64]; everything else keeps the incumbent `GQADecodeKernel`. sm80/89 stay constructible. - `forward()` dispatches by runtime `real_seqlen_kv`: `no_split` below 1024, ctx-split above (largest of {32,16,8} that divides `real`, balanced load). Non-aligned lengths (e.g. 2049) go through this kernel, not a fallback. ## Correctness 14 tests pass (`tests/ops/attention/test_gqa_decode.py`), including runtime context-switch and non-divisible-length cases; atol=rtol=1e-2 vs the fp32 reference. Hopper-only tests skip on other arches. ## Performance (H200 @ 1500MHz, CUPTI, H=32 Hkv=4 D=128, fp16) | KV len | new (µs) | incumbent | FlashInfer single | new/incumbent | new/FI | |-------:|---------:|----------:|------------------:|--------------:|-------:| | 1024 | 7.97 | 16.16 | 10.35 | 0.49 | 0.77 | | 2048 | 8.94 | 10.65 | 11.12 | 0.84 | 0.80 | | 3072 | 11.93 | 24.07 | 11.86 | 0.50 | 1.01 | | 4096 | 10.70 | 12.78 | 12.40 | 0.84 | 0.86 | | 8192 | 14.00 | 16.20 | 14.81 | 0.86 | 0.95 | | 32768 | 29.49 | 36.52 | 35.29 | 0.81 | 0.84 | | 131072 | 77.53 | 119.36 | 83.88 | 0.65 | 0.92 | Faster than the incumbent across the board; faster than or on par with FlashInfer `single_decode` everywhere (roughly tied at 3072). ## Trust-model note Kept as a single PR per request. It mixes a manifest workloads change (normally a separate human-reviewed PR) with the implementation; the manifest edits are the eight workloads plus one `source.kernel_map` line (the latter is within the implementation-PR carve-out).
## Summary Closes tile-ai#1659. Refs tile-ai#1650. Promotes implemented special elementwise ops from `spec-only` to `implemented`: - `SiluAndMulFwdOp`, `GeluAndMulFwdOp`, `GeluTanhAndMulFwdOp` - `AlibiFwdOp`, `SinusoidalFwdOp` This is a code/manifest cleanup PR, not a release-plan doc update. tile-ai#1659 is tracked as a sub-issue of the broader tile-ai#1650 release-facing manifest coverage tracker. ## Changes - Allow fused gated ops to bind `M`, `N`, and `dtype` from the runtime input tensor, while preserving the old explicit ctor path. - Keep lazy fused gated ops reusable across runtime shapes/dtypes by separating explicit ctor contracts from the last runtime-bound kernel specialization. - Ensure `_eager_forward` validates the runtime input and calls `_ensure_kernel(M, N, x.dtype)` before launch. - Add fused gated `eval_roofline()` support and the correct GELU-tanh FLOPs coefficient. - Add zero-input `_validate_dtypes()`, `_infer_output_shapes()`, `total_memory`, and `eval_roofline()` for `AlibiFwdOp` and `SinusoidalFwdOp`. - Keep generative op `dtype` explicit in the manifest because these ops have no input tensor to infer dtype from. - Add smoke tests for manifest-style fused gated lazy binding and shape rebinding. ## Validation Ran in `ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10`: ```bash python -m pytest tests/test_ops_manifest.py -q # 21 passed python -m pytest tests/test_ops_manifest.py::TestOpSchema::test_workloads_include_all_required_params -q # 1 passed for op in SiluAndMulFwdOp GeluAndMulFwdOp GeluTanhAndMulFwdOp AlibiFwdOp SinusoidalFwdOp; do python scripts/validate_manifest.py --check-op "$op" --strict done # all targeted checks passed python -m pytest \ tests/ops/test_fused_gated.py::test_silu_and_mul_infers_manifest_shape_contract \ tests/ops/test_fused_gated.py::test_silu_and_mul_lazy_op_rebinds_shape \ tests/ops/test_fused_gated.py::test_fused_gated_rejects_integer_dtype -q # 3 passed python -m ruff check \ tileops/ops/elementwise/_base.py \ tileops/ops/elementwise/fused_gated.py \ tileops/ops/elementwise/alibi.py \ tileops/ops/elementwise/sinusoidal.py \ tests/ops/test_fused_gated.py # All checks passed ``` GitHub CI is green, including `validate-manifest`, `benchmark-contract-tests`, and `gpu-smoke`.
…ock factorization (tile-ai#1646) ## Summary This PR optimizes the Mamba-2 SSD chunk scan forward kernel with two key improvements: 1. **Auto-pipeline optimization**: Leveraging TileLang's built-in auto-pipeline for better instruction overlap 2. **Micro-block factorization**: Reducing exponential operations in diagonal blocks via factorized computation **Performance**: ~4-5% speedup on representative workloads while maintaining numerical correctness. ## Motivation The SSD chunk scan kernel is compute-bound on the diagonal path exponential operations. The original implementation had two bottlenecks: 1. **Manual pipelining limitations**: Hand-written double-buffering didn't use async copies, limiting memory-compute overlap 2. **Redundant exponentials**: Each element in diagonal blocks computed `exp(dA[l] - dA[s])` independently, with ~6% theoretical overhead ## Changes ### 1. Auto-Pipeline Optimization **Before**: Manual double-buffer implementation without async memory operations ```python # Manual pipeline with synchronous loads for n_blk in range(T.ceildiv(N, block_n)): # Load, compute, store in lock-step ``` **After**: TileLang auto-pipeline with optimal staging ```python for n_blk in T.Pipelined(T.ceildiv(N, block_n), num_stages=3): # Auto-managed async loads with optimal overlap ``` **Impact**: - Optimal `num_stages=3` determined via benchmarking (tested 1, 3, 5) - Better overlap of global memory loads with tensor core operations - Cleaner code with equivalent or better performance ### 2. Micro-block Factorization for Diagonal Path **Optimization**: For diagonal blocks where `s` and `l` are close, factorize the exponential: ``` exp(dA[l] - dA[s]) = exp(dA[l] - anchor) * exp(anchor - dA[s]) ``` where `anchor = dA[l_block_start]` for each 32×32 micro-block. **Before**: 64×64 direct computation ```python for ll, ss in T.Parallel(64, 64): # 4096 exp calls exp_factor = T.exp(dA[l] - dA[s]) lcb_cast[ll, ss] = cb[ll, ss] * exp_factor * dt[ss] ``` **After**: 2×2 micro-block factorization (for block_l=block_s=64) ```python # Precompute factorized terms (64 + 2*64 = 192 exp calls) for ll in T.Parallel(64): exp_l_micro[ll] = T.exp(dA[ll] - anchor[r]) for r, ss in T.Parallel(2, 64): exp_s_micro[r, ss] = T.exp(anchor[r] - dA[ss]) * dt[ss] # Four 32×32 micro-blocks: # - Diagonal blocks (mr=0,mc=0 and mr=1,mc=1): direct exp (2*1024 = 2048) # - Lower block (mr=1,mc=0): factorized (1024, reuses precomputed) # - Upper block (mr=1,mc=1): zeros (masked out by causality) ``` **Exponential count reduction**: - Original: 4096 exp operations per 64×64 block - Optimized: 192 (precompute) + 2048 (diagonal micro-blocks) = 2240 exp operations - **Reduction: ~45% fewer exp calls in diagonal blocks** **Implementation note**: Currently uses Python static unrolling for the 2×2 decomposition, limiting applicability to `block_l=block_s=64` configuration. Future work could generalize this with dynamic tiling. ## Numerical Correctness Micro-block factorization introduces two-step rounding but maintains acceptable numerical accuracy: **Error Analysis**: - Mean absolute error: 1.97e-06 (negligible) - Max absolute error: 7.75e-02 (rare outliers, 0.0003% of elements) - All outputs pass `torch.allclose(rtol=1e-3, atol=1e-3)` against reference Error sources: 1. Two-stage rounding: `exp(a-b)` vs `exp(a-c) * exp(c-b)` 2. Both mathematically equivalent but different floating-point paths 3. Verified no data races; errors are deterministic from FP arithmetic The error is acceptable for ML workloads and significantly below typical FP16 training noise. ## Testing - ✅ All existing unit tests pass (`test_ssd_chunk_scan_fwd`) - ✅ Numerical correctness verified against PyTorch reference - ✅ Tested on smoke and full test configurations - ✅ No regressions on non-optimized code paths ## Three-Way Comparison | Config | TileOPs-OLD (ms) | TileOPs-NEW (ms) | mamba_ssm (ms) | NEW vs OLD | NEW vs mamba_ssm | |--------|------------------|------------------|----------------|------------|------------------| | latency-130m-4k | 0.0951 | 0.0930 | 0.1420 | **1.02x faster** ✅ | **1.53x faster** ✅ | | serving-130m-4k | 0.3991 | 0.3403 | 0.4395 | **1.17x faster** ✅ | **1.29x faster** ✅ | | longctx-130m-32k | 0.7445 | 0.6456 | 0.7952 | **1.15x faster** ✅ | **1.23x faster** ✅ | | latency-370m-4k | 0.1692 | 0.1465 | 0.2106 | **1.15x faster** ✅ | **1.44x faster** ✅ | | serving-370m-4k | 0.9359 | 0.7946 | 1.0044 | **1.18x faster** ✅ | **1.26x faster** ✅ | | longctx-370m-32k | 1.8240 | 1.5885 | 1.9419 | **1.15x faster** ✅ | **1.22x faster** ✅ | | throughput-370m-2k | 1.8128 | 1.5905 | 1.9422 | **1.14x faster** ✅ | **1.22x faster** ✅ | | latency-2p7b-4k | 0.2012 | 0.1755 | 0.2466 | **1.15x faster** ✅ | **1.41x faster** ✅ | | serving-2p7b-4k | 1.2003 | 1.0108 | 1.2210 | **1.19x faster** ✅ | **1.21x faster** ✅ | ## Future Work 1. **Generalize micro-block size**: Make factorization work for arbitrary block sizes, not just 64×64 2. **Auto-tuning**: Dynamically select micro-block size based on problem dimensions 3. **Larger workloads**: Test on production-scale sequence lengths (4K-32K tokens) ## References - Original implementation: `ssd_chunk_scan_before_microtile.py` - Mamba-2 paper: https://arxiv.org/abs/2405.21060 - TileLang auto-pipeline: https://github.com/TileOPs/TileLang --- **Checklist**: - [x] Code follows project style guidelines - [x] Tests pass - [x] Performance benchmarks included - [x] Numerical accuracy verified - [x] Documentation updated
## Summary Adds a new performance profile for the **NVIDIA H20-3e** (SXM, 141GB HBM3e) at `tileops/perf/profiles/h20_3e.yaml`, following the same schema as the existing `h200.yaml`. ## What's included | Metric | Theoretical | Calibration | |---|---|---| | HBM bandwidth | 4.80 TB/s | **0.8083** | | Tensor Core FP16 (dense) | 148 TFLOPS | **0.9500** | | Tensor Core BF16 (dense) | 148 TFLOPS | **0.9500** | | Tensor Core FP8 (dense) | 296 TFLOPS | **0.9500** | | Tensor Core TF32 (dense) | 74 TFLOPS | 0.75 *(placeholder)* | ## Why calibration ratios differ from `h200.yaml` The H20 is an export-compliance SKU: NVIDIA **policy-caps** its advertised compute throughput to roughly 15% of H100, while the underlying silicon (and HBM3e stack) is essentially identical to H100/H200. Consequently: - **Tensor Core ~0.95** (vs. H200 ~0.75): the compute cap is the binding limit, so real GEMMs saturate the throttled peak almost entirely. - **HBM ~0.81** (vs. H200 ~0.85): the HBM3e stack is physically the same as H200's; the small delta reflects the lower advertised peak, not a memory-subsystem difference.
## Summary Promote existing code-present operators into the release-facing manifest surface for the bounded tile-ai#1660 batch, excluding Linear Attention and Mamba2/SSD work. This PR adds manifest coverage for: - RoPE variants, including `RopeNeoxPositionIdsOp` - `DropoutOp` - narrow FP8 helper surfaces: `FP8QuantOp`, `FP8LightingIndexerOp` - `TopkSelectorOp` - standalone `GroupedGemmOp` - Engram fwd/bwd/decode - `FFTC2COp` - `MHCPreOp` / `MHCPostOp` It also adds matching roofline formulas for these promoted entries. ## Scope notes This PR intentionally keeps the current runtime surface as the source of truth and avoids broad API redesigns. In particular: - Linear Attention and Mamba2/SSD promotion are left out of this batch. - New low-precision compute-core contracts such as FP8 GEMM, block-scaled GEMM, low-bit GEMM, sparse GEMM, and special swizzled layouts are out of scope. - `_infer_output_shapes` coverage and manifest-driven benchmark migration are kept as follow-up work. Existing validator warnings in those areas are advisory and do not block `validate_manifest.py --strict`. - Existing benchmark files are referenced with `bench_manifest_driven: false` to represent the current coverage honestly. ## Validation Ran in the TileOpsGov CI Docker image: ```bash docker run --rm --entrypoint /bin/bash --gpus all --ipc host \ -e PYTHONDONTWRITEBYTECODE=1 \ -e PYTHONPYCACHEPREFIX=/tmp/pycache \ -e RUFF_CACHE_DIR=/tmp/ruff-cache \ -v /home/ga/TileOPs-issue1650-wave0:/workspace \ -w /workspace \ ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10 \ -lc 'python scripts/validate_manifest.py --strict' ``` ```bash docker run --rm --entrypoint /bin/bash --gpus all --ipc host \ -e PYTHONDONTWRITEBYTECODE=1 \ -e PYTHONPYCACHEPREFIX=/tmp/pycache \ -e RUFF_CACHE_DIR=/tmp/ruff-cache \ -v /home/ga/TileOPs-issue1650-wave0:/workspace \ -w /workspace \ ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10 \ -lc 'python -m ruff check tileops/perf/formulas.py' ``` Both passed. `validate_manifest.py --strict` reports existing advisory warnings for shape-infer rollout and non-manifest-driven benches, but no hard errors. Closes tile-ai#1660. Part of tile-ai#1650.
Closes tile-ai#1654 ## Summary - Add an `AvgPool2dSpatialKernel` fast path for the common NCHW case with `ceil_mode=False`, `count_include_pad=True`, and no `divisor_override`. - Rewrite the general `AvgPool2dKernel` fallback to use a flat `N*C*H_out*W_out` launch grid and remove the `block_c` tuning dimension. - Extend `AvgPool2dFixture` coverage for `ceil_mode`, `count_include_pad`, and `divisor_override` combinations. ## Test plan - [x] Cleared TileLang caches before the final benchmark/profiling run: ```bash rm -rf /home/lyc/.tilelang \ /home/lyc/Project/TileOps-workspace2/.tmp/tilelang_clear_sync_cache \ /home/lyc/Project/TileOps-workspace2/.tmp/tilelang_ifelse_cache ``` - [x] Functional tests: ```bash CUDA_VISIBLE_DEVICES=1 conda run -n tileops-dev \ python -m pytest tests/ops/test_pool.py -k avg_pool2d -vvs ``` Result: ```text 23 passed, 28 deselected, 14 warnings in 57.08s ``` - [x] Benchmark: ```bash CUDA_VISIBLE_DEVICES=1 conda run -n tileops-dev \ python -m pytest benchmarks/ops/bench_pool.py::test_avg_pool2d_bench -vvs ``` Result: ```text 3 passed, 15 warnings in 34.03s ``` - [x] Whitespace check: ```bash git diff --check ``` ## Benchmark Final benchmark on NVIDIA H200: | workload | kernel path | before TileOps ms | after TileOps ms | torch-ref ms | after vs before | after vs torch | | --- | --- | ---: | ---: | ---: | ---: | ---: | | `vision-3x3-s2` | `AvgPool2dSpatialKernel` | `0.0505` | `0.0046` | `0.0094` | `10.98x faster` | `2.04x faster` | | `vision-5x5-s2` | `AvgPool2dSpatialKernel` | `0.0372` | `0.0045` | `0.0102` | `8.27x faster` | `2.27x faster` | | `ceil-divisor-bf16` | `AvgPool2dKernel` flat general | `0.0149` | `0.0038` | `0.0080` | `3.92x faster` | `2.11x faster` | Autotune best configs from the final run: | workload | kernel | best config | | --- | --- | --- | | `vision-3x3-s2` | `AvgPool2dSpatialKernel` | `{"block_m": 256, "threads": 256}` | | `vision-5x5-s2` | `AvgPool2dSpatialKernel` | `{"block_m": 512, "threads": 256}` | | `ceil-divisor-bf16` | `AvgPool2dKernel` | `{"block_m": 256, "threads": 256}` | Notes on the baseline values: - The `vision-3x3-s2` and `vision-5x5-s2` "before" values are from the original TileOps kernel measured during the avg_pool2d performance investigation. - The `ceil-divisor-bf16` "before" value is the local pre-flat-general fallback re-run used for the review comparison. ### Profiling summary The benchmark improvement is consistent with fixed-config `nsys` and `ncu` data. For `vision-3x3-s2`, `ncu` shows that TileOps executes much less work than torch: | metric | TileOps | torch | | --- | ---: | ---: | | Duration | `5.12 us` | `10.05 us` | | Waves/SM | `1.48` | `1.48` | | Achieved Occupancy | `73.61%` | `79.44%` | | SM Busy | `44.53%` | `67.92%` | | Memory Throughput | `628.30 GB/s` | `321.45 GB/s` | | Issued Instructions | `1.09M` | `4.24M` | Torch has higher SM busy, but it issues about `3.9x` more instructions. The TileOps fast path avoids the general pooling divisor/control path and reaches about `2x` higher effective memory throughput. For `vision-5x5-s2`, fixed-config `nsys` reports: | metric | TileOps | torch | | --- | ---: | ---: | | nsys median | `3.808 us` | `9.664 us` | | nsys avg | `3.822 us` | `9.672 us` | | Grid / Block | `392 x 256` | `196 x 1024` | | Registers/thread | `29` | `32` | For `ceil-divisor-bf16`, the new flat general fallback is also faster than torch: | metric | TileOps | torch | | --- | ---: | ---: | | nsys median | `3.456 us` | `7.488 us` | | ncu Duration | `4.48 us` | `8.96 us` | | Grid / Block | `914 x 256` | `229 x 1024` | | Waves/SM | `0.87` | `0.87` | | Achieved Occupancy | `68.44%` | `73.10%` | | SM Busy | `43.02%` | `65.50%` | | Eligible Warps/Scheduler | `1.42` | `3.65` | | Memory Throughput | `404.46 GB/s` | `203.60 GB/s` | | Issued Instructions | `875K` | `3.28M` | Torch again keeps the SMs busier, but it issues about `3.75x` more instructions. The flat general kernel removes the old `block_c` split, launches over `N*C*H_out*W_out`, and completes the same workload with a shorter instruction path and higher effective memory throughput.
…1596) Closes tile-ai#1595. ## Summary Adds `GatedDeltaNetPrefillFwdOp`, an optimized zero-state inference prefill operator for Gated DeltaNet. The op returns `(o, final_state)` without training-only intermediates and provides a TileOps-owned serving path for long-context prefill. ## Implementation - Adds TileLang kernels for the BTHD serving path: chunk cumsum, blocked/Neumann-style A producer, CP initial-state correction, and fused replay/output. - Adds shape-aware dispatch for the primary serving contract: `layout="bthd"`, `B=1`, `fp16`, `DK=DV=128`, `chunk=64`, with an environment kill switch. - Keeps API compatibility for BTHD plus head-major BHTD/BHSD aliases. - Adds op registration, tests, benchmark coverage, manifest workload, and roofline formula support. - Fixes the manifest roofline path to decode BTHD workload shapes as `[B, T, H, D]` instead of `[B, H, T, D]`. ## Third-party Attribution This PR also adds `THIRD_PARTY_NOTICES.md` and clarifies source-file headers for the GDN prefill path. Qwen FlashQLA is the MIT-licensed schedule-level reference for the h-state / corrected-segment-start CP-split scheduling. The local utility file is adapted from FlashQLA's utility source and keeps its upstream copyright notice. TileOps-specific work includes the public op wrapper, BTHD/head-major API integration, shape-aware dispatch, blocked/Neumann-style A producer integration, local replay/output implementation, TileLang compatibility changes, manifest/roofline support, tests, benchmarks, and verification. ## Shape Contract Primary BTHD path: - `q, k`: `[B, T, H, DK]` - `v, o`: `[B, T, H, DV]` - `g, beta`: `[B, T, H]` - `initial_state, final_state`: `[B, H, DK, DV]` Head-major BHTD/BHSD layouts remain supported through the public op API, but the optimized long-context serving dispatch targets BTHD. ## Performance Headline H200/GPU4, BTHD, fp16, `B=1`, `DK=DV=128`, `chunk=64`, TileOps benchmark timer/CUPTI path. FlashQLA is reported as a public TL0.1.8 anchor and should be read as a public-environment comparison, not a same-lowering attribution experiment. Qwen FlashQLA is credited as the reference for the CP-split schedule family. | Shape | TileOps | FLA | FlashQLA | TileOps / FLA | TileOps / FlashQLA | | --- | ---: | ---: | ---: | ---: | ---: | | 32K / H16 | 0.369577 ms | 2.364970 ms | 0.583178 ms | 6.399x | 1.578x | | 64K / H16 | 0.690859 ms | 5.265457 ms | 1.331969 ms | 7.622x | 1.928x | | 128K / H16 | 1.227298 ms | 10.899666 ms | 2.623757 ms | 8.881x | 2.138x | | 128K / H32 | 2.384966 ms | 16.687767 ms | 5.395038 ms | 6.997x | 2.262x | ## Verification Host checks: - `python -m pytest tests/perf/test_gated_deltanet_prefill_roofline.py -q`: `2 passed` - `python -m pytest tests/test_validate_manifest.py -q`: `223 passed, 3 warnings` - `pre-commit run --files tileops/perf/formulas.py tests/perf/test_gated_deltanet_prefill_roofline.py`: passed - `PYTHONPYCACHEPREFIX=<tmp> python -m py_compile tileops/perf/formulas.py tests/perf/test_gated_deltanet_prefill_roofline.py`: passed - Manifest roofline smoke for `GatedDeltaNetPrefillFwdOp`: all listed BTHD workloads now report nonzero FLOPs and bytes. TileOpsGov CI docker image `ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10` (`tilelang 0.1.11+cu129.git65dbc983`, `torch 2.10.0+cu129`): - `python -m pytest tests/perf/test_gated_deltanet_prefill_roofline.py -q`: `2 passed` - Manifest roofline smoke for `GatedDeltaNetPrefillFwdOp`: all listed BTHD workloads report nonzero FLOPs and bytes. - `python -m pytest tests/test_validate_manifest.py -q`: `223 passed, 15 warnings` - GPU smoke on H200: `tests/ops/test_gated_deltanet_prefill.py::test_gated_deltanet_prefill_bthd_layout_matches_bhtd` and `::test_gated_deltanet_prefill_bthd_blocksolve_path_matches_bhtd`: `2 passed, 14 warnings`
… API (tile-ai#1661) ## Add payload support and implicit thread blocks to trace API ### Summary This PR adds two enhancements to the trace API: 1. **Payload support for duration ranges** - Brings `trace.range()` in line with `trace.record()` for payload capability. The `AnnoToken` now carries payload through both `RANGE_BEGIN` and `RANGE_END` markers. 2. **Implicit thread blocks support** - Enables tracing in kernels using simple `T.Kernel(..., threads=N)` syntax without explicit threadIdx.x binding. Uses a new `__tl_thread_idx_x()` CUDA helper when explicit binding is not present. ### Changes #### 1. Payload Support **API Changes**: - **`trace.range(name, lane, payload=None)`** - Added optional `payload` parameter - **`trace.range_start(name, lane, payload=None)`** - Added optional `payload` parameter - Payload is recorded in both `RANGE_BEGIN` and `RANGE_END` events - `AnnoToken` now stores payload value for the entire range scope **Implementation**: - Updated `AnnoToken` to store payload value - Modified `open_range()` and `close_range()` to accept and emit payload - Updated `RangeScope` to pass payload through context manager - Simplified payload handling by letting `_emit()` handle None conversion **Usage Example**: ```python # Tag iterations with their index for i in range(N): with trace.range("iteration", payload=i): work(i) # Payload appears in timeline UI (displayed separately and in hover tooltip) ``` #### 2. Implicit Thread Blocks Support **Problem**: Kernels using simple `T.Kernel(1, threads=16)` syntax have no explicit threadIdx.x binding, causing trace lowering to fail with `KeyError: 'threadIdx.x'`. **Solution**: - Added `__tl_thread_idx_x()` CUDA device helper in `tileops/trace/device.py` - Modified `tileops/trace/passes.py` to detect missing threadIdx.x binding - Falls back to `T.call_extern("int32", "__tl_thread_idx_x")` when binding not found - Emits a warning to stderr when fallback is used **Before** (required explicit binding): ```python with T.Kernel(1, threads=(1, 16)): ty = T.thread_binding(0, 1, thread="threadIdx.y") tx = T.thread_binding(0, 16, thread="threadIdx.x") # Required! with trace.range("work"): work(tx) ``` **After** (simple syntax works): ```python with T.Kernel(1, threads=16): tx = T.get_thread_binding() # Works now! with trace.range("work"): work(tx) ``` ### Backward Compatibility ✅ Fully backward compatible: - `payload` parameter is optional and defaults to `None` (recorded as 0) - Explicit threadIdx.x binding still works as before - Existing code continues to work without changes ```python # Both still work with trace.range("compute"): # No payload work() with T.Kernel(1, threads=(1, 16)): # Explicit binding tx = T.thread_binding(0, 16, thread="threadIdx.x") with trace.range("work"): work(tx) ``` ### Testing **Test coverage** (`tests/trace/test_payload.py`): - ✅ API signature verification (compilation tests) - ✅ Constant payload with decode (payload=42) - ✅ Dynamic payload with runtime PrimExpr (loop index → [0, 1, 2, 3]) - ✅ range_start/range_end with payload decode - ✅ Backward compatibility without payload - ✅ Implicit thread blocks with trace enabled/disabled - ✅ Process-level trace state preservation (no test pollution) - ✅ All tests marked with `pytest.mark.full` tier **Test strategy**: - GPU end-to-end verification with trace.decode() - Tests verify payload is written to slots and can be decoded - Tests verify implicit thread blocks fallback compiles and runs - Precise assertions (e.g., "exactly one slice with payload=42") - Coverage of both constant and runtime expression payloads ### Known Limitations⚠️ **Unverified behavior**: Payload behavior inside `T.Pipelined()` loops has not been validated. Compiler optimizations (such as loop-invariant code motion or loop unrolling) may affect whether payload values from individual iterations are captured. **Recommendation**: Use payload for: - ✅ Non-pipelined loops - ✅ Conditional tracing - ✅ High-level event tagging - ✅ Simple kernels with implicit thread blocks **Use with caution**: -⚠️ `T.Pipelined()` loop iterations (behavior not yet verified) For precise pipeline profiling, consider using NVIDIA NSight Compute. ### Future Work Potential improvements (out of scope for this PR): - Add tests that decode trace slots and assert payload values - Validate payload behavior in pipelined loops - Compiler pass to preserve trace markers through optimizations - Pipeline-aware trace API - Eliminate warning for implicit thread blocks (make it silent) ### Checklist - [x] API changes are backward compatible - [x] Unit tests added (compilation and API verification) - [x] Tests preserve trace state (no pollution) - [x] Implicit thread blocks support working - [x] Code follows project style - [x] Known limitations documented - [ ] Tests that decode and assert payload values (follow-up) - [ ] Documentation updated (follow-up) --- **Note**: This PR combines two related trace API enhancements. Both features improve the usability of the trace system for common kernel patterns. Feedback welcome on: 1. API design and payload propagation through token 2. Implicit thread blocks fallback strategy 3. Test coverage approach
superAngGao
force-pushed
the
issue-1597-gdn-prefill-bench
branch
from
ffcf855 to
95a312b
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
6 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stacked on tile-ai#1596. Addresses tile-ai#1597.
Summary
Validation
Notes