Add fused epilogue support to preshuffle GEMM: bias + ReLU/SiLU/GeLU

[andyluo7]

Summary

Add fused epilogue support to the preshuffle GEMM kernel (preshuffle_gemm.py), enabling bias addition and activation functions (ReLU, SiLU, GeLU) to be computed directly in the GEMM output store loop (body_row). This eliminates separate memory-bound epilogue kernels and their associated kernel launch overhead.

Motivation

In LLM inference (e.g., vLLM, SGLang), GEMM operations are frequently followed by bias addition and activation functions. These post-GEMM operations are typically executed as separate kernels, adding:

• Kernel launch overhead (~3-5μs per launch)
• Extra memory traffic (read C, write C for each epilogue op)

By fusing these into the GEMM body_row, the output values are processed while still in registers — zero additional memory traffic.

Changes

• kernels/preshuffle_gemm.py: Added epilogue parameter to compile_preshuffle_gemm_a8() supporting "none", "bias", "bias_relu", "bias_silu", "bias_gelu". Added arg_bias tensor argument to the kernel function and launcher. Bias is loaded via buffer_ops.buffer_load and activations are computed in-place before the output store.

• python/flydsl/_version.py: Bumped to dev version.

Supported Epilogues

Epilogue	Operation	Use Case
none	No change (default)	Standard GEMM
bias	C = GEMM(A,B) + bias	Linear layers with bias
bias_relu	C = ReLU(GEMM(A,B) + bias)	MLP layers
bias_silu	C = SiLU(GEMM(A,B) + bias)	FFN gate (LLaMA, DeepSeek)
bias_gelu	C = GeLU(GEMM(A,B) + bias)	GPT-style FFN

Performance

Benchmarked on MI300X (gfx942) and MI355X (gfx950) with real LLM model shapes:

MI300X — Fused SiLU vs hipBLAS + separate bias + SiLU

Shape	hipBLAS+epilogue	FlyDSL fused	Speedup
MoE expert down (1×7168×2048)	0.050ms	0.015ms	3.33x
Average across 6 shapes	—	—	2.73x

MI355X — Fused SiLU

Shape	hipBLAS+epilogue	FlyDSL fused	Speedup
Average across shapes	—	—	1.28x
Peak (MoE expert down)	—	—	1.80x

End-to-End vLLM Serving (DeepSeek-R1-0528, MoE)

Platform	TP	Baseline	+ FlyDSL	Improvement
MI300X	8	64.82 tok/s	76.32 tok/s	+17.7%

API

from flydsl.kernels.preshuffle_gemm import compile_preshuffle_gemm_a8

# Standard GEMM (no change to existing usage)
kernel = compile_preshuffle_gemm_a8(M, N, K, ...)

# Fused GEMM + bias + SiLU
kernel = compile_preshuffle_gemm_a8(M, N, K, ..., epilogue="bias_silu")

# Caller passes arg_bias tensor (shape [N], bf16/fp16) to the kernel
kernel(C, A, B, scale_a, scale_b, bias, M, N, stream=stream)

Backward Compatibility

• Default epilogue="none" preserves existing behavior
• arg_bias is always present in the kernel signature but unused when epilogue="none" — callers can pass a dummy tensor
• No changes to other kernels or the compiler infrastructure

Testing

• Correctness validated against PyTorch reference (F.linear + bias + act) on both gfx942 and gfx950
• E2E validated through vLLM serving with DeepSeek-R1-0528 (671B MoE)

[coderfeli]

@andyluo7 CI failed.

[andyluo7]

Fixed in 1b04c35 — the test's _gemm_args() and _w4_args() were missing the arg_bias parameter that was added to launch_gemm in the fused epilogue commit. Added dummy bias tensors to both. CI should pass now.