Skip to content

vllm.model_executor.warmup.minimax_m3_msa_warmup

_msa_prebuild_kwargs(vllm_config, module)

fmha_sm100.warmup arguments for the kernels this rank can request.

Source code in vllm/model_executor/warmup/minimax_m3_msa_warmup.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
def _msa_prebuild_kwargs(
    vllm_config: "VllmConfig", module: MiniMaxM3SparseAttention
) -> dict[str, Any] | None:
    """``fmha_sm100.warmup`` arguments for the kernels this rank can request."""
    impl = module.impl
    indexer = module.indexer.impl
    msa = (
        impl
        if isinstance(impl, MiniMaxM3SparseMSABaseImpl)
        and module.kv_cache_dtype in _MSA_KV_CACHE_DTYPES
        else None
    )
    attention = msa is not None
    if isinstance(indexer, MiniMaxM3IndexerNvfp4Impl):
        index_cache_dtype = "nvfp4"
    elif isinstance(indexer, MiniMaxM3IndexerMSAImpl):
        index_cache_dtype = module.indexer_kv_dtype
    else:
        index_cache_dtype = None
    if not attention and index_cache_dtype is None:
        return None

    spec_config = vllm_config.speculative_config
    num_spec = spec_config.num_speculative_tokens if spec_config is not None else 0
    q8kv4_decode = attention and uses_q8kv4_decode(
        kv_cache_dtype=module.kv_cache_dtype,
        num_q_heads=module.num_heads,
        num_kv_heads=module.num_kv_heads,
    )
    return dict(
        kv_cache_dtype=module.kv_cache_dtype if attention else "auto",
        index_cache_dtype=index_cache_dtype,
        num_q_heads=module.num_heads,
        num_kv_heads=module.num_kv_heads,
        topk=impl.topk_blocks,
        decode_query_lens=range(1, num_spec + 2),
        num_index_heads=indexer.num_index_heads,
        page_size=impl.block_size,
        decode_backend="q8kv4" if q8kv4_decode else "kv_mode3",
        block_scale_shift=NVFP4_BLOCK_SCALE_SHIFT,
        sparse_decode=msa is not None and msa.use_cutlass_decode,
        sparse_prefill=attention,
        prefill_backend=msa.prefill_backend if msa is not None else "cute_dsl",
        device=torch.device("cuda", torch.accelerator.current_device_index()),
    )