Skip to content

vllm.models.deepseek_v41.amd.mhc

Overlap the DSV4.1 mHC gate projection with the next sublayer on ROCm.

Classes:

  • MHCOverlapStream –

    The global aux stream with the fork and join events of the overlap.

Functions:

MHCOverlapStream

The global aux stream with the fork and join events of the overlap.

Every fork and join waits on its event before the next records it, so all layers share the two events.

Source code in vllm/models/deepseek_v41/amd/mhc.py 
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
class MHCOverlapStream:
    """The global aux stream with the fork and join events of the overlap.

    Every fork and join waits on its event before the next records it, so all
    layers share the two events.
    """

    def __init__(self) -> None:
        stream = aux_stream()
        assert stream is not None
        self.stream = stream
        self._fork_event = torch.cuda.Event()
        self._join_event = torch.cuda.Event()

    def fork(self) -> None:
        self._fork_event.record()
        self.stream.wait_event(self._fork_event)

    def join(self) -> None:
        self._join_event.record(self.stream)
        torch.cuda.current_stream().wait_event(self._join_event)

mhc_pre_delayed_overlap(pre_delayed, residual, fn, hc_scale, hc_base, rms_eps, hc_pre_eps, hc_sinkhorn_eps, hc_post_mult_value, sinkhorn_repeat, pre_mix=None, x=None, norm_weight=None, norm_eps=1e-06, sublayer_out=None, post_layer_mix=None, comb_res_mix=None, *, overlap)

pre_delayed with the gate projection moved onto overlap.stream.

Returns residual, post_mix, comb_mix, layer_input, next_pre_mix. Only the residual and layer input are ready on the caller stream; call overlap.join() before using the gates. Seams without a post block, a carried pre-mix or a fused norm run pre_delayed on the caller stream.

Source code in vllm/models/deepseek_v41/amd/mhc.py 
 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
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
def mhc_pre_delayed_overlap(
    pre_delayed: MHCPreDelayedOp,
    residual: torch.Tensor,
    fn: torch.Tensor,
    hc_scale: torch.Tensor,
    hc_base: torch.Tensor,
    rms_eps: float,
    hc_pre_eps: float,
    hc_sinkhorn_eps: float,
    hc_post_mult_value: float,
    sinkhorn_repeat: int,
    pre_mix: torch.Tensor | None = None,
    x: torch.Tensor | None = None,
    norm_weight: torch.Tensor | None = None,
    norm_eps: float = 1e-6,
    sublayer_out: torch.Tensor | None = None,
    post_layer_mix: torch.Tensor | None = None,
    comb_res_mix: torch.Tensor | None = None,
    *,
    overlap: MHCOverlapStream,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
    """``pre_delayed`` with the gate projection moved onto ``overlap.stream``.

    Returns residual, post_mix, comb_mix, layer_input, next_pre_mix. Only the
    residual and layer input are ready on the caller stream; call
    ``overlap.join()`` before using the gates. Seams without a post block, a
    carried pre-mix or a fused norm run ``pre_delayed`` on the caller stream.
    """
    main = torch.cuda.current_stream()
    if (
        x is not None
        or pre_mix is None
        or norm_weight is None
        or sublayer_out is None
        or post_layer_mix is None
        or comb_res_mix is None
    ):
        # Fork anyway: the caller joins after every seam, and a join needs the
        # stream to be part of the capture.
        overlap.fork()
        return pre_delayed(
            residual,
            fn,
            hc_scale,
            hc_base,
            rms_eps,
            hc_pre_eps,
            hc_sinkhorn_eps,
            hc_post_mult_value,
            sinkhorn_repeat,
            pre_mix=pre_mix,
            x=x,
            norm_weight=norm_weight,
            norm_eps=norm_eps,
            sublayer_out=sublayer_out,
            post_layer_mix=post_layer_mix,
            comb_res_mix=comb_res_mix,
        )
    residual, layer_input = mhc_post_collapse_rms_norm_triton(
        residual,
        sublayer_out,
        post_layer_mix,
        comb_res_mix,
        pre_mix,
        norm_weight,
        norm_eps,
    )
    stream = overlap.stream
    overlap.fork()
    with torch.cuda.stream(stream):
        _, post_mix, comb_mix, _, next_pre_mix = pre_delayed(
            residual,
            fn,
            hc_scale,
            hc_base,
            rms_eps,
            hc_pre_eps,
            hc_sinkhorn_eps,
            hc_post_mult_value,
            sinkhorn_repeat,
            pre_mix=pre_mix,
            norm_weight=norm_weight,
            norm_eps=norm_eps,
        )
    for tensor in (residual, pre_mix):
        tensor.record_stream(stream)
    for tensor in (post_mix, comb_mix, next_pre_mix):
        tensor.record_stream(main)
    return residual, post_mix, comb_mix, layer_input, next_pre_mix