feat(qwen35): add SM120 FlashInfer GDN prefill candidate

[qwzx-qwas]

Summary

This PR adds an SM120 FlashInfer GDN prefill candidate for Qwen3.5, including:

• native CUDA preparation kernels and Rust FFI;
• a validated FlashInfer PTX artifact generation and loading pipeline;
• pinned source/toolchain contracts and an HKV state-layout patch;
• explicit FlashInfer operator, HF golden, chunked-prefill, scheduler/CUDA Graph, and benchmark entry points;
• an ABBA benchmark harness for comparison with the existing Triton backend.

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:

• operator tests for T = 1, 2, 63, 64, 65, 127, and 128;
• short and long HF golden gates;
• chunked versus unchunked prefill;
• scheduler and CUDA Graph integration tests.

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:

• T=128, concurrency=8: FlashInfer 378.75 tok/s vs Triton 360.72 tok/s, approximately +5.0%.
• T=2048, concurrency=1: FlashInfer 54.15 tok/s vs Triton 52.82 tok/s, approximately +2.52%.
• T=2048 backend-owned scratch: FlashInfer 50,877,716 bytes including runtime workspace vs Triton 193,462,272 bytes, approximately 73.7% lower.
• FlashInfer runtime workspace: 21,760 bytes.

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:

225646b26dab488cdfd64dcf3fe189ba4b7ccaf2ba735eb7b68a47d13db96b68

The 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:

1. whether SM120 production serving should switch to FlashInfer;
2. whether Triton should remain as an internal rollback path;
3. whether artifacts should be prebuilt, locally generated, or support both;
4. whether the Hv48 upstream numerical tail should become a permanent diagnostic gate;
5. whether the candidate is ready for release integration.

Refs #691

[chatgpt-codex-connector]

💡 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".

[chatgpt-codex-connector]

Validate the benchmark's actual source tree

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 👍 / 👎.

[chatgpt-codex-connector]

Pin the standalone validator to the candidate PTX

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 👍 / 👎.

[chatgpt-codex-connector]

Initialize the pinned submodule before verification

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 👍 / 👎.

[chatgpt-codex-connector]

Add the required model-line decision record

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 👍 / 👎.

[xiaguan]

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.

[qwzx-qwas]

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