Skip to content

[Kernel][SM70] Add opt-in E4M3 grouped attention with FP32 partials - #524

Merged
yangzhuxinyzx merged 12 commits into
mainfrom
codex/v100-e4m3-fp32-integration-20260906-050258
Sep 7, 2026
Merged

yangzhuxinyzx merged 12 commits into
mainfrom
codex/v100-e4m3-fp32-integration-20260906-050258

Conversation

@yangzhuxinyzx

@yangzhuxinyzx yangzhuxinyzx commented Sep 6, 2026

Copy link
Copy Markdown
Contributor

2026-09-07 final status: memory repairs verified; model gate still fails

Latest evidence head 4cb6ec783b7339b0fd596ad07c943dfc2b2b25ca; model-tested source 0c34be5d6039101771ac492c4d96bb954ab97f33 includes host-only offset-Q protection; arithmetic source beb172ebd0278b7faa4118e2e0caf12047aa6cb2, integrated with main 099d9841f542f1b71121b4aff49e3aa29053a489. Keep Draft/default-off.

Latest model result: FP32 SSM + FP32 logits passes 128K but fails 261888+256 at one-based token 26. Both references are stable; candidate chooses -level, reference -, continuing as "low-level kernel changes" versus "low-precision arithmetic". Both are coherent; this is a failed strict token gate, not a claim of universal semantic degradation. All four ranks confirm actual FP32 recurrent state on 48 GDN layers, FP32 logits, page1616, and the native attention route. All captured rank-zero attention outputs are finite across 16 layers. Raw result SHA256: 452a5a1845cae49b8aa26470a2fffacc4e8009a715ebb7625ec3b0e990c3a722.

SSM Head Page 128K Boundary-256K
FP16 FP16 848 pass fail at token 126
FP16 FP32 848 fail at token 191 not run: stopped by 128K gate
FP32 FP32 1616 pass fail at token 26

The token position does not rank different configurations' general quality. Changing downstream precision alone did not close admission. Next localize the first-divergence prefix; do not repeat this unchanged full-model cohort or relax the token criterion. Two private recording setup failures (Dynamo hook / subclass filter) are retained separately and are not numerical results. The completed run's boundary wait coincides with an on-demand strided GDN build and recovers after it; that dependency uses CUDA 12.0.140, DSO 89337e7055cc8ba8f9bd972341f43010ce23a5a7cb991a84eb48e60bc5bbfaf9. Prebuild/warm it before production timing. Our GPU jobs are finished and released; max concurrent task GPUs was four.

  • The alignment repair below is validated. The synchronized revision-2 build subsequently matched all 256 tokens in a fresh 128K reference/native/reference bracket, with stable references and four-rank route evidence (result 79001ddfa0517c2fc8d1cfb2f017d174f2e3648506905847da69f867f905c22b). This does not establish the cause of the earlier token-77 difference.
  • Revision 3 retains FP32 max/sum and unnormalized PV until the final combine. The old normalized/LSE representation rounds the wrong way in analytically exact FP16-midpoint cases at 128K/256K; new regression tests cover the fix. Production accumulation stays FP32.
  • The workspace changes from [80,8,6] LSE to [80,8,6,2] max/sum (+15 KiB). Both wrapper and backend require native precision version >=3 and reject stale extensions. E5M2 and the ordinary E4M3 q8 route retain their original workspace contract.
  • Fresh DSO 8880b0405d6d7d212c4a738d45ab441815b6c948ab6e4e8ac5e1b83d754a9d49: 102 kernel, 142 routing, 38 planner checks pass; 69 memcheck cases, zero errors. Core/stable/FA2 dependencies remain explicitly pinned to archived libraries, not claimed as rebuilt.
  • All 200 real-input and five midpoint outputs match the screened prototype bitwise. Real-input FP16 disagreements decrease 1872→1640, but 45 groups worsen (125 improve, 30 tie): not uniform L2 improvement. Integrated 100-ABBA q5 speed ratios are 0.99928/1.00000/1.00041/1.00007, essentially unchanged. Neither these counts nor timings are model-quality/60T-prefill claims.
  • The frozen-driver same-process model run passes 128K but fails 261888+256 at one-based token 126. Both FP64-attention references are stable; reference chooses speed, candidate chooses benchmark. Their reference scores tie, while candidate favors benchmark by 0.046875; this is not a waiver of the deterministic gate or proof of broad semantic degradation. All four TP ranks select the route; rank-zero captures cover 16 attention layers and are finite. Result SHA256: a94d356ce29273fba3a202428ea77c737e7f51bd5152be1885d3ed46b6fb700e.
  • Boundary broadcast waits recover and all eight requests finish. The original q5-only summarizer fails on valid late q3 rows; its width-aware correction retains strict token/finite checks, with original failure logs preserved. The earlier startup with a changed driver hash remains excluded from accepted evidence.
  • The FP32-head-only counterfactual also fails, at 128K token 191: candidate validation, stable references scrutiny. Four ranks confirm actual FP32 logits; all recorded attention outputs are finite. The predeclared gate prevents extending this failure to 256K. Result SHA256: 9571a0b77a2dfa161b8e8f0aa90097e3e6f9115b89c5eab06f295394113b043c. Synthetic head arithmetic improves about 80x, but that did not establish model admission.
  • Both earlier failing cohorts explicitly fixed FP16 SSM state. Current main's GDN auto uses FP32 SSM. The completed FP32-SSM counterfactual is reported above; it does not retrospectively approve the FP16-SSM failures or close boundary-256K admission.
  • Additional Q-base safety repair: contiguous half-offset views can violate the uint4 load alignment. The native host entry now copies only those exceptional views (<=24 KiB); aligned inputs retain the same arithmetic and workspace. Fresh DSO c33a84443d42621060a90ff7ecdbe2af4910e7774317c44c74777a13d7b56aaf passes 105 kernel checks, 6 targeted Q/KV alignment memcheck cases with zero errors, and matches all 200 old aligned real-input outputs bitwise. Four q5 100-ABBA speed ratios are 1.01834/1.00065/1.00036/1.00000, not a new model-speed claim. The separate 142/38/69 counts above retain their earlier artifact attribution.

Revision 3 is not model-admitted. No merge, auto-merge, default KV promotion, service restart, or private v18/v37 prefill integration is claimed. Historical results below remain artifact-specific.

Earlier 2026-09-07 alignment update (revision 2)

Recorded head 1eeb7fb7a674d6b5b44682dbe5210d8be41cf76c, integrated with main 099d9841f542f1b71121b4aff49e3aa29053a489. At that update the refreshed same-prompt 128K run was pending; its later result is stated above.

  • Integration commit 4d889d0c1dde3589bced4211b984047a3d028609 retains both native capability registrations and prevents ROW_SEQLENS from inheriting the dense-q8 16-byte paired loader. Valid 8-byte KV strides keep their tested loader; main's q8 E4M3 acceleration remains intact.
  • Fresh DSO b65ee698fa992b6ff933f495c3827390aef3376641464c399b27939a3f9132b3: 98 kernel, 142 routing, 38 planner checks pass. The planner now uses explicitly pinned archived core dependencies; this is not a full core rebuild.
  • Three new token/block/both padded-stride graph cases match contiguous controls bitwise and pass memcheck with zero errors. This closes the reviewed alignment blocker for this source/artifact.
  • All 200 real-input outputs match the prior R6 bitwise. Same-GPU 20-warmup/100-ABBA q5 speed ratios at approximately 8K/64K/128K/256K are 1.00255/1.00029/1.00064/0.99990: essentially unchanged operator speed, not a model throughput claim.
  • Current model source is commit 4d889d0c1d (the next commit updates documentation only), frozen original prompt IDs, TP4/MTP4/E4M3, unchanged deterministic diagnostic settings, and the new DSO. No production default, service, or global KV alias changed. A private max/sum-state precision experiment has only compiled and is not included in this PR.

The older test/model evidence below remains historical and artifact-specific. AI code review is not human review/test confirmation; no approval or auto-merge is asserted.

Historical R6 purpose and evidence (superseded by revision 3 above)

Integrate opt-in SM70 E4M3 dense small-query attention with FP32 partial state, an explicit E4M3 paged-to-FP16 bridge API, and follow-up arithmetic repairs. Preserve FP32 accumulation, current production defaults, and existing services.

The historical R6 same-process 128K model gate failed at output token 77. Operator improvements did not override that failure. This section records that artifact; it does not describe the latest model result. Human line-by-line review and relevant test confirmation remain outstanding.

Scope and non-duplication

  • Dense single-request q2–8, GQA6/D256, explicit E4M3 KV, per-query device lengths, and graph padding. Independent batches and unsupported shapes retain the existing fallback.
  • VLLM_FLASH_V100_E4M3_GROUPED_FP32 remains off. The opt-in requires native precision capability version >=2; stale libraries fail closed with a rebuild warning.
  • Repair QK with compensated FP32 summation, represent PV probability with high/residual FP16 operands, and accumulate each N32 PV tile separately before updating FP32 online state. The residual is multiplied by 2048, with an exactly representable inverse scale on E4M3-derived V.
  • The native/Python E4M3 bridge API is tested but not yet selected by the backend. No global KV alias/default change.
  • Existing E5M2 q8/q16 and sparse-page4 template contracts remain unchanged. Main's independent allocation-invariance planner repair is preserved.
  • This updates existing [Kernel][SM70] Add opt-in E4M3 grouped attention with FP32 partials #524. [Kernel] Preserve DFlash2 sampling precision and repair QUASAR TP2 #517 has a different DFlash2 q8/logits/context-pipeline scope; no duplicate PR was opened or modified.
  • Private writer/B1 decode, private FP32 long-prefill v18/v37, and rejected paged SplitKV are not claimed as integrated. The older stable 60/61T prefill already in main is not replaced by this small-Q PR.

Original base: 755baae1d075ee04fa9096b23fc0225b23589a86. Arithmetic follow-up: 52dcdc4d523d8e7331897f073bada070ad468cc1. Synced main: 95205a2d9952813aa7469f63ff65b8f2813c027a, merged without rewriting branch history.

Historical R6 model gate: candidate rejected

Qwen3.8-27B-FP8, TP4/MTP4, explicit E4M3 KV, FP16 activations/SSM state, graph execution, fixed diagnostic GEMM dispatch, one cached 128K prefix, and three 256-token greedy requests in one process: FP64 attention reference–native candidate–FP64 attention reference. The reference replaces only admitted small-Q attention on identical quantized KV; this is not an all-FP64 model or a performance benchmark.

Both reference requests produce identical tokens. The candidate first differs at one-based token 77: references choose token 799 (one), candidate chooses 23438 (engineers). The reference top-two logprobs tie; the candidate margin is 0.015625. Near-ties do not waive the no-divergence gate. This is not proof of broad semantic degradation or that E4M3 is inferior to E5M2.

All four workers hit the native route; recorded rank-zero outputs cover all 16 full-attention layers and are finite. Result SHA256: 911c9906aed7ecf48d2673ce8267f2d27d2de22da65e55a0ec7ae36bbbff1926, private bundle e4m3-scaled-residual-r6-paired128-20260906.

An earlier unscaled-residual revision passed its own same-process 128K bracket. Its reference sequence differs from this cohort's reference, so the two cohorts cannot establish that residual scaling alone caused the model difference. The earlier cross-process token-151 counterexample and original token-102 failure remain recorded in the design report. Full approximately-256K reference/control acceptance is incomplete.

Test plan and results

Build the Flash-V100 extension with setup.py build_ext into a task-owned output directory using CUDA 12.8.93/GCC12, Torch 2.10+cu128, Python 3.12, SM70, and standard fast-math flags. Set PYTHONPATH to the new DSO and this checkout's Flash-V100 package. No installed runtime binary is overwritten.

Model-tested arithmetic artifact: 87ffc1fd626e3d0fa46f6e2bda741a88d79d53809685a4d569c37e4d35b325cd.

  • 89 GPU regression checks pass; 62 Compute Sanitizer checks pass with zero errors.
  • Negative controls reproduce the old biased-PV and small-probability-residual failures before final FP16 rounding; repaired assertions are unchanged and pass.
  • 200 real-input groups: disagreements with the FP64 result rounded to FP16 decrease from 2092 to 1872 elements; another 32 real 128K groups decrease from 348 to 320. These are attention tensor elements, not token errors or total FP8 quantization loss.
  • Same GPU, q5/page848, 20 warmups/100 ABBA graph samples at approximately 8K/64K/128K/256K: speed ratios 1.0007/1.0000/1.0007/1.0007 versus the preceding repair. This is essentially unchanged operator speed, not 60T prefill or TP4 throughput admission.

Fresh post-main-sync artifact: 76aa9a19f197fe805bb97b037f8a7934129ef0752b5c66922246c241a7e97b42.

  • Rebuild successful; 89 kernel and 138 routing-policy checks pass.
  • All 200 retained real-input small-Q outputs are finite and bitwise identical to the pre-sync R6 library. Replay SHA256: eb187e7fea4c81e4f01bb5ebb3752c888e583b0d62fc3bd63dec37bf2cde73f9.
  • Main's planner suite: 30 native checks pass; 8 integration checks fail at import because this source-only worktree lacks vllm._C, before numerical assertions. The failing setup log is retained; the full 38-test suite is not claimed as passed.
  • The post-sync extension has not received new whole-model admission. Archived model results are pinned to their own DSO and diagnostic launch contract, not silently transferred to new main.
.venv/bin/python -m pytest --confcutdir=tests/kernels/attention -q \
  tests/kernels/attention/test_sm70_fp8_bridge_formats.py \
  tests/kernels/attention/test_sm70_grouped_e4m3_fp32.py \
  tests/kernels/attention/test_sm70_flash_v100_grouped_verify.py \
  tests/kernels/attention/test_sm70_qsa_grouped_page4.py

.venv/bin/python -m pytest --confcutdir=tests/v1/attention -q \
  tests/v1/attention/test_sm70_e4m3_grouped.py \
  tests/v1/attention/test_sm70_flash_v100_policy.py

.venv/bin/python -m pytest --confcutdir=tests/kernels -q \
  tests/kernels/test_sm70_qsa_page4_plan.py

Remaining admission

Resolve the within-process token counterexample; complete approximately-256K reference/control, current-main model checks, and matched prefill/decode performance. Complete separate clean-source integration/routing for private long-prefill and remaining E4M3 paths. Do not enable the global alias based on arithmetic microbenchmarks. This PR makes no current-main model-quality or production-speed acceptance claim.

See docs/design/sm70_e4m3_grouped_fp32.md and the migration control log. Paper drafts, model weights, private activations, binary build products and machine-specific raw paths stay out of Git. Tests use at most four physical GPUs concurrently; no unrelated jobs or services are stopped.

AI assistance: OpenAI Codex. Human review/test confirmation is not claimed, and automated merge is not enabled.

Preserve explicit per-query causal lengths and zero graph padding for the single-request small-Q route. Keep legacy E5M2 and sparse defaults unchanged.

Assisted-by: OpenAI Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Assisted-by: OpenAI Codex
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
yangzhuxinyzx and others added 3 commits September 6, 2026 19:05
Compensate QK and probability residuals, keep PV accumulation tile-local, and reject stale precision binaries. Preserve default-off routing and record the failed same-process 128K token gate; this is not model promotion.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
…integration-20260906-050258

Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>

# Conflicts:
#	docs/design/sm70_v100_migration_control.md
Retain current-main migration history and distinguish verified small-Q replay from incomplete full-model and planner integration admission.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
@yangzhuxinyzx

Copy link
Copy Markdown
Contributor Author

AI integration review — 2026-09-07 — changes required

Reviewed head 42462927db3ece21e97bc31d287a0bb64265a8f6; fetched integration main 099d9841f542f1b71121b4aff49e3aa29053a489. This is an OpenAI Codex code/evidence review, not human approval or a passed model gate. Both published CI checks now pass.

P1: reconcile the new main E4M3 paired-load contract before integration

Main #517 enables FP8_PAIR_LOAD for dense E4M3 instantiations of the shared grouped kernel. A plain conflict resolution would also enable it for this PR's ROW_SEQLENS/FP32 instance. The loader addresses uint4 via physical_offset / 16, but this PR's Python admission and native entry only require KV strides divisible by 8 (sm70_e4m3_grouped.py:70, flash_decode_paged.cu:4273, line numbers at reviewed head).

A CPU-only counterexample uses valid non-overlapping uint8 KV shape [1,848,1,256], strides [223872,264,256,1], and a 16-byte-aligned base. The current admission returns true. The second token begins at byte 264; main's paired loader would instead address byte 256. This is a proposed-merge interaction, not a reproduced defect in the previously tested 8-byte-loader DSO and not an explanation of its token-77 difference.

Retain the tested 8-byte loader for the explicit-row FP32 specialization, or strengthen both dispatch and native stride guards to the new 16-byte contract with an appropriate fallback. Add an 8-but-not-16-byte stride negative/control case before rebuilding; existing contiguous/interleaved cases do not exercise it. Verify E5M2 and sparse-page4 remain unaffected.

P1: deterministic model admission remains failed

Rechecked the archived R6 same-process 128K FP64-reference/native/FP64-reference summary and raw-result hash 911c9906aed7ecf48d2673ce8267f2d27d2de22da65e55a0ec7ae36bbbff1926: both reference requests match all 256 tokens; candidate first differs at token 77. Native selection is recorded on all four TP ranks. Finite outputs and lower L2 do not clear the declared no-divergence gate. Near-ties do not prove broad semantic degradation, nor waive this gate. No new model run was performed in this review.

Integration/test limitations

git merge-tree --write-tree --name-only HEAD onecat/main reports conflicts in docs/design/sm70_v100_migration_control.md and flash-attention-v100/kernel/fused_mha_api.cpp. The registry resolution must preserve main's grouped_verify_e4m3 capability attribute alongside the new FP32 entry and revision query. Do not resolve it wholesale from one side.

The archived 89 kernel/138 routing passes and 200/200 bitwise real-input replay concern the prior 95205a2d99 integration artifact. They do not validate today's merge with #517/#532. The additional planner run previously had 30 passes and 8 setup failures because vllm._C was absent. Approximately-256K full reference/control admission is still incomplete.

Recommendation: keep Draft and do not merge/default-enable until the integration contract is repaired and rebuilt, affected regression/model gates pass, and repository-required human review/test confirmation is supplied. No source, production default, service, or GPU allocation was changed during this review. Private v18/v37 long-prefill integration remains outside #524.

yangzhuxinyzx and others added 7 commits September 7, 2026 00:13
Retain the q8 E4M3 capability and add padded-stride graph regressions. Runtime/model acceptance remains pending; no default promotion.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Keep model admission pending and distinguish pinned core dependencies from the rebuilt Flash-V100 extension.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Preserve max/sum and unnormalized PV, version the workspace contract, and cover exact FP16 midpoint rounding. Keep the route experimental pending integration and model gates.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Retain the non-monotonic per-input results and the aborted startup provenance; model admission remains pending.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Preserve the completed 128K pass and boundary-256K failure without transferring operator admission to model acceptance.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Clone only exceptional unaligned contiguous Q on the active stream. Preserve aligned arithmetic and graph behavior; retain failed model counterfactuals in the admission log.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
Record the completed FP32 SSM/head brackets and late GDN build provenance; do not approve the failing boundary route.

Co-authored-by: OpenAI Codex <noreply@openai.com>
Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
yangzhuxinyzx added a commit that referenced this pull request Sep 7, 2026
…20260907-100628

[Kernel][SM70] Integrate repaired E4M3 FP32 small-query attention (#524)
@yangzhuxinyzx
yangzhuxinyzx merged commit d12d4e4 into main Sep 7, 2026
2 checks passed
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