feat(qwen35): add SM120 FlashInfer GDN prefill candidate - #862
feat(qwen35): add SM120 FlashInfer GDN prefill candidate#862qwzx-qwas wants to merge 15 commits into
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 8c88c6c9c4
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| git rev-parse HEAD | ||
| git status --short -- pegainfer-qwen35 |
There was a problem hiding this comment.
When Stage 9 is run from a working checkout, this records the current commit but neither compares it with PEGAINFER_STAGE9_COMMIT nor checks changes outside pegainfer-qwen35. In particular, edits under pegainfer-kernels—which contains the candidate CUDA kernel and FFI—are compiled while remaining invisible in the provenance log, so results can be attributed to a clean commit that did not produce them. Validate the supplied commit against git rev-parse HEAD and reject or capture the complete dirty tree before benchmarking.
Useful? React with 👍 / 👎.
| _require_equal(artifact.get("size_bytes"), len(data), "artifact size") | ||
| _require_equal(artifact.get("sha256"), sha256_bytes(data), "artifact hash") |
There was a problem hiding this comment.
When a generated, downloaded, or copied bundle is checked without --flashinfer-dir as documented, this only verifies that the manifest's hash matches its accompanying bytes. A modified PTX can therefore be accepted after recomputing the artifact and bundle hashes; even source-assisted validation does not bind those instructions to the pinned source. The Rust loader later rejects such a bundle using its hard-coded candidate hash, making the advertised validation produce a false success. Check the pinned release PTX hash here as well.
Useful? React with 👍 / 👎.
| Run source and host-side contract checks without CuTe: | ||
|
|
||
| ```bash | ||
| python3 pegainfer-kernels/tools/flashinfer_gdn/artifact_contract.py verify-source |
There was a problem hiding this comment.
On a fresh checkout whose FlashInfer submodule has not been initialized, this documented first command fails: git -C pegainfer-kernels/third_party/flashinfer rev-parse HEAD walks up to the parent repository and reports its unrelated commit as a FlashInfer SHA mismatch. Initialize/update the pinned submodule before this command or make verify-source detect and initialize the missing gitlink; the documented command was reproduced failing in that context.
AGENTS.md reference: AGENTS.md:L136-L137
Useful? React with 👍 / 👎.
| @@ -0,0 +1,100 @@ | |||
| # FlashInfer GDN SM120 artifact generation | |||
There was a problem hiding this comment.
This commit introduces a large Qwen3.5 backend candidate, artifact contract, accuracy gates, and benchmark workflow, but records it only in a tool-local README; no relevant docs/models/qwen35/ document or docs/index.md route is created or updated. Add the model-line task/decision record so the rationale, validation status, blockers, and next action remain discoverable under the repository's required documentation workflow.
AGENTS.md reference: AGENTS.md:L160-L164
Useful? React with 👍 / 👎.
|
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, Could we reshape the PR around this smaller end-to-end boundary?
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
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 validationThe 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 validationI 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. CorrectnessOn both machines, the Hv32 production candidate passed:
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. PerformanceMachine 1 completed a 15-case same-production-path ABBA matrix covering:
FlashInfer won TTFT p50 in 14/15 cases and throughput in 12/15 cases. The case-level geometric-mean improvements were:
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:
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. |
…fer-gdn-sm120-pr # Conflicts: # pegainfer-qwen35/tests/serving_tp2.rs
|
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 boundaryThe integration now follows this path:
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:
PR scope reductionThe 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:
Serving validationThe 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 validationThe 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. CorrectnessThe Hv32 production candidate passed:
Production tests reported: and successful-launch counters increased as expected, confirming that the tests did not silently fall back to Triton. The latest final gate on commit
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. PerformanceMachine 1 completed a 15-case same-production-path ABBA matrix covering FlashInfer won:
The case-level geometric-mean results were:
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:
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 neededThe remaining policy and scope decisions are:
The current implementation keeps the conservative Triton capability fallback and retains the Hv48 diagnostic path until these decisions are made. |
2 participants
Summary
This PR adds an SM120 FlashInfer GDN prefill candidate for Qwen3.5, including:
The production
prefill_chunk_forward()path remains hard-coded to Triton. This PR does not switch serving traffic to FlashInfer and does not introduce an automatic fallback policy. The two backends share the surrounding embedding, full-attention, MLP, residual, and layer-loop logic; they diverge in backend-owned scratch allocation and the linear-attention prefill operation.Correctness
Validated on RTX 5090 / SM120 with driver 580.126.09 and CUDA 12.8.
The production Qwen3.5-4B geometry, Hq/Hk/Hv/D = 16/16/32/128, passes:
The non-production Hv48 generalization diagnostic passes output checks but has 5/786432 localized final-state tail violations at T=128 relative to the FP64 oracle. Triton has 2 violations for this case. FlashInfer nevertheless has better overall max_abs, mean_abs, and p99_abs state error at T=128.
Patched HKV and unmodified upstream HVK artifacts produce bitwise-identical outputs and states after layout conversion, including the same five violation coordinates. The Hv48 tail therefore originates in the upstream FlashInfer SM120 numerical path rather than the OpenInfer layout patch, TMA indexing, or alias handling.
Performance
Same-machine non-profiled ABBA results:
The non-profiled ABBA results are used for end-to-end latency and throughput. Nsight runs are retained as diagnostic kernel evidence rather than mixed into the latency comparison.
Artifact
Validated artifact SHA-256:
225646b26dab488cdfd64dcf3fe189ba4b7ccaf2ba735eb7b68a47d13db96b68The repository contains the generator, pinned source/toolchain metadata, validation contract, and local-generation documentation. Generated PTX, manifests, bundles, model weights, build outputs, logs, and Nsight reports are not included in this PR.
The frozen HKV patch currently contains one trailing-whitespace line. Changing it alters the pinned source-set hash and requires artifact regeneration and GPU revalidation.
Maintainer decisions requested
This PR intentionally leaves the following policy decisions open:
Refs #691