Skip to content

[Bugfix] Align SM70 paged prefill accumulators - #226

Merged
yangzhuxinyzx merged 1 commit into
mainfrom
codex/v100-paged-prefill-align-20260816-2310
Aug 16, 2026
Merged

yangzhuxinyzx merged 1 commit into
mainfrom
codex/v100-paged-prefill-align-20260816-2310

Conversation

@yangzhuxinyzx

Copy link
Copy Markdown
Contributor

Purpose

Port and complete the hardware gate for #123 / #98 on current main. The D=128 paged-prefill FP32 WMMA accumulator buffer started at a 16-byte boundary, while WMMA matrix load/store requires 256-bit (32-byte) alignment. The resulting undefined behavior was observed as run-to-run prefill jitter and could fork greedy output after server restart.

Change

  • Align the accumulator buffer to 32 bytes.
  • Explicitly align the dynamic shared-memory base and assert the accumulator offset at compile time.
  • Add the six-replay physical-SM70 regression from [Bugfix] Align SM70 paged prefill accumulators (#98) #123, updated to the current accelerator synchronization API.

Test Plan

  • Build the bundled Flash-V100 extension from this branch for sm_70.
  • On a physical V100, run direct-extension and public-wrapper D=128 paged prefill six times for L=300 and L=512 and require bitwise equality.
  • Compare both shapes against an independent FP32 causal-GQA reference.
  • Run focused Ruff and diff checks.

Test Result

  • CUDA 12.8 / Torch 2.9.1, V100-SXM2-16GB: extension build passed.
  • pytest --confcutdir=tests/kernels/attention tests/kernels/attention/test_sm70_flash_v100_paged_prefill_determinism.py: 2 passed.
  • All six direct and public replays were bit-exact for both L=300 and L=512.
  • FP32 reference:
    • L=300: finite, max abs error 0.00086688995, mean abs error 3.2426e-05, cosine 0.9999999404.
    • L=512: finite, max abs error 0.00109624863, mean abs error 2.6168e-05, cosine 1.0.
  • Ruff check and format check passed; git diff --check passed.

Supersedes #123 without changing the kernel algorithm or dispatch.

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
(cherry picked from commit e9a7f9a)
@yangzhuxinyzx
yangzhuxinyzx marked this pull request as ready for review August 16, 2026 14:46
@yangzhuxinyzx
yangzhuxinyzx merged commit f2a7937 into main Aug 16, 2026
3 of 4 checks passed
@yangzhuxinyzx
yangzhuxinyzx deleted the codex/v100-paged-prefill-align-20260816-2310 branch August 16, 2026 14:55
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant