Skip to content

vllm.models.deepseek_v41.nvidia.specialized_decode.boundary.boundary

The mHC boundaries of the specialized decode.

Every decoder sublayer ends in a TP-partial output (the attention's wo_b, the MoE's routed and shared experts). The next boundary reduces it across the TP group, mixes it into the hc residual streams (mHC post) and collapses and normalizes the streams into the next sublayer's input (mHC pre). Here the fused all-reduce + mHC kernel (BoundaryAllReduce) does all of that for every step the executor runs; the next sublayer's mHC coefficients then come from the stock mhc_pre_delayed_overlap, on the shared mHC side stream for up to MHC_OVERLAP_MAX_ROWS rows and inline after the fused kernel above.

The fused kernel also finalizes a MoE output left unfinalized, collects a wo_b output the publishing GEMM already stored into its mailbox, writes the attention input as MXFP8 for fused_wqa_wkv, and writes only the streams (residual-only) where the next consumer is not a sublayer's pre: an Engram injection and the final collapse.

Classes:

  • Boundaries –

    The boundaries of one decode step's layer loop, on one fused kernel.

  • BoundaryState –

    The hc residual streams and the mixes carried to the next boundary.

Functions:

Boundaries

The boundaries of one decode step's layer loop, on one fused kernel.

Methods:

  • first –

    The first layer's attention entry: the embedding broadcast to the

  • join –

    Make the current stream wait for the side-stream coefficients.

  • post_only –

    Reduce state.pending into the streams and return them, for an

  • pre –

    An Engram layer's attention entry, on the injected streams.

  • sublayer –

    Reduce state.pending into the streams and prepare the attention

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
 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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
class Boundaries:
    """The boundaries of one decode step's layer loop, on one fused kernel."""

    def __init__(self, all_reduce: "BoundaryAllReduce") -> None:
        self.all_reduce = all_reduce

    def first(
        self, layer: "DeepseekV4DecoderLayer", x: torch.Tensor, step: "DecodeStep"
    ) -> tuple[BoundaryState, torch.Tensor]:
        """The first layer's attention entry: the embedding broadcast to the
        hc streams, whose identity pre-mix selects stream zero."""
        assert layer.hc_attn_fn_broadcast is not None
        residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous()
        return self._pre(layer, residual, layer.hc_attn_fn_broadcast, None, x, step)

    def pre(
        self,
        layer: "DeepseekV4DecoderLayer",
        residual: torch.Tensor,
        state: BoundaryState,
        step: "DecodeStep",
    ) -> tuple[BoundaryState, torch.Tensor]:
        """An Engram layer's attention entry, on the injected streams."""
        return self._pre(layer, residual, layer.hc_attn_fn, state.pre, None, step)

    def sublayer(
        self,
        layer: "DeepseekV4DecoderLayer",
        ffn: bool,
        state: BoundaryState,
        step: "DecodeStep",
        *,
        quantize: bool = False,
        capture_aux: bool = False,
    ) -> tuple[
        BoundaryState, torch.Tensor, QuantizedActivation | None, torch.Tensor | None
    ]:
        """Reduce ``state.pending`` into the streams and prepare the attention
        (or with ``ffn``, the MoE) input of ``layer``.

        Returns the new state, the layer input, the same input as MXFP8 with
        ``quantize``, and with ``capture_aux`` the mean of the streams (on the
        side stream when the coefficients run there: join it before reading).
        """
        if ffn:
            fn, scale, base = layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base
            norm = layer.ffn_norm
        else:
            fn, scale, base = layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base
            norm = layer.attn_norm
        rows, _, hidden = state.residual.shape
        quant = (
            empty_quantized_layer_input(rows, hidden, state.residual.device)
            if quantize
            else None
        )
        residual, layer_input = self._reduce(
            state,
            state.pre,
            norm,
            None if quant is None else (quant.data, quant.scale),
        )
        assert layer_input is not None
        stream = layer.mhc_stream if step.side_stream else torch.cuda.current_stream()
        assert stream is not None
        aux = None
        if capture_aux and step.side_stream:
            streams_ready = torch.cuda.Event()
            streams_ready.record()
        elif capture_aux:
            aux = residual.mean(dim=1)
        post, comb, layer_input, next_pre = mhc_pre_delayed_overlap(
            residual,
            fn,
            scale,
            base,
            layer.rms_norm_eps,
            layer.hc_eps,
            layer.hc_eps,
            layer.hc_post_alpha,
            layer.hc_sinkhorn_iters,
            pre_mix=state.pre,
            norm_weight=norm.weight,
            norm_eps=norm.variance_epsilon,
            stream=stream,
            layer_input=layer_input,
        )
        if capture_aux and step.side_stream:
            # residual.mean(dim=1), after the coefficients on the side stream.
            aux = residual.new_empty(rows, hidden)
            stream.wait_event(streams_ready)
            with torch.cuda.stream(stream):
                torch.mean(residual, dim=1, out=aux)
            residual.record_stream(stream)
            aux.record_stream(stream)
        return BoundaryState(residual, post, comb, next_pre), layer_input, quant, aux

    def post_only(self, state: BoundaryState) -> torch.Tensor:
        """Reduce ``state.pending`` into the streams and return them, for an
        Engram injection or the final collapse."""
        if isinstance(state.pending, MoEOutput):
            residual, _ = self._reduce(state, None, None, None)
            return residual
        assert isinstance(state.pending, torch.Tensor)
        x = tensor_model_parallel_all_reduce(state.pending)
        return mhc_post_tilelang(x, state.residual, state.post, state.comb)

    def join(self, layer: "DeepseekV4DecoderLayer", step: "DecodeStep") -> None:
        """Make the current stream wait for the side-stream coefficients."""
        if step.side_stream:
            assert layer.mhc_stream is not None
            torch.cuda.current_stream().wait_stream(layer.mhc_stream)

    def _pre(
        self,
        layer: "DeepseekV4DecoderLayer",
        residual: torch.Tensor,
        fn: torch.Tensor,
        pre_mix: torch.Tensor | None,
        x: torch.Tensor | None,
        step: "DecodeStep",
    ) -> tuple[BoundaryState, torch.Tensor]:
        """The stock attention entry without a reduction: mHC pre of
        ``residual`` on the side stream, or inline with TileLang."""
        mhc_pre = (
            partial(mhc_pre_delayed_overlap, stream=layer.mhc_stream)
            if step.side_stream
            else mhc_pre_delayed_tilelang
        )
        post, comb, layer_input, next_pre = mhc_pre(
            residual,
            fn,
            layer.hc_attn_scale,
            layer.hc_attn_base,
            layer.rms_norm_eps,
            layer.hc_eps,
            layer.hc_eps,
            layer.hc_post_alpha,
            layer.hc_sinkhorn_iters,
            pre_mix=pre_mix,
            x=x,
            norm_weight=layer.attn_norm.weight,
            norm_eps=layer.attn_norm.variance_epsilon,
        )
        return BoundaryState(residual, post, comb, next_pre), layer_input

    def _reduce(
        self,
        state: BoundaryState,
        pre: torch.Tensor | None,
        norm: RMSNorm | None,
        mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None,
    ) -> tuple[torch.Tensor, torch.Tensor | None]:
        """The fused kernel on ``state.pending``; residual-only without ``pre``."""
        norm_weight = None if norm is None else norm.weight
        norm_eps = 0.0 if norm is None else norm.variance_epsilon
        x = state.pending
        if isinstance(x, MoEOutput):
            routed = x.routed
            assert isinstance(routed, UnfinalizedMoEOutput)
            assert x.shared_output is not None
            return self.all_reduce.finalize(
                routed.gemm2_permuted,
                routed.expert_weights,
                routed.expanded_idx_to_permuted_idx,
                x.shared_output,
                state.residual,
                state.post,
                state.comb,
                pre,
                norm_weight,
                norm_eps,
                mxfp8_out,
            )
        assert x is not None
        return self.all_reduce(
            x,
            state.residual,
            state.post,
            state.comb,
            pre,
            norm_weight,
            norm_eps,
            mxfp8_out,
        )

_pre(layer, residual, fn, pre_mix, x, step)

The stock attention entry without a reduction: mHC pre of residual on the side stream, or inline with TileLang.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
def _pre(
    self,
    layer: "DeepseekV4DecoderLayer",
    residual: torch.Tensor,
    fn: torch.Tensor,
    pre_mix: torch.Tensor | None,
    x: torch.Tensor | None,
    step: "DecodeStep",
) -> tuple[BoundaryState, torch.Tensor]:
    """The stock attention entry without a reduction: mHC pre of
    ``residual`` on the side stream, or inline with TileLang."""
    mhc_pre = (
        partial(mhc_pre_delayed_overlap, stream=layer.mhc_stream)
        if step.side_stream
        else mhc_pre_delayed_tilelang
    )
    post, comb, layer_input, next_pre = mhc_pre(
        residual,
        fn,
        layer.hc_attn_scale,
        layer.hc_attn_base,
        layer.rms_norm_eps,
        layer.hc_eps,
        layer.hc_eps,
        layer.hc_post_alpha,
        layer.hc_sinkhorn_iters,
        pre_mix=pre_mix,
        x=x,
        norm_weight=layer.attn_norm.weight,
        norm_eps=layer.attn_norm.variance_epsilon,
    )
    return BoundaryState(residual, post, comb, next_pre), layer_input

_reduce(state, pre, norm, mxfp8_out)

The fused kernel on state.pending; residual-only without pre.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
def _reduce(
    self,
    state: BoundaryState,
    pre: torch.Tensor | None,
    norm: RMSNorm | None,
    mxfp8_out: tuple[torch.Tensor, torch.Tensor] | None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
    """The fused kernel on ``state.pending``; residual-only without ``pre``."""
    norm_weight = None if norm is None else norm.weight
    norm_eps = 0.0 if norm is None else norm.variance_epsilon
    x = state.pending
    if isinstance(x, MoEOutput):
        routed = x.routed
        assert isinstance(routed, UnfinalizedMoEOutput)
        assert x.shared_output is not None
        return self.all_reduce.finalize(
            routed.gemm2_permuted,
            routed.expert_weights,
            routed.expanded_idx_to_permuted_idx,
            x.shared_output,
            state.residual,
            state.post,
            state.comb,
            pre,
            norm_weight,
            norm_eps,
            mxfp8_out,
        )
    assert x is not None
    return self.all_reduce(
        x,
        state.residual,
        state.post,
        state.comb,
        pre,
        norm_weight,
        norm_eps,
        mxfp8_out,
    )

first(layer, x, step)

The first layer's attention entry: the embedding broadcast to the hc streams, whose identity pre-mix selects stream zero.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
86
87
88
89
90
91
92
93
def first(
    self, layer: "DeepseekV4DecoderLayer", x: torch.Tensor, step: "DecodeStep"
) -> tuple[BoundaryState, torch.Tensor]:
    """The first layer's attention entry: the embedding broadcast to the
    hc streams, whose identity pre-mix selects stream zero."""
    assert layer.hc_attn_fn_broadcast is not None
    residual = x.unsqueeze(1).expand(-1, layer.hc_mult, -1).contiguous()
    return self._pre(layer, residual, layer.hc_attn_fn_broadcast, None, x, step)

join(layer, step)

Make the current stream wait for the side-stream coefficients.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
187
188
189
190
191
def join(self, layer: "DeepseekV4DecoderLayer", step: "DecodeStep") -> None:
    """Make the current stream wait for the side-stream coefficients."""
    if step.side_stream:
        assert layer.mhc_stream is not None
        torch.cuda.current_stream().wait_stream(layer.mhc_stream)

post_only(state)

Reduce state.pending into the streams and return them, for an Engram injection or the final collapse.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
177
178
179
180
181
182
183
184
185
def post_only(self, state: BoundaryState) -> torch.Tensor:
    """Reduce ``state.pending`` into the streams and return them, for an
    Engram injection or the final collapse."""
    if isinstance(state.pending, MoEOutput):
        residual, _ = self._reduce(state, None, None, None)
        return residual
    assert isinstance(state.pending, torch.Tensor)
    x = tensor_model_parallel_all_reduce(state.pending)
    return mhc_post_tilelang(x, state.residual, state.post, state.comb)

pre(layer, residual, state, step)

An Engram layer's attention entry, on the injected streams.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
 95
 96
 97
 98
 99
100
101
102
103
def pre(
    self,
    layer: "DeepseekV4DecoderLayer",
    residual: torch.Tensor,
    state: BoundaryState,
    step: "DecodeStep",
) -> tuple[BoundaryState, torch.Tensor]:
    """An Engram layer's attention entry, on the injected streams."""
    return self._pre(layer, residual, layer.hc_attn_fn, state.pre, None, step)

sublayer(layer, ffn, state, step, *, quantize=False, capture_aux=False)

Reduce state.pending into the streams and prepare the attention (or with ffn, the MoE) input of layer.

Returns the new state, the layer input, the same input as MXFP8 with quantize, and with capture_aux the mean of the streams (on the side stream when the coefficients run there: join it before reading).

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def sublayer(
    self,
    layer: "DeepseekV4DecoderLayer",
    ffn: bool,
    state: BoundaryState,
    step: "DecodeStep",
    *,
    quantize: bool = False,
    capture_aux: bool = False,
) -> tuple[
    BoundaryState, torch.Tensor, QuantizedActivation | None, torch.Tensor | None
]:
    """Reduce ``state.pending`` into the streams and prepare the attention
    (or with ``ffn``, the MoE) input of ``layer``.

    Returns the new state, the layer input, the same input as MXFP8 with
    ``quantize``, and with ``capture_aux`` the mean of the streams (on the
    side stream when the coefficients run there: join it before reading).
    """
    if ffn:
        fn, scale, base = layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base
        norm = layer.ffn_norm
    else:
        fn, scale, base = layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base
        norm = layer.attn_norm
    rows, _, hidden = state.residual.shape
    quant = (
        empty_quantized_layer_input(rows, hidden, state.residual.device)
        if quantize
        else None
    )
    residual, layer_input = self._reduce(
        state,
        state.pre,
        norm,
        None if quant is None else (quant.data, quant.scale),
    )
    assert layer_input is not None
    stream = layer.mhc_stream if step.side_stream else torch.cuda.current_stream()
    assert stream is not None
    aux = None
    if capture_aux and step.side_stream:
        streams_ready = torch.cuda.Event()
        streams_ready.record()
    elif capture_aux:
        aux = residual.mean(dim=1)
    post, comb, layer_input, next_pre = mhc_pre_delayed_overlap(
        residual,
        fn,
        scale,
        base,
        layer.rms_norm_eps,
        layer.hc_eps,
        layer.hc_eps,
        layer.hc_post_alpha,
        layer.hc_sinkhorn_iters,
        pre_mix=state.pre,
        norm_weight=norm.weight,
        norm_eps=norm.variance_epsilon,
        stream=stream,
        layer_input=layer_input,
    )
    if capture_aux and step.side_stream:
        # residual.mean(dim=1), after the coefficients on the side stream.
        aux = residual.new_empty(rows, hidden)
        stream.wait_event(streams_ready)
        with torch.cuda.stream(stream):
            torch.mean(residual, dim=1, out=aux)
        residual.record_stream(stream)
        aux.record_stream(stream)
    return BoundaryState(residual, post, comb, next_pre), layer_input, quant, aux

BoundaryState dataclass

The hc residual streams and the mixes carried to the next boundary.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
49
50
51
52
53
54
55
56
57
58
59
60
61
@dataclass
class BoundaryState:
    """The hc residual streams and the mixes carried to the next boundary."""

    residual: torch.Tensor
    # Post-mix and residual mix of the sublayer that produces ``pending``.
    post: torch.Tensor
    comb: torch.Tensor
    # The pre-mix the next collapse uses.
    pre: torch.Tensor
    # The sublayer's TP-partial output: a tensor (maybe published) or a MoE
    # output left unfinalized.
    pending: torch.Tensor | MoEOutput | None = None

empty_quantized_layer_input(num_tokens, hidden_size, device)

An unfilled MXFP8 layer input with FlashInfer's swizzled 128x4 scales.

Source code in vllm/models/deepseek_v41/nvidia/specialized_decode/boundary/boundary.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
def empty_quantized_layer_input(
    num_tokens: int, hidden_size: int, device: torch.device
) -> QuantizedActivation:
    """An unfilled MXFP8 layer input with FlashInfer's swizzled 128x4 scales."""
    shape = torch.Size((num_tokens, hidden_size))
    return QuantizedActivation(
        data=torch.empty(shape, dtype=torch.float8_e4m3fn, device=device),
        scale=torch.empty(
            mxfp8_scale_bytes(num_tokens, hidden_size), dtype=torch.uint8, device=device
        ),
        orig_dtype=torch.bfloat16,
        orig_shape=shape,
        quant_key=kMxfp8Dynamic,
    )