feat(kernels): Rust wrapper for the hd256 windowed paged prefill

[FeathBow]

Description

Closes #882.

The windowed C entry points landed with no callers; this wraps the prefill one in the Gemma 4 serving shape — attention read only, the pool-write prep owns the scatter — with sm_scale explicit (Gemma 4 runs unscaled attention; the scale is model policy, not geometry) and window_left under its inclusive-distance contract.

• Driven by a PrefillPagedPlan, cross-checking the plan's recorded cta_tile_q against the kernel's own derivation, so a plan tiled under an override cannot be silently misread.
• Pool geometry comes through the shared checked helper — renamed checked_paged_geometry, now serving prep writes and attention reads — and scalar narrowing on this path goes through checked conversions throughout, in the wrapper and in the plan alike, so a length past i32 fails rather than truncating into a legal-looking bound. The batch tile query also moves its positivity check into the i32 domain, where the single-request path already had it.
• Two contracts move into PrefillPagedPlan, where the arithmetic that needs them lives. The GQA group size is an integer division by num_kv_heads — in the plan's own tile arithmetic, in the C tile queries and again in the kernel, none of which guard the divisor — so every plan settles the head relation before dividing, and the wrapper is left checking only that the plan and the layout agree. And because the attention kernels compute addresses from page ids and the last-page length with no device-side bounds check, the plan records both bounds from its host inputs, at construction and on every in-place update, committed only after the update fully succeeds; the wrapper compares them against the pool geometry before every launch.

The windowed decode C entry deliberately stays wrapper-less: serving decodes through this entry at one token, and the dedicated decode form (non-partitioned versus split-KV) is a measurement-gated decision for the batching stage.

Test Env

Single GPU (sm_89, x86_64), CUDA 12.9, PEGAINFER_REQUIRE_GPU=1.

Verification

• No numeric gate lands here by design: the window mask is the kernels' behaviour, instantiated by the already-merged C entries; the wrapper surface is parameter plumbing plus validation. The end-to-end numeric gate for this read path arrives with the KV serving slice's HF oracle (greedy decode against the reference), which exercises the wrapper on the real checkpoint.
• The validation contracts are the deliverable: contradictory layouts, out-of-range layers, plan page metadata past the pool geometry and non-finite scales all fail as Err before any launch, through the same checked-geometry seam the prep wrappers use.
• A device gate (hd256_window_prefill_reject) drives the reject paths through the public surface: a page id past the pool and a last-page length past page_size both fail as Err before any launch and the same shapes in range launch; a batch-plan length past i32 fails rather than truncating into a legal-looking bound; and a zero KV head count fails when the plan is built, before anything divides by it.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update

[FeathBow]

@codex review please

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: ac3f8b8a9b

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".