ZYFA: Fast and Accurate FP8 Attention on AMD MI355X
Zyphra presents ZYFA, FP8 and BF16 prefill attention kernels for AMD Instinct MI355X GPUs.
Read the post
ZHA: Zyphra Hierarchical Attention
We introduce Zyphra Hierarchical Attention (ZHA), a sparse attention method that constructs and searches through a hierarchical tree to determine the relevant set of tokens for each query. ZHA achieves subquadratic complexity with suitably chosen top-p thresholds for routing, outperforming existing sparse-attention methods which maintain a quadratic indexer. We demonstrate that we can adapt existing pretrained models to ZHA through a distillation phase using only a few billion tokens, and that ZHA maintains performance on challenging long-context reasoning, agentic coding, and retrieval benchmarks at high levels of sparsity. Our performance model shows that ZHA dramatically outperforms sparse attention methods in speed and computational cost at ultra-long contexts.
Zyphra Inference Wave 2: Faster MoE and Long-Context Serving on MI355X
Zyphra presents the second wave of inference optimizations in the Zyphra Optimization Engine (ZOE) for AMD Instinct MI355X GPUs. We measure the empirical compute and memory-bandwidth limits of the MI355X, write custom HIP kernels for mixture-of-experts layers and sparse attention in our production forks of vLLM and SGLang, and evaluate MXFP4 quantization and speculative decoding. ZOE speeds up GLM-5.3 time to first token at 512K tokens of context by 2.8x with no prefix-cache hits and by 3.6x at a 90% hit rate, and it raises DeepSeek-V4-Pro single-user decode throughput at 1M tokens of context by 1.40x.
Expert Coupling in MoE Pretraining
Mixture-of-Experts (MoE) models scale model capacity efficiently by routing each token to only a small subset of experts. But at large scale, the cost of moving tokens between GPUs can become a major bottleneck, accounting for as much as 60.4% of total training step time in our experiments. In this work, we show that MoE routing contains substantial structure that can be exploited to reduce this communication. Experts selected by the same token are strongly correlated within a layer, while a token’s expert choices in one layer predict where it will be routed in later layers. We use these two properties to increase token–expert locality in two complementary ways: correlated expert placement co-locates experts that are frequently selected together, while token shuffling proactively moves tokens to the GPUs likely to contain their next experts. Neither technique changes the model architecture or its routing decisions. Moreover, we find that these routing patterns are cheap to estimate and emerge early enough in training to remain useful over long stretches of pretraining. In Megatron-LM experiments on MI300X GPUs, our methods reduce all-to-all communication time by up to 2.63× and improve end-to-end training step time by up to 1.41×.