flashrt/fp4-fused-ops

FlashRT fused FP16/BF16-to-NVFP4 producer kernels for transformer and diffuser low-bit paths.

Functions

• sfa_size_bytes
• silu_mul_quantize_fp4_sfa_bf16
• rms_norm_quantize_fp4_sfa_bf16
• rms_norm_fp4_sfa_fp16
• residual_add_rms_norm_fp4_sfa_fp16
• residual_add_rms_norm_fp4_sfa_v2_fp16
• residual_add_rms_norm_mul_fp4_sfa_fp16
• silu_mul_fp4_sfa_fp16
• silu_mul_fp4_sfa_v2_fp16
• silu_mul_mul_fp4_sfa_v2_fp16
• silu_mul_two_fp4_to_fp4
• silu_mul_two_mul_fp4_to_fp4
• dequantize_fp4_sfa_fp16
• quantize_bf16_to_nvfp4_linear
• rms_silu_nvfp4_ndhwc_bf16
• bf16_rms_norm_ncdhw
• bf16_rms_silu_ncdhw
• adaptive_rms_norm_e0m3_fp16
• gated_residual_adaptive_rms_norm_e0m3_fp16
• adaptive_rms_norm_nvfp4_fp16 / adaptive_rms_norm_nvfp4_bf16
• gated_residual_adaptive_rms_norm_nvfp4_fp16 / gated_residual_adaptive_rms_norm_nvfp4_bf16
• canonical aliases ada_rms_norm_quant_nvfp4_swizzled_{fp16,bf16} and gate_res_ada_rms_norm_quant_nvfp4_swizzled_{fp16,bf16}
• gelu_mul_e0m3_fp16
• residual_add_rms_norm_quant_nvfp4_swizzled_bf16
• relu2_quant_nvfp4_swizzled_fp16

This package targets Blackwell sm_110a and sm_120a and uses CUTLASS/CUTE SFA layouts. SM110 artifacts require CUDA 13+.

Example

from kernels import get_kernel
import torch

ops = get_kernel("flashrt/fp4-fused-ops", version=2, trust_remote_code=True)

merged = torch.randn((16, 4096), device="cuda", dtype=torch.float16)
packed, sfa = ops.silu_mul_fp4_sfa_v2_fp16(merged)

# Debug only; normal low-bit pipelines should pass packed/SFA to FP4 GEMM.
bf16_view = ops.dequantize_fp4_sfa_fp16(packed, sfa)

Shape Contract

• CUDA tensors only.
• FP16/BF16 producer inputs use dtype-specific entry points; outputs are uint8 FP4 packed tensors plus uint8 CUTLASS SFA buffers.
• Dimensions must be divisible by 16.
• v1 RMS producer paths support dim <= 2048.
• Larger residual/RMS producer shapes should use residual_add_rms_norm_fp4_sfa_v2_fp16.
• NCDHW BF16 RMS kernels require even C <= 1024; the fused NCDHW-to-NDHWC NVFP4 producer requires C % 128 == 0.