feat(phyai-kernel): add Thor NVFP4 GeGLU kernel for pi0.5

[Aharrypotter]

Summary

This Draft PR adds an opt-in SM110 NVFP4 Gate+Up+GeGLU kernel for the fixed
pi0.5 language MLP shapes. The kernel fuses the merged gate/up projection,
GeGLU fold, and compact NVFP4 output store. Its output can feed the existing
NVFP4 down projection without a BF16 intermediate or a separate quantization
launch.

The compact epilogue is implemented in phyAI against CUTLASS callback
interfaces. The PR does not vendor, include, or depend on FlashRT source.

This is a low-level kernel API. It does not change pi0.5 model dispatch, weight
loading, conversion, or the default DenseMLP path.

Review map

1. phyai_kernel/cuda/pi05_thor_nvfp4_geglu.py defines the public contract,
   fail-closed validation, caller-owned buffers, SM110 guard, and JIT build.
2. cuda/csrc/pi05_thor_nvfp4_geglu.cuh contains the fixed
   128x256x256 tile, 1x2x1 cluster schedule, and TVM-FFI boundary.
3. cuda/csrc/pi05_thor_nvfp4_geglu_epilogue.cuh implements the phyAI-owned
   GeGLU fold, per-block scale generation, packed E2M1 store, and CUTLASS
   fusion callback.
4. tests/test_pi05_thor_nvfp4_geglu.py covers the shape and storage contract,
   numerical comparison, failure paths, and CUDA Graph capture.
5. benchmark/bench_pi05_thor_nvfp4_geglu.py defines the matched staged
   baseline and raw-sample timing protocol.

Supported contract

Surface	Contract
Device	NVIDIA SM110
Shapes	M in {784, 816, 880, 968}, K=2048, merged N=32768, hidden H=16384
Values	Packed NVFP4 E2M1 activation and weight values
Scales	FP8 E4M3 in CUTLASS/FlashInfer 128x4 layout, global scale 1
Weight layout	Pairwise-interleaved gate/up rows: gate0, up0, gate1, up1, ...
Output	Packed NVFP4 hidden activation and scale factors for the down projection
Ownership	Caller supplies output and workspace buffers
Fallback	No fallback is hidden inside the API. Unsupported inputs fail before launch
Dispatch	Explicit opt-in only. Existing model dispatch remains unchanged

Numerical checks compare the consumed down-projection result with the staged
NVFP4 path. The frozen gate is relative-L2 below 0.05 and cosine above 0.999.

Validation

Current product commit: e44875cce92374002339164b123624cb9295fc35.

Gate	Status	Evidence boundary
Repository pre-commit checks	PASS	clang-format, codespell, ruff-format, and diff checks on the seven-file product diff
Fresh Thor JIT build	PASS	Current source built native SM110 and SM110a images on Jetson Thor
Focused Thor tests	PASS, 4 tests	Host rejection paths, numerical tolerance through the NVFP4 down projection, launch, and CUDA Graph capture/replay
Code generation and resources	PASS with changed code	Both images use 255 registers and 1024 bytes static shared memory. The SM110a image has no stack or LDL instruction
Accepted performance matrix	PENDING	Device code changed after the source cleanup, so earlier timing is not inherited
Local macOS target pytest	BLOCKED	The host environment has no Triton installation. Target correctness is covered by the fresh Thor run
Upstream CI	PENDING	Draft PR gate

The current source contains no FlashRT vendored file or include. The first
published Draft revision did, and its performance and model screens are no
longer used as current-source evidence.

Performance evidence

Evidence boundary: performance qualification is pending an uncontended Thor
window. Two unrelated GPU-resident policy servers are active on the shared
device, so their timing is excluded from the PR claim.

Setup: the checked-in benchmark freezes the intended run contract. Timing
is reported in ms.

• CUDA Graph replay timed with CUDA events.
• 25 warmups, 30 raw samples per shape, and 20 replays per sample.
• Alternating candidate and baseline order with no outlier filtering.
• Baseline: FlashInfer NVFP4 merged GEMM, PyTorch GeGLU, then FlashInfer NVFP4
  quantization.
• Matched endpoints: both paths start with the same packed NVFP4 activation and
  interleaved weight, and end with packed NVFP4 hidden activation and scales.

Results in milliseconds, per-shape speedups, the primary geometric mean, and
the worst row will be added after the current source passes that full matrix.
The 18-layer checkpoint workset and temporary model integration screen will
also be rerun before this Draft is promoted.

Files changed

• Public API and exports: three Python files.
• CUDA implementation: the fixed GEMM header and one phyAI-owned epilogue
  header.
• Validation: one focused test and one benchmark.

The PR contains 7 changed files with 1,063 additions and 1 deletion.
Experimental harnesses, generated artifacts, model weights, rejected schedules,
and model integration patches are excluded.

Known limits

• The fast path supports SM110 and the four observed pi0.5 token counts only.
• CUTLASS 4.4.2 or newer is required. The validated environment selected the
  copy bundled with FlashInfer.
• Callers must provide pairwise-interleaved weights, global scale 1, and
  caller-owned workspace and output buffers.
• No production model dispatch, automatic fallback, or weight preparation is
  included. Unsupported contracts fail closed and the existing MLP path is
  unchanged.
• Numerical validation establishes parity with the staged NVFP4 recipe. It is
  not calibrated model-quality or task-quality evidence.
• The PR remains Draft until the current source completes uncontended kernel,
  18-layer MLP, and temporary model integration performance screens.

[coderabbitai]

Important

Review skipped

Draft detected.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro Plus

Run ID: 7317663e-015f-4d7d-9b7a-4b560882fc15

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}