[Refactor][Attention] Establish dense GQA Op boundary by superAngGao · Pull Request #1975 · tile-ai/TileOPs

superAngGao · 2026-08-24T09:09:10Z

Summary

Define the spec-only public boundary for GroupedQueryAttentionDenseFwdOp over dense BSHD tensors.
Standardize rectangular causal/window/RoPE semantics on bottom-right alignment.
Keep validation and optional-input normalization in the Op, then resolve implementations through the existing Op.get_or_build_kernel boundary. A future BUILTIN PR will provide the shape/dtype key and concrete kernel construction.

The ABI covers rectangular prefill and contiguous decode, FP16/BF16, scaled FP8, causal/window/softcap, and NeoX/interleaved RoPE. This PR adds no BUILTIN kernel, benchmark, compatibility adapter, design-document change, Op base change, runtime test, or torch.compile promise; kernel families and their public execution tests will migrate separately under #1916.

Validation

Manifest validator suite: 154 passed
Ruff, Python compilation, and git diff --check passed
Independent review found no remaining high- or medium-severity issue

zhen8838 · 2026-08-25T03:00:54Z

+    def default_kernel_map(self) -> Dict[str, Kernel]:
+        return {}
+
+    def _infer_output_shapes(


Why do these two functions use only one parameter but pass so many unused ones?

The full signatures are intentional because they mirror the manifest input slots. _infer_output_shapes is required by the Op abstract contract and manifest parity checks; Dense GQA output shape depends only on q_shape. _validate_dtypes uses the tensor inputs to enforce Q/K/V, scale, and RoPE dtype relations.

lcy-seso · 2026-08-25T03:14:18Z

+        v_scale: Optional[torch.Tensor],
+        rope_cos: Optional[torch.Tensor],
+        rope_sin: Optional[torch.Tensor],
+    ) -> torch.Tensor: ...


A protocol feels too heavy for a single-use interface — worth it only if the interface is meant to be shared across many APIs.

Agreed. I removed _DenseFwdCallable, the Protocol, and the cast. The Op now uses the existing get_or_build_kernel boundary directly.

lcy-seso

Reviewed the Dense GQA boundary. One blocking item, the rest are placement and duplication.

Blocking

Query-to-key alignment is undefined for S_q < S_kv (inline on gqa.py:291). It decides both the mask and the RoPE positions, and every target will guess differently until it is written down.

Also

forward carries 21 argument checks inline while two sibling ops in the same file use _validate_forward_inputs.
Three checks duplicate __init__; resolved_sm_scale is computed and never used.
No test. grep finds the class only in ops/__init__.py, attention/__init__.py, and the manifest; source.test names tests/ops/attention/test_gqa.py, which has no case for it. status: spec-only covers the missing kernel, but the 95 lines of validation ship constructible and callable, unexercised.
The class docstring never states what the op computes.

lcy-seso · 2026-08-25T03:37:45Z

+        rope_sin: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Run Dense BSHD GQA through the selected callable."""
+        batch, seq_len_q, heads, dim = q.shape


forward runs 26 branches; 21 of them are raise ValueError on the arguments.

Two sibling ops in this file already extract that: GroupedQueryAttentionPrefillVarlenFwdOp._validate_forward_inputs (L891) and GroupedQueryAttentionPrefillPagedWithKVCacheFwdOp._validate_forward_inputs (L1178).

Here only dtype is extracted (_validate_dtypes); shape, device, and combination checks stay inline. Half-extracted reads worse than not extracting.

Suggested shape: _validate_forward_inputs + _resolve_optional_inputs returning the 8-tuple, leaving forward as unpack shapes -> validate -> resolve -> self._get_callable(inputs)(*inputs).

Done. Runtime checks now live in _validate_forward_inputs; contiguous conversion and fixed manifest-slot construction live in _canonicalize_inputs. forward is now validate -> canonicalize -> get kernel -> invoke.

lcy-seso · 2026-08-25T03:37:45Z

+    ) -> torch.Tensor:
+        """Run Dense BSHD GQA through the selected callable."""
+        batch, seq_len_q, heads, dim = q.shape
+        batch_kv, seq_len_kv, heads_kv, dim_kv = k.shape


No ndim check on q / k / v.

A 3-D q fails here as ValueError: not enough values to unpack, which does not name the argument.

Both sibling _validate_forward_inputs methods check ndim explicitly.

Done. q, k, and v now receive explicit rank-4 BSHD checks before shape unpacking, with the argument name in each error.

lcy-seso · 2026-08-25T03:37:45Z

+        _validate_gqa_dims(heads, heads_kv, dim)
+
+        if self.is_causal and seq_len_q > seq_len_kv:
+            raise ValueError("causal dense attention requires seq_len_q <= seq_len_kv")


Blocking: seq_len_q < seq_len_kv is admitted, but nothing defines where query i sits on the key axis.

The existing square kernel masks with q_idx >= k_idx, i.e. top-left (kernels/attention/gqa_fwd.py:102).

Decode needs bottom-right (i + S_kv - S_q >= j), otherwise it attends to key 0 only.

The same offset decides which RoPE position q_i gets, and there is no position/offset input.

Neither the docstring nor shape_rules states it, so every target will pick its own. Please fix the alignment here.

Fixed. The class contract now defines bottom-right alignment as p_i = i + S_kv - S_q; causal/window masking and fused-RoPE query positions are defined from p_i. The manifest records the same constraints.

lcy-seso · 2026-08-25T03:37:45Z

+        if self.pos_encoding_mode == "rope" and seq_len_q > seq_len_kv:
+            raise ValueError("fused RoPE requires seq_len_q <= seq_len_kv")
+
+        resolved_sm_scale = _attention_scale(dim, self.sm_scale)


resolved_sm_scale is computed and discarded.

Its only use is the isfinite check on the next line, which __init__:183 already performs.

_manifest_params() hands the builder self.sm_scale, possibly None, so every target re-derives 1/sqrt(D).

softcap is handled the other way: normalized once in __init__ via _score_softcap. Pick one convention for both.

The discarded local has been removed. sm_scale=None remains constructor state because its default depends on the current head dimension; softcap is shape-independent and is normalized at construction. The future BUILTIN build closure (or a target builder using TensorSpec) resolves the dimension-dependent scale.

lcy-seso · 2026-08-25T03:37:45Z

+        if not math.isfinite(resolved_sm_scale):
+            raise ValueError("sm_scale must be finite")
+        if self.pos_encoding_mode not in ("none", "rope"):
+            raise ValueError("pos_encoding_mode must be 'none' or 'rope'")


Duplicate of __init__:179. pos_encoding_mode is not mutated between construction and call.

Done. pos_encoding_mode is validated only at construction now.

lcy-seso · 2026-08-25T03:37:45Z

+            if tuple(scale.shape) != (batch, heads_kv):
+                raise ValueError(f"{name} must have shape {(batch, heads_kv)}")
+            normalized_scales.append(scale.contiguous())
+        resolved_scales = tuple(normalized_scales) if normalized_scales else (None, None, None)


normalized_scales is built by append, so its length depends on the all-or-none check at L314 holding.

If that check moves, this splats a 1- or 2-tuple into inputs at L359 and the failure surfaces at the callable, not at validation.

tuple(s.contiguous() if s is not None else None for s in scales) removes the coupling.

Done. _canonicalize_inputs now constructs all eight manifest slots with one tuple comprehension, preserving None independently of the all-or-none validation.

lcy-seso · 2026-08-25T03:37:45Z

+                raise ValueError(f"{name} must have dtype {output_dtype}")
+
+    def eval_roofline(self) -> tuple[int, int]:
+        raise NotImplementedError("Dense GQA has no in-tree implementation yet")


This override can go: Op.eval_roofline (op_base.py:191) already raises NotImplementedError, with a message pointing at roofline.md 4.4.6.

I retained this minimal override after checking the class contract: Op.eval_roofline is abstract, and roofline codegen intentionally skips status: spec-only. Removing the override makes the public class non-instantiable. The override can disappear once an implemented entry supplies generated roofline code.

lcy-seso · 2026-08-25T03:37:45Z

+    selection and caching of concrete kernels belongs here, not in the Op.
+    """
+
+    def __call__(


_DenseFwdCallable has one use site, and it is the target of a cast, which does not check at runtime.

As written it buys what a Callable[...] alias would. Worth a Protocol only if a target-side implementation gets checked against it.

Agreed and removed. There is no dedicated callable type or second callable-owned cache now.

lcy-seso · 2026-08-25T03:37:45Z

+    FP16/BF16 and scaled FP8 inputs, and optional caller-owned RoPE tables.
+    This target-neutral layer caches a target-owned :class:`_DenseFwdCallable`
+    under the standard input-signature rules; concrete kernel dispatch remains
+    inside that callable. The in-tree callable is deferred to a follow-up.


This class defines the public contract for dense GQA but never states what it computes. Please write the linear algebra into the class docstring.

The docs site renders LaTeX now; ops/norm/layer_norm.py:22 is the in-repo precedent for the format ($$ ... $$, escaped backslashes). The equation should pin down what prose leaves open:

which KV group head h reads

where query i sits on the key axis, i.e. the offset the causal and window masks compare against (see the alignment comment on L291)

the score scale and its default

how softcap enters

where the FP8 q_scale / k_scale / v_scale factors apply

what accumulates in FP32, and where the cast to dtype happens

Done. The class docstring now specifies KV-group head mapping, bottom-right coordinates, mask/window semantics, scale default, softcap placement, FP8 scale placement, FP32 accumulation, and output casting.

lcy-seso · 2026-08-25T03:37:45Z

+      - {q: bfloat16, k: bfloat16, v: bfloat16, o: bfloat16}
+      - {q: float8_e4m3fn, k: float8_e4m3fn, v: float8_e4m3fn, o: float16}
+      - {q: float8_e4m3fn, k: float8_e4m3fn, v: float8_e4m3fn, o: bfloat16}
+    shape_rules:


forward enforces two constraints missing from shape_rules:

is_causal requires S_q <= S_kv (gqa.py:291)

pos_encoding_mode == 'rope' requires S_q <= S_kv (gqa.py:293)

The manifest is the spec, so both belong here, plus the rule stating query-to-key alignment once it is decided.

Done. The manifest now includes the causal and fused-RoPE S_q <= S_kv rules and documents bottom-right query-to-key alignment.

superAngGao · 2026-08-25T07:15:03Z

Updated at 192ec155:

defined bottom-right alignment (p_i = i + S_kv - S_q) in the Op contract and manifest;
moved runtime checks into _validate_forward_inputs and constructor-only invariants into __init__;
added explicit BSHD rank errors;
replaced append-based optional handling with _canonicalize_inputs, which returns contiguous tensors in all eight manifest slots while preserving None;
removed the single-use Protocol, cast, and second-cache framing;
documented the Dense GQA linear algebra, FP8 scales, softcap, accumulation, and output-cast semantics;
added the missing causal/RoPE manifest constraints.

I retained the minimal eval_roofline override because the base method is abstract and codegen skips status: spec-only; without it the public Op cannot be instantiated.

The PR has also been narrowed: it no longer changes Op base, design docs, backend-seam tests, or existing GQA tests, and it adds no BUILTIN kernel. Runtime correctness tests will accompany the first concrete kernel migration. Manifest validation, Ruff, Python compilation, and git diff --check pass.

## Problems - `GroupedQueryAttentionDenseFwdOp` (tile-ai/TileOPs#1975) has docstring formulas but no `:::` entry on `api/attention.md`, so the op and its math never render. - A scan of TileOPs top-level exports and manifest entries against the API pages found 24 public ops in the same state. - `api/quantization.md` and `api/topk.md` are coming-soon placeholders while both ops are `implemented` in the manifest and exported at the top level. - The RoPE, pooling, convolution, and dropout families have no API page at all. ## Changes - `attention.md`: add `GroupedQueryAttentionDenseFwdOp` under Grouped-query attention, and an Attention indexing section for `FP8LightningIndexerFwdOp`. - `quantization.md`, `topk.md`: replace the placeholders with `FP8QuantFwdOp` and `TopkSelectorFwdOp`. - New pages: `rope.md` (6 ops), `pool.md` (12 ops), `convolution.md` (3 ops), `dropout.md` (1 op), each following the existing page pattern. - `mkdocs.yml`: four new nav entries plus `nav_translations` labels; `api/index.md`: table rows for the new pages, stale-interface caveats dropped. - Verified: `mkdocs build` renders all six pages; the Dense GQA $$ formulas arrive as arithmatex blocks. 6 new griffe warnings are missing type annotations in TileOPs docstrings (`fp8_quant.py`, `topk_selector.py`, `fp8_lightning_indexer.py`), same class as the 16 pre-existing ones — the fix belongs upstream. ## Left out on purpose - MoE and Engram ops: the MoE family is mid-refactor with spec-only staged ops, and neither family is exported at the `tileops.ops` top level. - `MeanPoolingForwardOp`: an `UnmanifestedOp`.

## Summary - document that Dense GQA derives `B`, `Sq`, `Skv`, `H`, `Hkv`, `D`, and input dtype from every `forward` call - separate construction-time attention semantics from call-time tensor metadata - define the full computation: FP8 dequantization, optional RoPE, scaled QK, softcap, causal/window masking, softmax, and PV reduction - map every constructor parameter to its effect in the equations - clarify that one Op instance may resolve and cache implementations for multiple input shapes ## Scope This draft does **not** change the constructor, `forward` signature, manifest, validation, dispatch, or kernels. It only clarifies the input-inferred contract established in #1975, following the documentation structure used by `GemmFwdOp`.

## Summary - Define the spec-only public boundary for `GroupedQueryAttentionDenseFwdOp` over dense BSHD tensors. - Standardize rectangular causal/window/RoPE semantics on bottom-right alignment. - Keep validation and optional-input normalization in the Op, then resolve implementations through the existing `Op.get_or_build_kernel` boundary. A future BUILTIN PR will provide the shape/dtype key and concrete kernel construction.

github-actions Bot added the refactor Code restructuring without behavior change label Aug 24, 2026

superAngGao force-pushed the refactor/attention/gqa-dense-op-boundary branch 6 times, most recently from 1e21de7 to d8dabe8 Compare August 24, 2026 11:09

superAngGao marked this pull request as ready for review August 24, 2026 11:20

superAngGao requested a review from a team August 24, 2026 11:20

superAngGao force-pushed the refactor/attention/gqa-dense-op-boundary branch from d8dabe8 to 2e32818 Compare August 24, 2026 11:27

[Refactor][Attention] Establish dense GQA Op boundary

75ebfe3

superAngGao force-pushed the refactor/attention/gqa-dense-op-boundary branch from 2e32818 to 75ebfe3 Compare August 24, 2026 11:35

zhen8838 previously approved these changes Aug 25, 2026

View reviewed changes

lcy-seso reviewed Aug 25, 2026

View reviewed changes

superAngGao added 3 commits August 25, 2026 15:09

[Refactor][Attention] Address dense GQA boundary review

1cc7e04

[Refactor][Backend] Allow resolved external builder params

cff870d

[Refactor][Attention] Resolve dense GQA builder inputs

883c2b5

superAngGao dismissed zhen8838’s stale review via 883c2b5 August 25, 2026 07:14

superAngGao added 7 commits August 25, 2026 15:17

[Doc][Ops] Format kernel cache reference table

a62a6c8

[Refactor][Attention] Narrow dense GQA boundary tests

a2ff6af

[Refactor][Attention] Keep dense boundary independent of Op base

82ba30b

[Test][Attention] Exercise dense validation without a fake kernel

1492a7b

[Test][Attention] Defer dense GQA runtime coverage

bcaeef0

[Refactor][Attention] Clarify dense input canonicalization

f2a994e

[Style][Attention] Format dense input canonicalization

192ec15

lcy-seso approved these changes Aug 25, 2026

View reviewed changes

lcy-seso merged commit e0adc15 into tile-ai:main Aug 25, 2026
15 checks passed

lcy-seso mentioned this pull request Aug 26, 2026

[Docs] Document every implemented op the API reference was missing tile-ai/TileOPs.github.io#31

Merged

superAngGao mentioned this pull request Aug 27, 2026

[Doc][Attention] Clarify dense GQA input-inferred contract #1994

Merged

Conversation

superAngGao commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lcy-seso Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

lcy-seso left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

superAngGao commented Aug 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

superAngGao commented Aug 24, 2026 •

edited

Loading

lcy-seso Aug 25, 2026 •

edited

Loading

superAngGao commented Aug 25, 2026 •

edited

Loading