vllm.v1.attention.backends.mla.amx_mla ¶

AMX-only, high-performance MLA backend for DeepSeek V2/V3/R1 on CPU.

Built on the AMX decode/extend/bmm kernels vendored under csrc/cpu/sgl-kernels/, plugged into vLLM's MLACommonBackend/ MLACommonImpl abstraction the same way every other concrete MLA backend (TritonMLAImpl, etc.) does.

This is a separate backend from the reference CPUMLABackend (vllm/v1/attention/backends/mla/cpu_mla.py): that one targets every CPU (any dtype, block_size=16, mla_decode_kvcache decode kernel + inherited SDPA-style prefill) as a functional/CI reference, not performance. This backend instead requires AMX (bf16 only, block_size a multiple of 32) and is selected by the platform layer in preference to the reference backend whenever the host supports it -- see CpuPlatform.get_attn_backend_cls.

Two points where this backend differs structurally from the GPU backends, both explained in the CPU MLA design plan:

• forward_mha is fully overridden (not inherited from MLACommonBaseImpl): the GPU "compute-friendly" prefill path depends on a pluggable MLAPrefillBackend and CUDA-only chunked-context gather ops, neither of which exist on CPU. Instead, this attends directly in latent-MQA space via the extend_attention_cpu kernel, which handles cached-prefix continuation and fresh prefill in one causal pass (the KV cache already contains the new tokens by the time this runs, since do_kv_cache_update executes first).
• Weight absorption for the prefill path is done with this impl's own VNNI-packed copies of W_UK/W_UV (computed in process_weights_after_loading, using bmm_cpu), rather than reading the generic layer.W_UK_T/layer.W_UV/layer._v_up_proj the way GPU backends do, since forward_mha's abstract signature has no layer parameter (only forward_mqa does) and that signature is left untouched. The decode path needs no such packing: MLAAttention.forward_impl already absorbs/de-absorbs Q around the forward_mqa call using the generic (unpacked) layer.W_UK_T/layer._v_up_proj, so forward_mqa here only has to invoke the decode kernel.

64
65
66
67
68
69
70
71

def _compute_num_kv_splits(max_seq_len: int, num_threads: int) -> int:
    """Mirrors TritonMLAImpl's _compute_num_kv_splits, using the CPU thread
    count in place of SM count."""
    ideal_splits = 1
    while ideal_splits < max(1, max_seq_len // _MIN_WORK_PER_SPLIT):
        ideal_splits *= 2
    max_splits = num_threads * _SPLIT_OCCUPANCY_MULTIPLIER
    return min(ideal_splits, max_splits)

74
75
76
77
78
79
80
81
82
83

def _expand_block_table(block_table: torch.Tensor, block_size: int) -> torch.Tensor:
    """Adapter: vLLM's block-table paging -> the flat per-(request, position)
    physical row index the decode/extend kernels expect. Called once per step
    from ``AMXMLAMetadataBuilder.build``, not per layer.
    """
    offsets = torch.arange(
        block_size, device=block_table.device, dtype=block_table.dtype
    )
    flat = block_table.unsqueeze(-1) * block_size + offsets
    return flat.reshape(block_table.size(0), -1).contiguous()