optimize mhc

[huangyiqun]

PR Category

Operator

Type of Change

Performance Optimization

Description

optimize mhc ops

Issue

N/A

Progress

• Change is properly reviewed (1 reviewer required, 2 recommended).
• Change is responded to an issue.
• Change is fully covered by a UT.

Performance

Performance

Environment:

- GPU: NVIDIA H20
- Driver: 610.43.02
- PyTorch: 2.13.0+cu130
- Triton: 3.7.1
- vLLM: 0.27.1
- Workload: H=4096, HC=4, BF16 activation/norm, FP32 mix/weight
- Sinkhorn iterations: 20
- Measurement: CUDA Graph steady-state, 16 copies × 200 replays × 11 ABBA rounds
- Metric: median latency; Speedup = baseline / Triton candidate
- Candidate commit: `5c9a2f061`

| Tokens | vLLM fused baseline | FlagGems Triton candidate | Speedup | Latency reduction |
| ---: | ---: | ---: | ---: | ---: |
| 64 | 11.0601 µs | 8.2169 µs | 1.3460x | 25.71% |
| 96 | 13.3748 µs | 9.3943 µs | 1.4237x | 29.76% |
| 128 | 13.4894 µs | 9.6523 µs | 1.3975x | 28.45% |

- Geometric mean speedup: **1.3887x**
- Mean latency reduction: **27.97%**
- Correctness: passed for all tested shapes
- Triton-only source gate: **11 passed**
- The candidate steady-state GPU path is implemented entirely with `@triton.jit`.
- One-time Triton weight packing is excluded from steady-state latency.