Skip to content

[Refactor][Attention] Consolidate GQA prefill ops and kernels - #1913

Draft
superAngGao wants to merge 8 commits into
tile-ai:mainfrom
superAngGao:refactor/gqa-prefill-fp8-goal
Draft

[Refactor][Attention] Consolidate GQA prefill ops and kernels#1913
superAngGao wants to merge 8 commits into
tile-ai:mainfrom
superAngGao:refactor/gqa-prefill-fp8-goal

Conversation

@superAngGao

@superAngGao superAngGao commented Aug 14, 2026

Copy link
Copy Markdown
Collaborator

Reference implementation only. Mergeable work is tracked in #1915 and will be split into topology-scoped PRs; this draft remains available for integrated review and validation. KV-cache mutation/append is a required final follow-up and is not implemented by this draft.

Summary

Consolidate the GQA prefill operator and kernel surface around three data topologies:

  • GroupedQueryAttentionPrefillDenseFwdOp for dense BSHD inputs
  • GroupedQueryAttentionPrefillVarlenFwdOp for packed THD inputs
  • GroupedQueryAttentionPrefillPagedWithKVCacheFwdOp for paged KV cache and optional append

Full/causal/sliding-window attention, softcap, NeoX/interleaved RoPE, FP16/BF16, FP8 cache storage, and native FP8 compute are selected inside those topology-level Ops instead of being exposed as separate public Ops. This does not add a new manifest schema.

The kernel layer is reorganized accordingly:

  • share mask, logical-position, RoPE, and FP8 prefill building blocks
  • keep Dense, Varlen, and Paged loaders as separate topology specializations
  • add native-FP8 Dense, Varlen, and Paged paths, including causal/window/softcap/tail coverage
  • consolidate Paged FP16/BF16, FP8-cache dequantization, append, and fused-RoPE dispatch
  • remove obsolete standalone sliding-window Op/benchmark/test surfaces and unused GQA kernels

Finally, all five GQA forward Ops now use the target seam: external targets receive the exact normalized manifest ABI, while NVIDIA architecture inspection and in-tree kernel selection remain inside builtin factories. External callable memoization is device-aware and bounded.

This broader refactor supersedes #1896, incorporates the applicable FP8 kernel work from the still-open #1873, and covers the contiguous prefill-with-KV-cache functionality proposed in #1879.

Validation

Targeted suites completed on H200:

  • backend seam, Op core, and manifest: 59 passed
  • GQA smoke suite: 108 passed, 15 deselected
  • multi-dtype and kernel-map coverage: 70 passed
  • final Dense/MHA wrapper regression: 44 passed, 8 deselected
  • ruff check passed on the changed Python surface before the final upstream-only manifest-label rebase; git diff --check passes after rebase

Native-CUPTI measurements on a clock-locked H200 (SM 1500 MHz) confirmed no performance regression across the covered Dense, Varlen, and Paged GQA prefill paths.

This is intentionally a broad draft so the three-topology ABI and the Op-to-kernel ownership boundary can be reviewed before compatibility cleanup is finalized.

@lcy-seso
lcy-seso force-pushed the refactor/gqa-prefill-fp8-goal branch from df8a94b to 4a595c8 Compare August 19, 2026 10:47
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

refactor Code restructuring without behavior change

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants