FlashInfer FP8 MLA decode exampleΒΆ

The wrapper repository contains the runnable end-to-end example referenced by the FlashInfer wrapper guide:

 1"""Run fused FP8 MLA RoPE quantization followed by sparse decode on MUSA."""
 2
 3from __future__ import annotations
 4
 5import math
 6
 7import torch
 8
 9import flashinfer
10
11
12def _make_cos_sin_cache(tokens: int) -> torch.Tensor:
13    positions = torch.arange(tokens, device="musa", dtype=torch.float32).unsqueeze(1)
14    dimensions = torch.arange(0, 64, 2, device="musa", dtype=torch.float32)
15    inv_freq = torch.pow(10000.0, -dimensions / 64.0)
16    angles = positions * inv_freq.unsqueeze(0)
17    return torch.cat((angles.cos(), angles.sin()), dim=-1)
18
19
20def main() -> None:
21    torch.manual_seed(20260722)
22    tokens, heads, topk = 64, 64, 64
23    q_quant_scale = 3.0
24    kv_quant_scale = 1.5
25    attention_scale = 1.0 / math.sqrt(576)
26
27    q_rope = torch.randn(tokens, heads, 64, device="musa", dtype=torch.bfloat16) * 0.1
28    q_nope = torch.randn(tokens, heads, 512, device="musa", dtype=torch.bfloat16) * 0.1
29    k_rope = torch.randn(tokens, 64, device="musa", dtype=torch.bfloat16) * 0.1
30    k_nope = torch.randn(tokens, 512, device="musa", dtype=torch.bfloat16) * 0.1
31    cos_sin_cache = _make_cos_sin_cache(tokens)
32    pos_ids = torch.arange(tokens, device="musa", dtype=torch.int32)
33
34    q_fp8 = torch.empty(tokens, heads, 576, device="musa", dtype=torch.float8_e4m3fn)
35    kv_fp8 = torch.empty(tokens, 576, device="musa", dtype=torch.float8_e4m3fn)
36
37    flashinfer.rope.mla_rope_quantize_fp8(
38        q_rope=q_rope,
39        k_rope=k_rope,
40        q_nope=q_nope,
41        k_nope=k_nope,
42        cos_sin_cache=cos_sin_cache,
43        pos_ids=pos_ids,
44        is_neox=True,
45        quantize_dtype=torch.float8_e4m3fn,
46        quant_scale_q=q_quant_scale,
47        quant_scale_kv=kv_quant_scale,
48        q_rope_out=q_fp8[..., 512:],
49        k_rope_out=kv_fp8[..., 512:],
50        q_nope_out=q_fp8[..., :512],
51        k_nope_out=kv_fp8[..., :512],
52    )
53
54    query = q_fp8[:1].view(1, 1, heads, 576)
55    kv_cache = kv_fp8.view(1, 1, tokens, 576)
56    block_tables = torch.arange(topk, device="musa", dtype=torch.int32).view(1, 1, topk)
57    seq_lens = torch.tensor([topk], device="musa", dtype=torch.int32)
58    # workspace_buffer is retained only for FlashInfer call compatibility.
59    # Prepare explicit metadata once while seq_lens and top-k stay unchanged.
60    metadata = flashinfer.decode.get_batch_decode_metadata_mla(
61        query=query,
62        seq_lens=seq_lens,
63        sparse_mla_top_k=topk,
64    )
65
66    q_descale = 1.0 / q_quant_scale
67    kv_descale = 1.0 / kv_quant_scale
68    out, lse = flashinfer.decode.trtllm_batch_decode_with_kv_cache_mla(
69        query=query,
70        kv_cache=kv_cache,
71        workspace_buffer=None,
72        qk_nope_head_dim=128,
73        kv_lora_rank=512,
74        qk_rope_head_dim=64,
75        block_tables=block_tables,
76        seq_lens=seq_lens,
77        max_seq_len=tokens,
78        sparse_mla_top_k=topk,
79        bmm1_scale=attention_scale * q_descale * kv_descale,
80        bmm2_scale=kv_descale,
81        backend="trtllm-gen",
82        skip_softmax_threshold_scale_factor=None,
83        return_lse=True,
84        metadata=metadata,
85    )
86    torch.musa.synchronize()
87
88    assert out.shape == (1, 1, heads, 512)
89    assert out.dtype == torch.bfloat16
90    assert lse.shape == (1, heads)
91    print(f"query: {tuple(query.shape)} {query.dtype}")
92    print(f"kv_cache: {tuple(kv_cache.shape)} {kv_cache.dtype}")
93    print(f"out: {tuple(out.shape)} {out.dtype}; lse: {tuple(lse.shape)}")
94
95
96if __name__ == "__main__":
97    main()