[Perf][Linear-Attn] Add optimized Gated DeltaNet prefill op by superAngGao · Pull Request #1596 · tile-ai/TileOPs

superAngGao · 2026-06-17T09:11:41Z

Closes #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

Gabbering

goose skimmed 58a4ea4 — nothing to honk about.

gemini-code-assist

Code Review

This pull request implements the Gated DeltaNet inference prefill operator (GatedDeltaNetPrefillFwdOp) and its corresponding kernel, along with associated tests, benchmarks, and performance formulas. Feedback on the changes highlights two critical issues: first, input tensors should be made contiguous in the operator's forward pass to prevent memory access errors or crashes; second, the optimized scan path condition (use_scan_h) in the prefill kernel needs an additional divisibility check on the sequence length to avoid silent correctness bugs on non-divisible lengths.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Gabbering

goose skimmed ec23c79 — nothing to honk about.

Gabbering

goose skimmed b617d52 — nothing to honk about.

Gabbering

goose skimmed 2170e25 — nothing to honk about.

Gabbering

goose review — `17bfa8b9`

honk. Both inline thread comments were addressed by the master. The goose has re-read the conversation and confirmed the following:

Thread 1 (tileops/ops/gated_deltanet.py:303): The master added .contiguous() calls for all five input tensors before validation and kernel dispatch, and wrote a test (test_gated_deltanet_prefill_bthd_layout_matches_bhtd) that passes non-contiguous permuted inputs through the BTHD path. The goose's contiguity concern is resolved.
Thread 2 (tileops/kernels/gated_deltanet/gated_deltanet_prefill.py:852): The master defined group_chunks = 64 before the use_scan_h guard and added and seq_len % (chunk_size * group_chunks) == 0 to the condition. Non-divisible long sequences now fall back to the standard recurrence path. The goose's scan-path divisibility concern is resolved.

Both fixes are present in the delta 17bfa8b9 and the goose sees no new issues introduced by this commit.

SILENT

Ibuki-wind

Overall

The new manifest-driven benchmark does not record an independent baseline, so the benchmark contract is not satisfied yet.

Gabbering

goose skimmed d09c8f2 — nothing to honk about.

Ibuki-wind

Overall

The benchmark baseline blocker is fixed, but the public BTHD layout contract is not represented in the manifest and the approval-gate authoring hygiene still needs cleanup. Update the manifest shape rules for layout, make the PR body follow the template with final verification facts including pre-commit/benchmark/test-node-delta, and edit author replies to outcome-only one-liners.

## Summary - update the GatedDeltaNetPrefillFwdOp spec-only manifest contract to allow chunk_size=None - add the layout parameter used by the implementation path - add an H16/S2K manifest workload row for the target prefill shape family ## Notes - This is a manifest-only precursor for PR #1596. - It intentionally keeps status: spec-only and leaves source paths null; the implementation PR should flip status/source/roofline once the kernel code lands. ## Validation - python scripts/validate_manifest.py --check-op GatedDeltaNetPrefillFwdOp

## Summary - encode the GatedDeltaNetPrefillFwdOp BHTD/BHSD and BTHD layout shape contracts in the manifest - add fixed-rank shape metadata for the manifest tensors - fill the planned roofline/source metadata while keeping the entry spec-only ## Notes - This is a manifest-only prerequisite for PR #1596. - The implementation PR should only flip status and keep any remaining manifest edit within the status-flip carve-out. ## Validation - python scripts/validate_manifest.py

zhen8838 · 2026-06-22T03:50:37Z

For the performance comparison, please include FlashQLA as a baseline/reference as well: https://github.com/QwenLM/FlashQLA

It would be useful to compare on the same target hardware and representative GDN prefill shapes/dtypes so the remaining gap vs. the current public optimized implementation is explicit.

superAngGao · 2026-06-25T07:24:22Z

Thanks for the pointer. I added FlashQLA as an explicit baseline in the PR body, measured on the same H200/GPU4 setup with the TileOps benchmark/CUPTI timer and representative Qwen-style GDN prefill rows.

The refreshed numbers show the TileOps path is faster than the public FlashQLA reference on these rows. The main reason is not just a different block/thread schedule tweak. The important change is the overall inference prefill path: the serving path materializes only o + final_state, avoids the old global w/u/S/v_new intermediates, uses CP-split replay with corrected segment initial states, and combines that schedule family with the TileOps-owned A producer, local replay/output implementation, and shape-aware dispatch. Parallel schedule tuning still matters, but the larger win comes from this algorithmic decomposition of the inference prefill path.

I also kept FlashQLA attribution explicit in the PR body: Qwen FlashQLA is credited as the schedule-level reference for the h-state / corrected-segment-start CP-split part of long-context GDN prefill. The TileOps-specific work is the owned operator integration, BTHD/head-major API, blocked/Neumann-style A-producer integration, local replay/output implementation, dispatch/tuning, tests, benchmarks, and validation against the recorded references.

…shqla-specialized # Conflicts: # tileops/kernels/gated_deltanet/gated_deltanet_fwd.py

zhen8838

Approved current head 7fb5f35. The earlier requested FlashQLA comparison is now in the PR body with H200/Qwen-style rows, the manifest now represents the BTHD public layout contract with layout-specific shape rules, and current GitHub checks are green including gpu-smoke, validate-manifest, benchmark-contract, pre-commit, actionlint, and gitleaks.

Ibuki-wind

Overall

Please fix the manifest roofline shape handling and rewrite the PR body into a commit-log-ready summary before merge.

Cross-cutting concerns

The current PR body is too noisy to become the squash commit log: it mixes implementation history, temporary verification logs, artifact paths, long benchmark tables, and explanatory back-and-forth. Replace it with a concise final-state commit message: short summary, key implementation bullets, supported dispatch shape, final performance headline, and final verification commands/results only.

superAngGao · 2026-07-06T02:28:51Z

Thanks for the review. Fixed the manifest roofline blocker and rewrote the PR body into a final-state summary.

What changed:

gated_deltanet_prefill_fwd_roofline now decodes manifest q_shape / v_shape according to layout; BTHD workloads are treated as [B, T, H, D] and head-major BHTD/BHSD remains supported.
Added CUDA-free roofline tests for both BTHD manifest shapes and head-major compatibility.
Replaced the PR body with a concise squash-summary-style description: implementation, supported shape contract, performance headline, and final verification only.

Verified:

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.

superAngGao · 2026-07-06T07:39:13Z

Also re-verified with the TileOpsGov CI runner image from the new launcher docs:

Image: ghcr.io/tile-ai/tileops-runner:65dbc98-torch2.10
Runtime: torch 2.10.0+cu129, tilelang 0.1.11+cu129.git65dbc983

Docker results:

python -m pytest tests/perf/test_gated_deltanet_prefill_roofline.py -q: 2 passed
Manifest roofline smoke for GatedDeltaNetPrefillFwdOp: all BTHD workloads report nonzero FLOPs/bytes.
python -m pytest tests/test_validate_manifest.py -q: 223 passed, 15 warnings
H200 GPU smoke: BTHD layout compatibility + BTHD blocksolve path tests: 2 passed, 14 warnings

I also added these docker results to the PR body.

superAngGao · 2026-07-07T09:40:26Z

Thanks for raising this. I kept the original license headers because this path includes third-party MIT-licensed lineage, but I tightened the wording to avoid implying that TileOps directly wraps or copies the complete FlashQLA kernels.

More specifically:

From Qwen FlashQLA, this PR uses the CP-split scheduling idea for the h-state / corrected-segment-start part of long-context GDN prefill. That schedule-level reference should remain credited, so the affected files keep the original MIT notice with an explicit TileOps modification note.
utils.py is adapted from the FlashQLA utility source and keeps the upstream copyright notice present there; I no longer describe it as a direct Flash Linear Attention import.
TileOps-specific work in this PR includes the public GatedDeltaNetPrefillFwdOp wrapper, BTHD/head-major API integration, shape-aware dispatch, blocked/Neumann-style A-producer integration, local replay/output implementation, TileLang compatibility/env changes, manifest and roofline support, tests, benchmark integration, and verification.

To make this explicit in-repo, I added THIRD_PARTY_NOTICES.md and changed the affected file headers to point there instead of the ambiguous see LICENSE for details wording. I also updated the PR body with a short third-party attribution section using the same schedule-level wording.

Ibuki-wind

Clean on current head under a strict correctness/project-contract review.

The prior blockers I cared about are addressed: benchmark rows now record an independent FLA or torch-ref baseline, the manifest represents the BTHD public layout contract, and roofline decoding branches on layout so BTHD workloads produce nonzero FLOPs/bytes. I also rechecked the earlier correctness hazards: forward makes inputs contiguous before validation/dispatch, and the grouped scan path now requires divisibility by chunk_size * group_chunks before taking that path.

I treated the remaining third-party attribution threads as already answered by the clarified headers plus THIRD_PARTY_NOTICES.md; I do not see a code correctness blocker there. GitHub CI is green, including gpu-smoke, validate-manifest, benchmark-contract, and pre-commit. I did not run the GPU tests locally.