feat(qwen35): add SM120 FlashInfer GDN prefill candidate - #862

Thanks for exploring this direction. I think keeping the FlashInfer CuTe DSL kernel is reasonable: if we AOT-compile it, the serving runtime can avoid Python, Triton, JIT compilation, and JIT cache management.

The main change I would suggest is to simplify the integration boundary before expanding the candidate further. Today, compile_sm120.py:86-88 extracts FlashInfer's patched/raw PTX, while flashinfer_gdn.rs:65-72 freezes the generated entry symbol and naked launch ABI and flashinfer_gdn.rs:426-459 reconstructs that ABI in Rust. At the same time, lib.rs:101-102 explicitly keeps production launch APIs Triton-only. This leaves PegaInfer maintaining a private CuTe/PTX/TMA contract without yet achieving the original runtime goal.

Could we reshape the PR around this smaller end-to-end boundary?

1. Let pegainfer-kernels own the pinned, reproducible CuTe build, patched artifact, manifest, and a stable generated C ABI.
2. Let pegainfer-qwen35 call only that wrapper; it should not know the generated CuTe symbol or PTX/TMA argument layout.
3. Wire the supported SM120 case into the real production dispatch, with unsupported shapes/configurations handled explicitly.
4. Validate correctness and report end-to-end A/B measurements on that exact production path before treating it as a performance win.

pegainfer-kernels/tools/cutedsl/export_glm52_fp8_dsl.py:181 is the closest repository pattern. FlashInfer's SM120 path additionally patches PTX after cute.compile, so the export flow must preserve that patched artifact (ideally through an upstream-supported export path) rather than applying export_to_c blindly.

This keeps the useful CuTe kernel and the goal of removing Triton from serving, while giving the runtime a boundary we can maintain.

Thanks, this makes sense. I’ll reshape the PR around the smaller integration boundary you suggested

Thanks — the new C ABI boundary is a substantial improvement.

One remaining request is scope reduction. The PR is still +6,711 lines across 43 files, and much of that is staged diagnostics and benchmark infrastructure rather than the production integration itself. Please keep this PR focused on the production wrapper/dispatch, one reproducible generation path, and the essential end-to-end correctness/performance checks. The stage7/stage13 diagnostics, ABBA harnesses, and additional research scaffolding can move to follow-up PRs or external artifacts.

That would make the runtime change much easier to review and maintain without losing the useful validation work.

Following the review feedback, I have completed the integration-boundary refactor and production-path validation for the SM120 FlashInfer GDN prefill candidate.

The review requested that the FlashInfer CuTe kernel remain AOT-compiled while moving all generated CuTe/PTX/TMA details out of pegainfer-qwen35. The current integration is now:

FlashInfer/CuTe generation environment
→ patched AOT object with embedded SM120 CUBIN
→ manifest + statically linked CuTe runtime
→ PegaInfer-owned versioned C ABI
→ pegainfer-kernels Rust wrapper
→ pegainfer-qwen35 semantic GDN operation
→ production dispatch

pegainfer-kernels now owns the pinned source, patch, generator, artifact validation, generated kernel symbol, TMA descriptors, raw launch ABI, native prepare kernel, and static runtime linkage.

pegainfer-qwen35 only sees the semantic Qwen35GdnAot interface, geometry, typed device buffers, workspace, and backend evidence. It no longer constructs or depends on generated CuTe symbols, TMA layouts, PTX, or the raw C launch ABI.

The supported single-GPU SM120 Qwen3.5-4B Hv32 case is now connected to production dispatch, while unsupported capabilities explicitly retain the Triton fallback. Invalid artifact/ABI/hash/geometry combinations fail validation rather than silently executing an incompatible kernel.

Serving-path validation

The serving path was also tested with an empty environment and with PTX JIT disabled.

Real inference completed successfully using the embedded native CUBIN and statically linked runtime, without requiring Python, Triton JIT, FlashInfer Python, or the CuTe compiler at serving time.

Cross-machine validation

I validated the refactored code on two RTX 5090 machines:

Both machines used CUDA Toolkit 12.8 and the same source commit, FlashInfer submodule, generator, patch, and pinned package versions.

Each machine generated its object reproducibly within that machine, but the object hashes differed across machines. Both objects were exactly 291,336 bytes and passed the complete correctness suite.

This validates the runtime ABI and numerical behavior, while bitwise cross-machine artifact reproducibility remains an open build-boundary issue.

Correctness

On both machines, the Hv32 production candidate passed:

• stable-ABI alias/separate-state tests;
• all operator lengths T=1,2,63,64,65,127,128;
• short- and long-sequence HF fixtures;
• chunked/resumed prefill;
• prefill-to-decode handoff;
• scheduler and slot-reuse coverage;
• CUDA Graph coverage.

Production dispatch selected FlashInfer, and the successful-launch counters increased as expected, confirming that the tests did not silently fall back to Triton.

The production HF short replay also remained stable across machines:

The broader HF comparison against the same oracle remains:

FlashInfer therefore has a slightly better overall error distribution, although it is not uniformly better on every tail metric.

The non-production Hv48/T128 result is unchanged: FlashInfer has five localized state-tail violations versus two for Triton. The A/B attribution continues to show that these originate from the upstream FlashInfer kernel rather than the PegaInfer HKV layout patch.

Performance

Machine 1 completed a 15-case same-production-path ABBA matrix covering:

• T=63/64/65/128/2048
• c1/c4/c8

FlashInfer won TTFT p50 in 14/15 cases and throughput in 12/15 cases.

The case-level geometric-mean improvements were:

• TTFT p50: +4.26%
• Throughput: +3.94%
• TPOT p50: -0.11% (effectively unchanged)

One case, T2048/c8, regressed by 1.95% in TTFT p50 and 2.59% in throughput. In that profile, the FlashInfer GDN operator was 7.94% slower despite reducing the GDN launch count by 3×.

Machine 2 repeated representative cases:

On Machine 2, the earlier c8 regression did not reproduce: T2048/c8 was effectively at parity.

Nsight Systems measured FlashInfer GDN GPU time as:

• 28.9% lower at T128/c1;
• 6.0% lower at T2048/c1;
• 6.3% lower at T2048/c8.

Scratch allocation remains approximately 48.52 MiB for FlashInfer versus 184.50 MiB for Triton, a 73.70% reduction.

The c8 traces show that per-launch FlashInfer time remains stable, while the current single-sequence operation is repeated once per sequence.

Based on this analysis, the next optimization directions are multi-sequence batching and native-prepare optimization.

The existing Triton capability fallback and Hv48 diagnostic scope remain unchanged while the maintainers decide the final release and distribution policy.

After the optimization work is complete, I will also clean up the PR itself by removing the staged diagnostic tooling, benchmark infrastructure, and other temporary research scaffolding that is not required for the production integration. The goal is to keep the final PR focused and minimal, preserving only the production integration, reproducible generation path, and essential correctness/performance validation needed for long-term maintainability.

Following the review feedback, I have completed the integration-boundary refactor, scope reduction, native-prepare optimization, and production-path validation for the SM120 FlashInfer GDN prefill candidate.

Current integration boundary

The integration now follows this path:

FlashInfer/CuTe generation environment
→ patched AOT object containing an SM120 CUBIN
→ validated manifest + statically linked CuTe runtime
→ PegaInfer-owned versioned C ABI
→ pegainfer-kernels Rust wrapper
→ pegainfer-qwen35 semantic GDN operation
→ production dispatch

pegainfer-kernels owns the pinned source, HKV patch, production-only generator, artifact validation, generated kernel symbol, TMA descriptors, raw launch ABI, native prepare kernel, and static runtime linkage.

pegainfer-qwen35 only sees the semantic Qwen35GdnAot interface, supported geometry, typed device buffers, workspace, and backend evidence. It no longer constructs or depends on generated CuTe symbols, TMA layouts, PTX, or the raw C launch ABI.

The supported single-GPU SM120 Qwen3.5-4B Hv32 case is connected to production dispatch. Unsupported capabilities explicitly retain the Triton fallback. If an eligible FlashInfer configuration selects an invalid, incomplete, or ABI-incompatible artifact, validation fails instead of silently executing another kernel.

The README now documents the complete reproducible local path:

1. Initialize the pinned FlashInfer submodule.
2. Create the isolated Python 3.12.3 / CUDA 13 generation environment.
3. Install the locked dependencies.
4. Generate the production-only bundle.
5. Validate the bundle.
6. Configure the build variables.
7. Link the release binary.

PR scope reduction

The PR has been reduced from 43 files and approximately 6,711 added lines to 34 files and approximately 3.7k added lines.

The stage7/stage13 harnesses, ABBA infrastructure, profiling scripts, generated artifacts, logs, benchmark JSON, and other temporary research scaffolding have been removed.

The remaining scope is limited to:

• the production wrapper and dispatch;
• one reproducible artifact-generation path;
• the stable ABI and manifest contract;
• essential operator and production-path correctness tests;
• the native Hv32 prepare implementation.

Serving validation

The serving path was tested with an empty environment and again with PTX JIT disabled.

Real inference completed using the embedded native CUBIN and statically linked runtime, without requiring Python, Triton JIT, FlashInfer Python, or the CuTe compiler at serving time.

Cross-machine validation

The refactored implementation was validated on two RTX 5090 machines:

Both used CUDA Toolkit 12.8 and the same source commit, FlashInfer submodule, generator, patch, and pinned package versions.

Generation was bitwise reproducible when repeated on the same machine. The objects differed across machines, although both were exactly 291,336 bytes and passed the complete correctness suite.

The runtime ABI and numerical behavior are therefore validated, while cross-machine bitwise artifact reproducibility remains an open build/distribution question.

Correctness

The Hv32 production candidate passed:

• stable-ABI alias/separate-state coverage;
• operator lengths T=1,2,63,64,65,127,128;
• short- and long-sequence HF fixtures;
• chunked/resumed prefill;
• prefill-to-decode handoff;
• scheduler and slot-reuse coverage;
• CUDA Graph coverage;
• native-prepare dynamic-length and non-finite-input coverage.

Production tests reported:

selected_backend=flashinfer

and successful-launch counters increased as expected, confirming that the tests did not silently fall back to Triton.

The latest final gate on commit 9e70e86a regenerated the production-only artifact and passed artifact validation, release linking, native-prepare correctness, and production HF short replay:

The broader comparison against the same HF oracle remains:

FlashInfer has a lower mean error in four of five replay surfaces and a lower maximum error in four of five, but a lower p99 in only two of five.

It is slightly better overall, but not uniformly better on every tail metric.

The non-production Hv48/T128 investigation remains unchanged: FlashInfer has five localized final-state tail violations versus two for Triton.

Patched-HKV versus upstream-HVK A/B testing produced identical results after layout conversion, attributing these violations to the upstream kernel rather than the PegaInfer HKV patch.

Performance

Machine 1 completed a 15-case same-production-path ABBA matrix covering T=63/64/65/128/2048 and c1/c4/c8.

FlashInfer won:

• TTFT p50 in 14/15 cases;
• throughput in 12/15 cases.

The case-level geometric-mean results were:

• TTFT p50: 4.26% better;
• throughput: 3.94% better;
• TPOT p50: 0.11% worse, effectively unchanged.

The T2048/c8 case regressed by 1.95% in TTFT p50 and 2.59% in throughput.

In that trace, the FlashInfer GDN operator was 7.94% slower despite reducing the GDN launch count by approximately 3×.

Machine 2 repeated representative cases:

The earlier c8 regression did not reproduce on Machine 2; T2048/c8 was effectively at parity.

Nsight Systems measured FlashInfer GDN GPU time as:

• 28.9% lower at T128/c1;
• 6.0% lower at T2048/c1;
• 6.3% lower at T2048/c8.

FlashInfer scratch allocation is approximately 48.52 MiB versus 184.50 MiB for Triton, a 73.70% reduction.

The native Hv32 prepare optimization was also validated independently. Relative to the previous production implementation, the ABBA means were:

The optimization therefore provides a small, consistent improvement without changing the selected artifact or production numerical gates.

Multi-sequence batching remains a follow-up and is not part of the current PR.

Maintainer decisions still needed

The remaining policy and scope decisions are:

1. Whether the current Triton capability fallback should remain long-term or eventually be removed for supported SM120 configurations.
2. Whether official artifacts should be distributed as prebuilt bundles, generated locally by users, or support both paths.
3. Which environment should be the canonical artifact build environment, given the observed cross-machine object-hash difference.
4. Whether the non-production Hv48 generic prepare and diagnostic support should remain in this PR or move to a follow-up. It is not required by the Hv32 production artifact or dispatch.
5. Which reference machine and variance threshold should define the final performance acceptance criterion, particularly for T2048/c8.

The current implementation keeps the conservative Triton capability fallback and retains the Hv48 diagnostic path until these decisions are made.

The test suite is still broader than the production confidence it buys. I suggest reducing it around five meaningful gates:

• one real AOT bundle exercised through the production build.rs / link / load path;
• the native prepare kernel checked against a CPU oracle;
• the production dispatch path checked against the HF golden;
• chunked vs. unchunked recurrent-state continuation;
• the real scheduler/CUDA-graph E2E.

Several micro-tests can be deleted or collapsed without reducing confidence: production_geometry_is_exact_hv32 compares a constant with the same constant; the token-extent test repeats a function that has no token-extent input; the patch-scope test relies on string counts; and the packaging reproducibility test uses identical fake header/object bytes rather than a real generated artifact. The duplicated source-lock/bundle-index/layout arithmetic tests also mostly validate helper code, not the contract consumed by production build.rs. The struct-layout test should check field offsets or use generated bindings; size/alignment alone will not catch field reordering.

Most importantly, the high-value GPU tests are all ignored, while the Python contract tests are not wired into GitHub CI. Current CI therefore does not prove that the SM120 production path can build, link, and launch. Replacing the mocked artifact micro-tests with one real production-boundary integration gate would reduce the diff and increase confidence.

Thanks, I learned a lot from this review. I’ll simplify the test coverage further and add a real integration gate that exercises the production path end to end, so the tests actually demonstrate that the SM120 path can build, link, launch, and produce correct results.

Thanks again for the thoughtful review and guidance！

Thanks — this is moving in the right direction. The fake artifact/string-contract tests are gone, the ABI offset checks are now meaningful, and the five-gate runner is a much better validation shape.

Could you do one more reduction pass while keeping those five gates?

• Gate 4 only needs to prove real-model continuation: unchunked vs. chunked recurrent state, last-token logits, and the first decode step. The new synthetic fixture, full CPU GDN recurrence, raw operator launcher, and duplicate comparison/log-softmax helpers overlap with Gates 1-3 (real AOT launch, native-prepare CPU oracle, and HF parity), so they can be removed.
• Please move the remaining large test modules out of the production files (a child prefill/tests.rs module can still access private parent items). prefill.rs is now 1,161 lines and recurrent.rs is 1,484 lines.
• The six CUDA-graph evidence counters currently add unconditional atomics to the normal serving path and expand the public runtime-evidence API only for this test. Please either remove them or put them behind a non-default validation feature used by the gate runner.
• unsupported_geometry_is_explicit can be folded into the real capability/load boundary rather than kept as a separate three-branch micro-test.

Please keep the real bundle/ABI gate, native-prepare oracle, HF golden, reduced real-model chunk-continuation gate, scheduler/CUDA-graph E2E, and the fail-closed manual runner.

Thanks — I really appreciate the review and the concrete suggestions.

This revision was pushed by Codex while I was still planning to do another cleanup pass before uploading it, especially to reduce the overall diff. Looking at the current version together with your feedback, there are clearly still several places where I can simplify the implementation further.

I’ll do another reduction pass with the following changes:

1. Reduce Gate 4 so that it only compares the real model under 128-token unchunked vs. 64+64 chunked execution, checking:

   • recurrent state and conv state for all linear layers;
   • last-token logits and token;
   • first-decode logits and token.

   I’ll remove the synthetic fixture, full CPU GDN recurrence, raw operator launcher, and duplicated comparison/log-softmax helpers that overlap with Gates 1–3.

2. Move the large tests currently embedded in prefill.rs and recurrent.rs into private child tests.rs modules. This keeps the tests able to access private parent items without expanding the production API.

3. Put the six CUDA Graph evidence counters and the associated runtime-evidence API behind a non-default gdn-validation feature used only by the validation runner, so the normal serving path does not pay for those atomics.

4. Fold unsupported_geometry_is_explicit into the real capability/load boundary instead of keeping it as a separate micro-test.

5. Continue reducing the overall diff and keep the production changes as small, clean, and consistent with the existing project structure as possible.

I’ll keep the real bundle/ABI gate, native-prepare CPU oracle, HF golden, the reduced real-model chunk-continuation gate, scheduler/CUDA-graph E2E, and the fail-closed manual runner.

Thanks again for taking the time to review this carefully. I’ve learned a lot from the feedback, especially about keeping validation focused on production confidence rather than adding test infrastructure for its own sake.

@xiaguan
The SM120 implementation and validation are complete on the final PR head, 21be3e50.

The production path now connects the generated FlashInfer/CuTe bundle through build.rs, static linkage, the versioned C ABI, the kernels-owned Rust wrapper, Qwen3.5 production dispatch, and the real scheduler path. The validated object is d92be4c9…efa17.

Production gates

All five exact-filtered gates ran on an RTX 5090 with driver 580.126.09 and CUDA Toolkit 12.8. Each matched exactly one test and completed with 1 passed / 0 failed / 0 ignored.

1. Real AOT bundle and ABI

   A real generated bundle passed manifest validation, production build/link/load, and an actual SM120 launch. The gate checks all C/Rust ABI field offsets, size and alignment, rejects Hv48 at the real capability/load boundary, and launches T={1,63,64,65,128}. At T=65, separate-state and in-place aliasing produced bitwise-identical output and final HKV state.

2. Native prepare CPU oracle

   Q, K, V, alpha and beta were checked against an independent CPU oracle for T={1,63,64,65,128,2048}. All comparisons passed. Representative Q/K/V and gate non-finite inputs were detected; the chunk-owned status remained sticky across layers and was cleared for a new chunk.

3. Production dispatch against HF golden

   The gate loads the pinned Qwen3.5-4B model through production Auto dispatch:

   • selected_backend=flashinfer
   • AOT launches: 0 → 288
   • 12 sequences, 108 positions and 864 head deltas
   • mean logprob delta: 0.0258 (≤ 0.06)
   • p99: 0.1025 (≤ 0.20)
   • max: 0.1367, reported diagnostically

4. Real-model chunk continuation

   Independent zero-initialized states executed the same prompt as T=128 and T=64+64. The gate compares every element of all 24 recurrent and convolution states, last-token full-vocabulary logits/token, and the first decode step.

   • maximum recurrent mean/p99: 0.00011414 / 0.00111641, within 0.00025 / 0.002
   • maximum convolution-state mean/p99: 0.01032218 / 0.046875, within 0.015625 / 0.0625
   • last-token logits mean/p99: 0.01413624 / 0.046875
   • last token: 220 == 220, regret 0
   • first-decode logits mean/p99: 0.01250439 / 0.046875
   • first decode token: 200 == 200, regret 0

5. Scheduler/CUDA Graph E2E

   The real scheduler path completed state handoff, FlashInfer launch, CUDA Graph capture and replay, state-slot copy/reuse and slot compaction. The gate requires eager fallback to remain zero; all assertions passed.

The final test layout keeps the large test bodies in private child/integration-test modules. CUDA Graph counters and runtime-evidence APIs exist only under the non-default gdn-validation feature and do not enter the default production object, hot path or public API. The standalone unsupported-geometry test is folded into Gate 1, and Gate 4 no longer contains the synthetic fixture, full CPU recurrence, raw operator launcher or duplicated log-softmax/comparison implementation.

The runner separately builds the default production configuration and the validation configuration, and fails closed on missing inputs, model/config mismatch, non-SM120 hardware, invalid bundles, zero exact-filter matches, ignored tests, fallback or any failed gate. This covers all five requested validation boundaries.

HTTP ABBA

Serving performance was measured separately using:

Triton A1 → FlashInfer B1 → FlashInfer B2 → Triton A2

The comparison used one release binary with a temporary, unmerged benchmark-only backend selector. Startup logs proved the selected backend for every leg. The selector, benchmark harness and result files are not part of the PR.

The matrix contained nine workloads, four legs and 32 requests per case: 36 result files and 1,152/1,152 successful requests with exact token counts.

The 63/64/65 cases are effectively at throughput/TPOT parity and show no 64-token cliff. The 2048/32 cases show small repeated improvements, while 128/128 c4/c8 improved throughput and TPOT in both FlashInfer legs.

TTFT p99 was noisier: 64/32 c1 was +7.39%, and 128/128 c1 was +8.19% due to a single B1 spike that did not repeat in B2. I am therefore not treating either as a stable regression. The c4/c8 improvement direction repeated, but the two FlashInfer samples still showed approximately 7%/6% run-to-run variation, so the exact improvement percentages should not be over-interpreted.

Peak sampled VRAM was effectively unchanged: 25,945 MiB for Triton and 25,992 MiB for FlashInfer, a 47 MiB (~0.18%) difference.

I did another audit and found there are still a few small bits of pollution left in the default build. I’ll go back and clean those up a bit more.

This pass reduces the boundary-cleanup diff by 409 net lines (+263/-672) and addresses the remaining production hygiene and redundant-contract issues:

• GDN build dependencies, build.rs logic, native-prepare CUDA source, raw FFI, and safe wrappers are gated by qwen35. The default kernels build graph contains only cc and pegainfer-build; it does not pull in the GDN serde_json/sha2 build dependencies or compile/link the GDN objects.
• Runtime evidence counters, graph/state fields, record hooks, and evidence APIs are isolated behind the non-default gdn-validation feature. They are absent from the default production object and hot path.
• Large tests remain in child tests.rs modules or integration tests instead of production source files.
• Artifact generation now produces one candidate with fixed artifact basenames. The bundle index, multi-variant plumbing, string-count contracts, self-declared production/serving fields, manifest-provided artifact paths, and redundant source preparation were removed.
• compile_sm120.py now returns only actual source/toolchain/PTX/AOT results. Geometry, dtype, token, and workspace constants remain owned by the parent production contract.
• build.rs independently verifies the candidate schema, target, geometry, dynamic-token contract, ABI/state layout, workspace requirements, and the hash and size of all three linked artifacts.
• Native prepare is limited to the actual Hv32 production specialization. The unused Hv48/generic native path and dynamic geometry FFI were removed.
• The stable C ABI was reduced from 128 to 112 bytes. The duplicated spec/supported API, fixed geometry launch fields, and runtime artifact-size API were removed. Field-offset guards remain, workspace sizing is cached at handle creation, and integer narrowing is explicitly bounded.
• Supported SM120/Hv32 loads remain fail-closed when the candidate is missing or invalid; unsupported SM, geometry, or TP configurations continue to use the existing fallback policy.

Validation

All five required gates passed on SM120 for this exact commit and the newly generated candidate:

1. Real AOT candidate through the production path

   • build.rs → C wrapper → link → load → launch
   • ABI offsets
   • Dynamic T={1,63,64,65,128}
   • Non-symmetric HKV state
   • Alias/separate-state parity
   • Unsupported-Hv48 rejection

2. Native prepare against the CPU oracle

   • T={1,63,64,65,128,2048}
   • Q/K normalization
   • V bit parity
   • Alpha/beta
   • Representative non-finite inputs
   • Sticky status
   • New-chunk clearing

3. Production Auto dispatch against the pinned Qwen3.5-4B HF golden

   • FlashInfer selection
   • HF mean/p99/regret checks
   • Greedy-token parity
   • Successful-launch evidence

4. Real-model continuation

   • T=128 unchunked vs. 64+64 chunked
   • Every linear layer's recurrent and convolution state
   • Last-token logits and token
   • First-decode logits and token

5. Real scheduler/CUDA-graph E2E

   • Graph capture
   • Graph replay
   • Zero eager fallback
   • State-slot copy/reuse/compaction
   • Successful FlashInfer launches

I approved and ran the full CI workflow for the current head. It is not merge-ready yet:

• DCO fails because 24 commits are missing Signed-off-by trailers.
• Qwen3 CUDA Clippy fails on the unused CUresult import in pegainfer-kernels/src/ffi/qwen35.rs:6.
• Qwen3.5 CUDA Clippy fails on format_collect (pegainfer-kernels/build.rs:54), a needless lifetime (:61), and a redundant clone (:183).

Both CUDA compile jobs and the CPU/format/metadata/sim-E2E jobs pass. Please fix the four Clippy findings, sign the commits, and rerun the complete workflow; once it is fully green, we can merge.