[jiapingW]

Implement registered rotary components and GQA/MHA projections with shared attention preparation and execution. Eager, SDPA, Flex and FlashAttention-2 implement the same prepare/attend interface; masks and rotary tables are reused across equivalent layers within each forward.

RoPE follows SGLang serving semantics: linear scaling divides positions, dynamic NTK uses a fixed factor × original-context length, and frequencies/tables remain FP32. Rotation computes in FP32 and casts once to the Q/K dtype. Static frequencies, including YaRN's ramp, are cached per device and invalidated on module moves/casts. Dynamic NTK has no position-dependent host sync.

This is an intentional numerical change from legacy DFlash's bf16-cast frequencies and eager per-operation rounding, and from EAGLE's rounded tables/history-dependent dynamic cache. Tests compare the legacy production rotation functions at 32k/65k positions and record the difference. SGLang table/native-reference comparisons use exact equality; the CUDA kernel comparison allows bf16 kernel rounding. Legacy EAGLE dynamic configurations remain explicitly unsupported.

Rotation pairing is explicit on the shared RoPE node. GQA takes a resolved RopeSpec and supports QK normalization before/after RoPE and per head/across the full projection, including weightless RMSNorm as used by SGLang Llama4. The legacy codec resolves old conventions explicitly; canonical schemas and golden fixtures require all structural fields. Input validation is shared by projection implementations.

Validation at fc2aff66af6538fc360773338d2ae7202f534ef6 on torch 2.14.1+cu130 and transformers 5.19.0:

• 222 configuration tests passed.
• CPU modeling: 136 passed, 15 environment-dependent tests skipped.
• GPU2: 55 attention/RoPE/GQA/reference tests passed, including short-query Flex and compiled-legacy rounding. The local SGLang is 0.5.23.dev1131+g4517a5ff1, not the CI pin; FlashAttention is not installed in this environment.
• Full-repository pre-commit checks passed.
• Fixed SGLang/FlashAttention dependency validation: PR Test workflow_dispatch, running at the same head; results pending.

Latest review fixes align backend names with training (eager, sdpa, flex_attention, fa), reject unsupported canonical USP, use explicit Flex kernels, validate padding across backends, share head-independent masks, and avoid tensor-ID Dynamo guards. Multi-axis interleaving, meta rotary execution, legacy field rejection and per-head inactive rows have regression tests. Canonical schema v1 remains provisional until the format cutover.

The branch includes #1116's 93c754bc initialization fixes through an ordinary merge. The RoPE numerics/schema wording is updated in #1076 (8fefec6b); QK-norm serving restrictions belong to #1162 and are implemented there (e7199330, two focused export tests passed).
The resolved-RoPE constructor, shared validation, boolean causal comparison and SDPA backend now belong to this PR. The current implementation and follow-ups are ordinary commits rather than implementation hidden in merge commits. Downstream consumers still need to be synchronized when the PR stack is refreshed.

Base: main (#1116 is merged). Next: #1130.

Latest main refresh: c63bd6a33c63145a1ac4540db5b3c8a35752a80f merges main 132d35d3, preserves the EAGLE regression case, and resolves the PR conflict. Post-merge CPU config/model/MTP validation: 411 passed, 17 skipped (428 total); full-repository pre-commit passed. The workflow_dispatch above tests the earlier fc2aff66 head, not this merge commit.

[maocheng23]

Inline notes below. Checked on CPU at the repo pins (torch 2.14.1, transformers 5.19.0): across the 9 changed test modules, two tests fail. They are test_draft_gqa.py (fp32 MHA, layout) and test_draft_rotary.py (device=). With the two fixes suggested inline, all 43 CPU tests pass. The GPU PR Test workflow only runs for PRs based on main, so CUDA precision and gradients haven't been exercised for this branch.

[maocheng23]

Thanks. The follow-ups address nearly everything from the first two rounds; I re-checked each item at e8d556c:

• Layout: a description (tensor, mask_mod or PaddingMask) plus active_queries.
• Backends: prepare/attend (Flex, SDPA, FA2), with the mask built once per forward and the LSE exposed.
• Errors: the eager backend is bool-only; FP32 RoPE tables are enforced; the apply_rotary length checks are in.
• Layout batches broadcast, and the registry test is isolated.
• RoPE semantics match SGLang: the interleaved convention lives on the RoPE node; linear scales positions; dynamic uses the fixed factor × original length; tables are cached per device.
• Stack hygiene: ordinary commits.

Verified on CPU at the repo pins:

• The changed tests pass: 184 passed, 12 skipped (CUDA/SGLang/flash).
• Attention and masks are still bit-equal to main's DFlash eager path when RoPE is factored out (identity tables, fp32 and bf16).
• The callable mask_mod path matches main's DFlash mask.
• Linear and dynamic RoPE match SGLang's formulas.
• Rotation arithmetic matches main's compiled EAGLE path, given the same tables.

Remaining:

1. The RoPE numerics change is now explicit and tested. I'm fine taking it as an intended deviation, since it follows SGLang serving and main's bf16-rounded frequencies are the bug. Let's record it in the design doc (#1076 §10 says steps 1–5 don't change numerics) and in the usage change list.
2. CI. Flex, FA2, the SGLang references and the compiled-legacy rounding test only run on CUDA, and the GPU PR Test workflow doesn't run for stacked PRs. Could you run it via workflow_dispatch on this branch at the repo pins (torch 2.14.1 / transformers 5.19.0) and note the head in the body?
3. Refresh onto #1116's current head. The merge is textually clean, but #1116's newest commit currently breaks on transformers 5.19 (_initialize_weights() gets is_remote_code), which will also fail test_meta_materialization_uses_generic_initialization here.
4. Serving boundary for the new qk-norm options. SGLang's DFlash and EAGLE drafts hard-code a per-head, weighted qk RMSNorm before RoPE. qk_norm_position="after_rope", qk_norm_scope="projection" and has_weight=False are accepted for every architecture, so such a draft trains but can't be served faithfully. The SGLang converter (or the architecture constraints) should reject them for those architectures.

[maocheng23]

Non-blocking nits:

• LayerInputs.attention_mask has no consumer left; could it go now that attention exists?
• RotaryEmbedding.forward still enters torch.autocast(device_type="meta"), which raises on meta tensors (RuntimeError: unsupported scalarType).
• The dense window path still allocates an int64 (q - k).abs() of shape [B, Q, K] (once per forward now). Comparing k > q - W and k < q + W keeps it boolean.
• No test covers per-head visibility for inactive-row zeroing.
• AttentionPreparation and PreparedAttention differ only in word order. Maybe AttentionPreparationScope?
• Dense masks from a PaddingMask are built with every head (create_mask(..., num_heads, ...)), which is H× the memory of [B, 1, Q, K]. Same for Flex create_block_mask.
• Canonical SCHEMA_VERSION is still 1 although v1, already on main, gains required fields (interleaved, qk_norm_*, has_weight). Bump it, or note that v1 is unfrozen until step 6.
• The legacy Qwen3/DFlash RoPE reader overwrites an interleaved key (codec/legacy/rules.py:57) instead of rejecting it, as the EAGLE reader does.
• DenseMask, FlashMask and AttentionResult use default dataclass equality over tensors; eq=False, like the other tensor holders.
• codec/legacy/context.py imports before its module docstring.

[maocheng23]

Approving. I re-checked at a5396f9:

• N9, N14, N15, N17 and N18 are fixed.
• The numerics decision is recorded here and in #1076.
• PR Test passed with every CUDA, FA and SGLang test running.

I pushed eight small follow-ups on top (d754c2ad..bf824a4d), each with a test:

• The backbone RoPE memo shares tables across modules only for the built-in rotary types. Before this, a custom subclass with per-module state silently got another layer's tables.
• FA2 casts FP32 Q/K/V to the autocast dtype when autocast is on, like SDPA and HF. Weighted qk-norm promotes them to FP32 under autocast, so FA2 used to reject them. It checks for CUDA first, so meta and CPU inputs still get the named error.
• FA2's LSE raises on backward, because flash_attn drops that gradient.
• The window rule compares distances. A window at or beyond the position dtype's limit no longer overflows or wraps; the layout resolves it once for dense, Flex and FA2. The shifted form came from my earlier nit, sorry.
• The docs say where canonical RoPE matches SGLang: the FP32 base cache for head sizes 64/128/256/512. Other sizes and multi-axis RoPE are rounded to the activation dtype in SGLang.

Follow-ups outside this PR:

• #1162's exporter has to honour RotarySpec.interleaved. It hardcodes NeoX and still reads MLASpec.rope_interleave.
• A custom RopeSpec without interleaved fails only at the first forward; better carried by #1130, which covers MLA too.
• The PRs above this one need main merged once it lands.