[DSv4.1] Add FROST prepared BF16 tail RoPE for prefill

[YangXu1990uiuc]

Stacked on #1114. Review the incremental diff.
This PR targets develop, so its Files changed tab currently includes the
parent change. Rebase after #1114 merges; review the linked increment separately.

Add TailRoPEForward and tail_rope for contiguous BF16 [T,D] and [T,H,D]
with D=128/512 on SM100. The kernel copies the unrotated prefix and rotates the
last 64 channels using caller-prepared FP32 cosine/sine tables. Explicit FP32
multiply/FMA ordering preserves the source BF16 rounding and signed zeros.

The prepared backend="frost" API supports caller-owned output, dynamic token
counts, explicit streams, zero workspace and CUDA Graph capture. Execute performs
no allocation, conversion, synchronization or compilation. CuTe DSL >=4.7.0 is
required. Lazy exports, documentation, API tests and a standalone benchmark are
included. This change is stacked on the group32 RoPE/QDQ change.

Two independent B200 runs measured the actual public API against the fastest
valid measured FlashInfer, compiled-Torch or source control per case and metric:

Prefill stage	Speedup across both runs and cache regimes
GPU event latency	1.04–1.69×
Graph replay wall latency	1.03–1.54×

Each run covers 64 cases: T=1/4/4096/16384; query, opposite-angle output,
index-query, KV and index-key geometries; TP1/TP8 local head counts; and source
base/compressed frequency tables. Activations are synthetic BF16. TP8 means
per-rank geometry, with no multi-rank execution or collectives. All prefill case
medians win on both metrics. Eight paired blocks and both cache regimes are
retained, including individual block win counts.

Decode has regressions: its GPU speedup spans 0.86–1.20×
and Graph wall spans 0.91–1.08×. The intended performance
use is prefill. Measurements include the full tensor copy, exclude table
preparation, and amortize three independent buffer pairs per replay. Graph wall
includes submission and completion waiting; it excludes eager API validation and
wrapper allocation. These are operator-stage measurements, with no complete
consumer, attention or full-model throughput claim.

Validation:

• All 32 new API tests and 29 parent QDQ tests pass on B200, including dynamic
  shapes, zero tokens, alias/alignment guards, explicit streams, changed-input
  Graph replay and a real tensor exceeding 2^31 elements.
• Four intentional faults are detected: device-to-host synchronization, ignored
  explicit stream, missing implicit-device binding and omitted Graph work.
• Both performance runs pass source, numerical, route, timing, GPU-isolation and
  thermal audits. Each records 3,456 numerical checks, 30,720 retained samples
  and 384 provider routes. Compiled PTX/CUBIN retention confirms all six public
  kernel specializations without changing execute.
• Frost matches the source bit for bit on seven input generations. Controls
  exact on both timed distributions remain competitors even when another stress
  generation fails; their witnesses and contract status remain recorded.
• CPU import/metadata checks pass, including declining old DSL before kernel
  import. The new kernel passes the naming guard; the pinned base has 17
  existing unrelated naming violations. The original GPU tests reused the
  pinned compiled FE binding; post-rebase host build validation is described
  below.

The API requires 16-byte-aligned, contiguous tensors and disjoint output. It has
no autograd rule and rejects gradients. Opposite-angle forward execution is not
a backward implementation. Quantization, table gathering and cache packing are
outside this change.

The arithmetic and frequency tables follow
DeepSeek-V4.1-Flash.
The vectorized kernel, API and benchmark glue are authored for this work. The
benchmark loads two unchanged source reference functions from a user-supplied,
hash-checked checkout. FlashInfer controls retain upstream attribution; the
no-FTZ control changes only a compiler flag on the same sources. Reproduction is
documented in benchmark/rope_tail/README.md.

Publication validation: rebased onto develop at
adb963c01b6f65a0622efda48092fbc853dd3cf5. Operation kernels, APIs,
benchmark controls and focused test files are byte-for-byte identical to the
validated versions described above. The rebase preserves the new upstream
lazy exports and documentation entries. Source identity, syntax, Black 26.3.1
(line length 160), SPDX headers and diff checks pass. Performance results above predate
this rebase.

Post-rebase host validation on these exact published heads now passes. A fresh
common C++ binding was built with GCC 11.4, CUDA 13.2.46, cuDNN 9.28.0 and
CPython 3.12.12; all binding sources, headers and CMake inputs are identical
across #1113–#1116. Each PR's byte-verified Python package passes actual import,
lazy-export and missing-dependency checks with that new binding. CPU metadata
checks also pass, with hardware queries explicitly mocked as SM100 and no GPU
kernel execution.

Fresh post-rebase B200 regression passed on 2026-09-17 at
941ee629ed6a1d83a0ec876ad56ec73039768b7b: 61 focused API tests passed,
with zero failures, errors or skips, using the fresh binding and cuDNN 9.28.0.
All 165 focused tests across #1113–#1116 passed on the same B200; loaded
package files and the binding were checked against their source hashes.
This adds correctness and integration validation. Performance figures above
still refer to the original measured versions; this regression did not time kernels.

[coderabbitai]

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

• Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

---

_{Comment @coderabbitai help to get the list of available commands.}