Skip to content

vllm.models.muse_spark.nvidia.ops.mnnvl_fusions

Workspaces for FlashInfer MNNVL collectives fused into Muse Spark boundaries.

Functions:

_create_on_all_tp_ranks(name, create)

Collectively create a workspace; every TP rank gets one or none.

Source code in vllm/models/muse_spark/nvidia/ops/mnnvl_fusions.py 
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def _create_on_all_tp_ranks(name: str, create: Callable[[], Any]) -> Any | None:
    """Collectively create a workspace; every TP rank gets one or none."""
    tp_group = get_tp_group()
    eligible = (
        current_platform.is_cuda()
        and current_platform.is_device_capability_family(100)
        and tp_group.world_size > 1
    )
    if eligible:
        try:
            from flashinfer.comm.mnnvl import is_multicast_supported

            eligible = is_multicast_supported(torch.accelerator.current_device_index())
        except ImportError:
            eligible = False
    if not _all_tp_ranks(eligible):
        logger.info_once("Muse Spark fused %s is unavailable.", name)
        return None
    try:
        workspace = create()
    except (ImportError, RuntimeError, ValueError) as exc:
        logger.warning("Muse Spark fused %s failed: %s", name, exc)
        workspace = None
    if not _all_tp_ranks(workspace is not None):
        return None
    logger.info_once("Muse Spark fused %s enabled (TP%d).", name, tp_group.world_size)
    return workspace

maybe_create_all_gather_dual_norm(hidden_size, max_tokens, *, enable_pdl)

FlashInfer MNNVLCuteDSLAllGatherDualNormWorkspace for the TP group.

Source code in vllm/models/muse_spark/nvidia/ops/mnnvl_fusions.py 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def maybe_create_all_gather_dual_norm(
    hidden_size: int, max_tokens: int, *, enable_pdl: bool
) -> Any | None:
    """FlashInfer ``MNNVLCuteDSLAllGatherDualNormWorkspace`` for the TP group."""

    def create() -> Any:
        from flashinfer.comm.mnnvl_cutedsl_ag import (
            LLAllGatherDualNormTuning,
            MNNVLCuteDSLAllGatherDualNormWorkspace,
        )

        tp_group = get_tp_group()
        return MNNVLCuteDSLAllGatherDualNormWorkspace(
            tp_group.world_size,
            tp_group.rank_in_group,
            max_tokens,
            hidden_size,
            group=tp_group.device_group,
            tuning=LLAllGatherDualNormTuning(enable_pdl=enable_pdl),
        )

    return _create_on_all_tp_ranks("unpooler all-gather + residual boundary", create)

maybe_create_all_gather_norm_quant(hidden_size, max_tokens)

FlashInfer MNNVLCuteDSLAllGatherNormQuantWorkspace for the TP group.

Source code in vllm/models/muse_spark/nvidia/ops/mnnvl_fusions.py 
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def maybe_create_all_gather_norm_quant(hidden_size: int, max_tokens: int) -> Any | None:
    """FlashInfer ``MNNVLCuteDSLAllGatherNormQuantWorkspace`` for the TP group."""

    def create() -> Any:
        from flashinfer.comm.mnnvl_cutedsl_ag import (
            MNNVLCuteDSLAllGatherNormQuantWorkspace,
        )

        tp_group = get_tp_group()
        return MNNVLCuteDSLAllGatherNormQuantWorkspace(
            tp_group.world_size,
            tp_group.rank_in_group,
            max_tokens,
            hidden_size,
            group=tp_group.device_group,
        )

    return _create_on_all_tp_ranks("pooler all-gather + norm + NVFP4 quant", create)

maybe_create_moe_finalize_all_reduce(hidden_size, top_k, max_tokens, *, enable_pdl)

FlashInfer MNNVLCuteDSLMoEFinalizeAllReduceWorkspace for the TP group.

Source code in vllm/models/muse_spark/nvidia/ops/mnnvl_fusions.py 
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
def maybe_create_moe_finalize_all_reduce(
    hidden_size: int, top_k: int, max_tokens: int, *, enable_pdl: bool
) -> Any | None:
    """FlashInfer ``MNNVLCuteDSLMoEFinalizeAllReduceWorkspace`` for the TP group."""

    def create() -> Any:
        from flashinfer.comm.mnnvl_cutedsl_moe_finalize import (
            MNNVLCuteDSLMoEFinalizeAllReduceWorkspace,
        )

        tp_group = get_tp_group()
        return MNNVLCuteDSLMoEFinalizeAllReduceWorkspace(
            tp_group.world_size,
            tp_group.rank_in_group,
            max_tokens,
            hidden_size,
            top_k,
            group=tp_group.device_group,
            enable_pdl=enable_pdl,
        )

    return _create_on_all_tp_ranks("MoE finalize + all-reduce", create)

tp_group_spans_nodes() cached

Whether the TP group crosses nodes; collective on the first call.

Source code in vllm/models/muse_spark/nvidia/ops/mnnvl_fusions.py 
20
21
22
23
@functools.cache
def tp_group_spans_nodes() -> bool:
    """Whether the TP group crosses nodes; collective on the first call."""
    return not all(in_the_same_node_as(get_tp_group().cpu_group, source_rank=0))