Enhanced FP8 Support with Significant Performance Optimizations

[justinSmileDate]

Description

This PR ports and enhances FP8 support from #82 to the latest branch, delivering substantial performance improvements through optimized computational patterns and memory access strategies.

Key Changes

Compared to Original FlashMLA:

• WGMMA FP8 Integration: Leverage WGMMA instructions for efficient FP8 matrix operations
• FP8 Data Types: Utilize FP8 data types for both Q and KV tensors
• Memory Optimization:
  • Save shared memory by allocating shared memory for sP1
  • Remove retrieve_rP_for_SP(sQ(8)) operations
  • Eliminate rQ(8)*sK RS WGMMA operations

Improvements over PR #82:

• Code Structure: Refactored for better integration with the latest branch
• Boundary Processing Optimization:
  • Operate directly at block index level instead of sequence length level
  • Eliminate redundant ceil_div calculations
  • Reduce division operations for improved computational efficiency
• Enhanced Memory Access:
  • Minimize intermediate variable calculations
  • Optimize register utilization
  • Improve performance across diverse workload patterns

Performance Benchmarks

Test Configuration:

• Batch Size: 128
• Sequence Lengths: 4096, 8192, 16384
• Heads: h_q=16, 32, 64, 128; h_kv=1
• Dimensions: d=576, dv=512
• Causal Mask: Enabled
• Variable Length: Both True and False
• Hardware: H20

Key Performance Highlights:

Metric	FlashMLA (bfloat16)	PR #82 (FP8)	This PR (FP8)	Improvement vs PR #82
Best Latency	0.553 ms	0.405 ms	0.345 ms	+15% faster
Best TFLOPS	140 TFLOPS	214 TFLOPS	234 TFLOPS	+9% higher
Best Bandwidth	1165 GB/s	802 GB/s	972 GB/s	+21% higher
Varlen Performance	0.577 ms	0.579 ms	0.353 ms	+39% faster

Representative Performance Comparison:

Configuration: b=128, s_q=1, mean_sk=4096, h_q=16

Implementation	varlen	Latency	TFLOPS	Bandwidth
FlashMLA	False	0.553 ms	33	1100 GB/s
PR #82	False	0.405 ms	45	752 GB/s
This PR	False	0.345 ms	53	881 GB/s
FlashMLA	True	0.577 ms	32	1060 GB/s
PR #82	True	0.579 ms	32	528 GB/s
This PR	True	0.353 ms	52	866 GB/s

Configuration: b=128, s_q=2, mean_sk=4096, h_q=32

Implementation	varlen	Latency	TFLOPS	Bandwidth
FlashMLA	False	0.561 ms	130	1108 GB/s
PR #82	False	0.412 ms	177	755 GB/s
This PR	False	0.351 ms	208	885 GB/s

Scalability Analysis:

Across Head Counts (mean_sk=4096, s_q=1):

Heads	PR #82 Latency	This PR Latency	Improvement
16	0.405 ms	0.345 ms	+15%
32	0.406 ms	0.347 ms	+14%
64	0.410 ms	0.350 ms	+15%
128	0.790 ms	0.665 ms	+16%

Across Sequence Lengths (h_q=16, s_q=1):

SeqLen	PR #82 Latency	This PR Latency	Improvement
4096	0.405 ms	0.345 ms	+15%
8192	0.771 ms	0.645 ms	+16%
16384	1.509 ms	1.245 ms	+17%

Performance Improvements Summary

vs Original FlashMLA (bfloat16):

• Latency Reduction: Up to 38% (0.553ms → 0.345ms)
• TFLOPS Improvement: Up to 67% (140 → 234 TFLOPS)
• Consistent gains across all head counts and sequence lengths

vs PR #82 (FP8):

• Latency Reduction: 14-17% across all configurations
• TFLOPS Improvement: 9-18% across various workloads
• Significant varlen improvement: 39% latency reduction
• Better bandwidth utilization: Up to 21% improvement

Technical Advantages

1. Computational Efficiency:

   • Reduced arithmetic operations through optimized boundary handling
   • Better utilization of GPU compute resources

2. Memory Optimization:

   • Reduced shared memory footprint
   • Improved memory access patterns
   • Better register allocation

3. Scalability:

   • Consistent performance improvements across different model configurations
   • Excellent scaling with increasing head counts and sequence lengths

Testing & Validation

python test/test_fp8.py --dtype float8_e4m3fn

Usage

export ENABLE_SWAPAB=1
python3 -m sglang.launch_server XXX --quantization fp8 --kv-cache-dtype fp8_e4m3

This PR delivers substantial performance improvements while maintaining full numerical correctness and compatibility with existing FP8 functionality.

[akhoroshev]

Are any sglang updates needed for testing?

[TianyuZhang1214]

Excellent work!

I'm a member of the LLM Inference Team at AntGroup Theta, and I also serve as the maintainer for this AntGroup-H20-PR.

We are interested in establishing a long-term connection. Would you be comfortable sharing your contact information (such as WeChat) by sending it to my email?

My email is: moyun.zty@antgroup.com

[llc-kc]

Look forward to this PR being merged.

[llc-kc]

@justinSmileDate Does sglang need some adaptations for this PR? thanks

[justinSmileDate]

Excellent work!

I'm a member of the LLM Inference Team at AntGroup Theta, and I also serve as the maintainer for this AntGroup-H20-PR.

We are interested in establishing a long-term connection. Would you be comfortable sharing your contact information (such as WeChat) by sending it to my email?

My email is: moyun.zty@antgroup.com

I apologize for the delay in responding due to my busy work. I have sent the relevant information to your email address. Please check your inbox.

[justinSmileDate]

Are any sglang updates needed for testing?

Actually, only the FlashMLA repository needs to be replaced, and there is no need to adjust sglang. The version of sglang I'm using is v0.5.3.

[justinSmileDate]

@justinSmileDate Does sglang need some adaptations for this PR? thanks

No, there is no need to modify sglang. The version of sglang I'm using is v0.5.3.