feat(qwen35): add opt-in SM120 FlashInfer GDN prefill - #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！

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.

I really appreciate the careful review and the detailed code-level analysis. The issues you identified are clear, and I agree that the current gates do not fully cover the production behavior changed by this PR.

I’ll address them as follows:

• Make the locally generated FlashInfer candidate an explicit opt-in. Without a selected bundle, SM120/Hv32/TP1 will retain the existing Triton path; an explicitly supplied but invalid bundle will fail closed.

• Resolve the backend before KV-cache budgeting and reserve the peak scratch/workspace required by the selected backend.

• Remove the complete finite-status path from normal serving, including the device status buffer, atomics, D2H copy, and ctx.sync(), so asynchronous Shared-SM prefill/decode overlap is preserved.

• Extend build.rs to validate the full frozen production contract before linking, including source/toolchain identity, dtypes, tensor views and strides, ABI/layout, workspace, and artifact metadata.

• Keep the project-owned in-place C ABI, native-prepare CPU oracle, real-model continuation check, and existing HF/scheduler E2Es. I’ll remove gdn-validation, validation-only counters and public APIs, duplicate backend-specific gates, the separate-state test ABI, dead helpers, unrelated micro-tests, and process-history comments.

• Restore the scheduler-level chunked-prefill E2E, since the private continuation test does not cover scheduler admission, request events, or the real scheduler chunk boundary.

• Revalidate the exact production path with FlashInfer explicitly selected and Shared-SM enabled, including functional scheduler coverage and GPU timeline/performance evidence. I’ll also retain exact-head upstream-HVK versus patched-HKV differential evidence at the generation boundary without introducing another production backend or test-only runtime API.

Thank you again for taking the time to review this so thoroughly！

This PR integrates an opt-in FlashInfer GDN prefill backend into the actual Qwen3.5 serving path. Triton remains the default, including when a candidate bundle is linked. The candidate supports SM120, TP1 and Hq/Hk/Hv/D = 16/16/32/128, validated with Qwen3.5-4B.

Changes following review

• Kernel ownership and artifact contract: pegainfer-kernels owns pinned generation, artifact validation and the stable in-place C ABI. Generated symbols and TMA layouts stay outside the model layer. Build-time validation checks the frozen source/toolchain, tensor views, dtypes, ABI and final artifact bytes before linking.
• Production dispatch and memory budgeting: The selected backend is checked and loaded before KV allocation, and its peak prefill scratch determines the reservation. Unsupported or unavailable candidates fail explicitly.
• Shared-SM execution: Removed the per-chunk diagnostic status-copy/host-synchronization path while preserving stream/event ordering and resource lifetimes.
• Production/test isolation: Removed validation-only features, counters, the separate-state test ABI and duplicate backend-specific gates. Restored the original three-argument executor API; existing HF backend setup now lives in a private cfg(test) module.
• Scope reduction: Retained the existing HF, real-model continuation, scheduler-level chunked-prefill and CUDA Graph gates. Shared FlashInfer headers retain main’s pin; generation uses an independent checkout.
• Documentation: Added the concise model integration record and its docs/index.md route. Rental scripts, profiling tools, private reports and generated artifacts remain outside the PR.

Validation

The final runtime/test source was validated on RTX 5090 with Qwen3.5-4B: release build, affected-package all-targets Clippy, short/long HF, Shared-SM lifecycle, and real default/explicit-candidate HTTP requests. The candidate HTTP check consumed 20,000 input tokens and generated eight output tokens.

Existing real-bundle, native-prepare oracle, upstream/patched layout, continuation and positive-overlap evidence retains its original source identities.

Earlier same-binary, unprofiled ABBA on aa9ce155 completed 12,288 requests with zero failures. Paired throughput gains were 0.64–2.39%, with workload-dependent tail-latency regressions. These measurements do not justify changing the default backend.

Thank you for the careful review, particularly for identifying the silent-success paths in artifact identity and candidate acceptance.

Following your latest review, I have added three changes:

• The runtime constructor now rejects null, invalid UTF-8, incorrectly sized, and non-hexadecimal artifact identities, and stores only the validated SHA.
• Candidate acceptance entries explicitly select the candidate and verify the identity of the actual model passed to the executor or scheduler.
• HF acceptance now fails when the model revision cannot be resolved or does not match the fixture.

The negative cases exercise the existing constructor and acceptance entries before the positive scenarios. Test setup remains private, and identity fault injection is confined to a single test executable. No production injection hook or additional public testing API was introduced.

During validation, two candidate short-HF runs failed at the same padded-decode position: seq 3 / pos 2 selected token 67, whose HF regret is 0.25 nat against the unchanged 0.20 limit. The third run and subsequent direct reruns passed. The latest main-integrated tree also passed all 19 acceptance stages, including short/long HF, real build/load/HTTP, layout and prepare checks, continuation, scheduler lifecycle, and actual Shared-SM overlap. These passes do not resolve the earlier failure.

The isolated investigations narrowed the issue. Fixed GEMM configurations reproduced the candidate failure with the two sensitive projection split-K settings at 9/9, while Triton at 9/9 and the candidate at 16/16 passed. Real-input GDN replay found prefill-state differences. With identical decode-entry state, changing the controlled GEMM configuration still changed the outcome; the first captured differences appeared in the shared decode b/a projections. Both backends use the same decode operators.

My strongest suspicion is an interaction between prefill-state differences, shared split-K rounding, and final BF16 logit materialization. For the failing hidden vector, FP64 projection gives logits of 11.14099662 and 11.11763134 for tokens 692 and 67, respectively. Both round to BF16 11.125, matching the GPU result, and the sampler selects 67. The observed GEMM recipes also store split-K partial results in BF16 despite requesting FP32 computation. This is a concrete precision lead, but its individual contribution has not been isolated, and a unique FlashInfer kernel defect has not been established.

No production fix has been implemented for this numerical issue. I have not hardcoded the passing split-K configuration, changed tie-breaking, relaxed the HF threshold, or promoted diagnostic helpers into this PR.

Triton remains the default backend.

Oh, thanks for testing all these options—let’s go with down, get the accuracy checks passing, and wrap this up.