[leejet]

Summary

• Reduce intermediate memory traffic by fusing RMSNorm, modulation, residual operations, and SwiGLU/scale with BF16 GEMM operand packing.
• Share fusion implementations across CUDA, HIP, and MUSA, retaining FP32 accumulation and GEMM outputs.
• Add a CUTLASS GEMM path for eligible matrix shapes on NVIDIA SM80+, with native BLAS fallback.

Related Issue / Discussion

N/A

Additional Information

Hardware: NVIDIA GeForce RTX 4090 24 GB, NVIDIA driver 596.36.

DiT files:

Model	DiT file
Z-Image Turbo	z_image_turbo_bf16.safetensors

Sampling throughput in it/s; higher is better:

Model	Resolution	sd.cpp before	sd.cpp after
Z-Image Turbo	512×512	9.02	10.14
Z-Image Turbo	1024×1024	2.20	2.42

Configuration:

• CUDA backend, batch 1, CFG 1, Euler, discrete scheduler, seed 42.
• 40 steps; throughput calculated from mean step time after excluding the first 5 steps.
• Same BF16 DiT weights, qwen_3_4b.safetensors text encoder, and ae.sft VAE.
• Prompt: a lovely cat

Checklist

• I have read and confirmed this PR follows the contribution guidelines.

Summary by CodeRabbit

• Performance
  • Added optimized BF16 matrix multiplication and fused SwiGLU operations for compatible GPU configurations.
  • Added a fused RMS normalization and affine-residual operation for supported CUDA configurations.
  • These optimizations are used only when the device and operation meet compatibility requirements.

[coderabbitai]

No actionable comments were generated in the recent review.

Walkthrough

The CUDA backend adds CUTLASS BF16 GEMM support, a fused RMS-affine kernel, and fused BF16 SwiGLU projection execution. It adds graph matching, execution, build configuration, and registration for these operations.

Changes: CUDA BF16 and RMS-affine operations

Layer / File(s)	Summary
CUTLASS BF16 GEMM support src/core/cuda/bf16_gemm.h, src/core/cuda/bf16_gemm.cu, CMakeLists.txt	Adds the BF16 GEMM result type and implementation. CUTLASS build configuration includes the new GEMM source and compile settings.
RMS-affine fused kernel src/core/cuda/rms_affine.h, src/core/cuda/rms_affine.cu, src/core/cuda/CMakeLists.txt	Adds RMS-affine graph matching and optimization, kernel execution, BF16 packing, and operation registration.
BF16 fused projections src/core/cuda/bf16.h, src/core/cuda/bf16.cu, src/core/cuda/registration.cu, src/core/cuda/CMakeLists.txt	Adds BF16 fused graph matching and execution. The implementation can pack an RMS-affine input and dispatch GEMMs through CUTLASS or platform BLAS. Extension registration adds the BF16 and RMS-affine operations.

Estimated code review effort: 3 (Moderate) | ~25 minutes

Change: Feature

Sequence Diagram(s)

sequenceDiagram
  participant GraphMatcher
  participant BF16Execution
  participant RMSAffinePack
  participant GEMMDispatch
  participant CUTLASSOrBLAS
  GraphMatcher->>BF16Execution: identify supported fused projection pattern
  BF16Execution->>RMSAffinePack: pack RMS-affine input when matched
  BF16Execution->>GEMMDispatch: run projection GEMMs
  GEMMDispatch->>CUTLASSOrBLAS: execute BF16 GEMM

Loading

Merge Risk: Minimal · up to 21f6d

No confirmed merge-blocking issue remains. MUSA compatibility has not been verified, so a supported MUSA build is advisable before release.

Pre-merge checks: 4 · 1

Failed checks: 1 warning

Check name	Status	Explanation	Resolution
Docstring Coverage		Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentat…	Write docstrings for the functions missing them to satisfy the coverage threshold.

Passed checks: 4 passed

Check name	Status	Explanation
Title check		The title clearly summarizes the main changes: fusing Z-Image GPU operations and optimizing BF16 GEMM.
Description check		The description includes the required Summary, Related Issue / Discussion, Additional Information, and Checklist sections. It explains the changes, provides benchmark details, and confirms the contrib…
Linked Issues check		Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check		Check skipped because no linked issues were found for this pull request.

Full details: Description check

Explanation

The description includes the required Summary, Related Issue / Discussion, Additional Information, and Checklist sections. It explains the changes, provides benchmark details, and confirms the contribution guidelines checklist.

Full details: Docstring Coverage

Explanation

Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 21 functions across 7 files.

Resolution

Write docstrings for the functions missing them to satisfy the coverage threshold.

Finishing Touches

Generate docstrings

• Commit to this branch
• Create a new PR

Generate unit tests (beta)

• Commit to this branch
• Create a new PR

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Tip

🎯 Customize your risk assessments

Guide how CodeRabbit assesses merge risk and security architecture risk. Highlight sensitive changes and explain your security boundaries using Merge Risk Instructions and Security Risk Instructions in your review settings. Security architecture risk assessments require the Advanced plan or above.

✅ Automatic approvals, with your rules

Let CodeRabbit approve pull requests once required threads are resolved, pre-merge checks pass, and your risk limits are met. Enable Automatic approvals in review settings, with overrides for teams or file paths.

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

CodeRabbit Verify

• Run CodeRabbit Verify to test the changes in this pull request.

Recent review info

Run configuration

• Source mode: Auto
• Repository: Unavailable
• Branch: Unavailable
• File: No configuration file from this repository applied
• Commit: Unavailable
• Layers: Repository UI
• Review profile: CHILL
• Plan: Advanced
• Run ID: 32110628-312c-446c-9f99-8b79ac92cd3b

Commits

Reviewing files that changed from the base of the PR and between e3a5649 and 21f6d43.

Files selected for processing: 1

• src/core/cuda/bf16.cu

Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 7 remain after this review.