Skip to content

Commit ae792e7

Browse files
tboerstadmodularbot
tboerstad
authored and
modularbot
committed
[Kernels][MAX] Fix Inkling conv state after rejected draft tokens
Rejected MTP draft tokens no longer leak into Inkling's conv history. The conv state was a shift register the kernel updated in place, so a verify step committed every draft token, accepted or not. It is now a ring of past inputs indexed by position. A rollback needs no work: the sequence overwrites those entries when it reaches the same positions again. Each request's ring holds one conv window plus the draft length. The MTP draft's per-request scratch pools and their per-step zeroing are gone, replaced by one shared zero slot. MODULAR_ORIG_COMMIT_REV_ID: c523d193f1fbedc7c888be981f129d979230bfae
1 parent 7be84fd commit ae792e7

18 files changed

Lines changed: 886 additions & 254 deletions

File tree

‎max/kernels/src/graph_compiler/builtin_kernels/short_conv_ring.mojo‎

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
"""Registers the ring-state short convolution graph ops; see
14+
`nn.short_conv_ring`."""
15+
16+
import extensibility
17+
from extensibility import InputTensor, OutputTensor
18+
from extensibility import (
19+
_MutableInputTensor as MutableInputTensor,
20+
)
21+
from max.gpu.host import DeviceContext
22+
23+
from nn.short_conv_ring import short_conv_ring_commit, short_conv_ring_fwd
24+
25+
26+
@extensibility.register("mo.short_conv_ring_fwd")
27+
struct Struct_short_conv_ring_fwd:
28+
"""`x + conv(x)` over a ragged batch, pre-chunk taps from the ring."""
29+
30+
@inline(.always)
31+
@staticmethod
32+
def execute[
33+
dtype: DType,
34+
ring_dtype: DType,
35+
//,
36+
target: StaticString,
37+
](
38+
output: OutputTensor[dtype=dtype, rank=2, ...],
39+
x: InputTensor[dtype=dtype, rank=2, ...],
40+
weight: InputTensor[dtype=dtype, rank=2, ...],
41+
# Only read here; mutable so the graph orders the commit after it.
42+
ring: MutableInputTensor[dtype=ring_dtype, rank=3, ...],
43+
input_row_offsets: InputTensor[dtype=.uint32, rank=1, ...],
44+
positions: InputTensor[dtype=.uint32, rank=1, ...],
45+
conv_row: InputTensor[dtype=.uint32, rank=1, ...],
46+
context: DeviceContext,
47+
) raises:
48+
short_conv_ring_fwd[target=target](
49+
x.to_tile_tensor[.int64](),
50+
weight.to_tile_tensor[.int64](),
51+
ring.to_tile_tensor[.int64](),
52+
input_row_offsets.to_tile_tensor[.int64](),
53+
positions.to_tile_tensor[.int64](),
54+
conv_row.to_tile_tensor[.int64](),
55+
output.to_tile_tensor[.int64](),
56+
context,
57+
)
58+
59+
60+
@extensibility.register("mo.short_conv_ring_commit")
61+
struct Struct_short_conv_ring_commit:
62+
"""Writes each sequence's last inputs into its conv ring slot."""
63+
64+
@inline(.always)
65+
@staticmethod
66+
def execute[
67+
dtype: DType,
68+
ring_dtype: DType,
69+
//,
70+
target: StaticString,
71+
](
72+
ring: MutableInputTensor[dtype=ring_dtype, rank=3, ...],
73+
x: InputTensor[dtype=dtype, rank=2, ...],
74+
input_row_offsets: InputTensor[dtype=.uint32, rank=1, ...],
75+
positions: InputTensor[dtype=.uint32, rank=1, ...],
76+
conv_row: InputTensor[dtype=.uint32, rank=1, ...],
77+
context: DeviceContext,
78+
) raises:
79+
short_conv_ring_commit[target=target](
80+
x.to_tile_tensor[.int64](),
81+
ring.to_tile_tensor[.int64](),
82+
input_row_offsets.to_tile_tensor[.int64](),
83+
positions.to_tile_tensor[.int64](),
84+
conv_row.to_tile_tensor[.int64](),
85+
context,
86+
)

‎max/kernels/src/nn/short_conv_ring.mojo‎

Lines changed: 267 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,267 @@
1+
# ===----------------------------------------------------------------------=== #
2+
# Copyright (c) 2026, Modular Inc. All rights reserved.
3+
#
4+
# Licensed under the Apache License v2.0 with LLVM Exceptions:
5+
# https://llvm.org/LICENSE.txt
6+
#
7+
# Unless required by applicable law or agreed to in writing, software
8+
# distributed under the License is distributed on an "AS IS" BASIS,
9+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
10+
# See the License for the specific language governing permissions and
11+
# limitations under the License.
12+
# ===----------------------------------------------------------------------=== #
13+
"""Depthwise causal short convolution over a ring of past inputs.
14+
15+
A request's conv state is a ring indexed by absolute position: the input at
16+
position `p` lives at `ring[p % R]`. A forward is two launches that never
17+
touch the same entry:
18+
19+
- `short_conv_ring_fwd`: `x + conv(x)`. A tap at `p - j` reads `x` when it
20+
is inside the chunk, the ring when it is before the chunk, and zero before
21+
position zero. Writes nothing.
22+
- `short_conv_ring_commit`: writes each sequence's last `R` inputs into
23+
their entries. Reads nothing.
24+
25+
One code path therefore covers decode, prefill, mixed batches and
26+
speculative verify. A rejected draft token's entry is overwritten once the
27+
sequence passes its position again, so `R` must be `width - 1` plus the
28+
largest rollback.
29+
"""
30+
31+
from std.math import ceildiv
32+
from std.utils.numerics import get_accum_type
33+
34+
from max.gpu import block_dim, block_idx, thread_idx
35+
from max.gpu.host import DeviceContext
36+
from max.gpu.host.info import is_gpu
37+
from layout import TileTensor
38+
from layout.coord import Coord
39+
from max.runtime.tracing import Trace, TraceLevel, get_safe_task_id
40+
41+
from nn._ragged_utils import get_batch_and_token_idx_from_row_offsets
42+
43+
comptime _BLOCK = 256
44+
"""Channels one block covers."""
45+
46+
47+
@inline(.always)
48+
def _short_conv_ring_step[
49+
x_dtype: DType, ring_dtype: DType, //, width: Int, accum_dtype: DType
50+
](
51+
x: TileTensor[mut=False, x_dtype, ...],
52+
weight: TileTensor[mut=False, x_dtype, ...],
53+
ring: TileTensor[ring_dtype, ...],
54+
token_idx: Int,
55+
channel: Int,
56+
slot: Int,
57+
position: Int,
58+
chunk_start: Int,
59+
) -> Scalar[accum_dtype]:
60+
"""`x + conv(x)` for one token of one channel.
61+
62+
`x[token_idx, channel]` is the token's input. Tap `j` reads position
63+
`position - j` from `x` when it is at or after `chunk_start`, from
64+
`ring[slot, (position - j) % R, channel]` when it is before the chunk,
65+
else zero. `weight[channel, width - 1]` multiplies the current input,
66+
`weight[channel, width - 1 - j]` the input `j` back, as in
67+
`causal_conv1d`.
68+
"""
69+
comptime ring_len = ring.static_shape[1]
70+
comptime assert ring_len != -1, "the conv ring needs a static length"
71+
comptime assert (
72+
ring_len >= width - 1
73+
), "the conv ring must hold at least width - 1 positions of history"
74+
75+
@inline(.always)
76+
def tap(w: Int) {imm} -> Scalar[accum_dtype]:
77+
return weight.load[width=1](Coord(channel, w)).cast[accum_dtype]()[0]
78+
79+
var x_cur = x.load[width=1](Coord(token_idx, channel)).cast[accum_dtype]()[
80+
0
81+
]
82+
var acc = x_cur + x_cur * tap(width - 1)
83+
comptime for j in range(1, width):
84+
var tap_pos = position - j
85+
var value = Scalar[accum_dtype](0)
86+
if tap_pos >= chunk_start:
87+
value = x.load[width=1](Coord(token_idx - j, channel)).cast[
88+
accum_dtype
89+
]()[0]
90+
elif tap_pos >= 0:
91+
value = ring.load[width=1](
92+
Coord(slot, tap_pos % ring_len, channel)
93+
).cast[accum_dtype]()[0]
94+
acc += value * tap(width - 1 - j)
95+
return acc
96+
97+
98+
def short_conv_ring_fwd[
99+
dtype: DType,
100+
ring_dtype: DType,
101+
//,
102+
target: StaticString,
103+
](
104+
x: TileTensor[mut=False, dtype, ...],
105+
weight: TileTensor[mut=False, dtype, ...],
106+
ring: TileTensor[ring_dtype, ...],
107+
input_row_offsets: TileTensor[mut=False, .uint32, ...],
108+
positions: TileTensor[mut=False, .uint32, ...],
109+
conv_row: TileTensor[mut=False, .uint32, ...],
110+
output: TileTensor[mut=True, dtype, ...],
111+
context: DeviceContext,
112+
) raises:
113+
"""`x + conv(x)` over a ragged batch; reads the ring, writes nothing.
114+
115+
`x`, `output`: `[total_seq_len, channels]`. `weight`: `[channels,
116+
width]`. `ring`: `[slots, R, channels]`. `positions` is per token,
117+
`conv_row` per sequence.
118+
"""
119+
comptime assert is_gpu[target](), "short_conv_ring_fwd is GPU-only"
120+
comptime assert x.flat_rank == 2, "x must be [total_seq_len, channels]"
121+
comptime assert output.flat_rank == 2, "output must match x"
122+
comptime assert weight.flat_rank == 2, "weight must be [channels, width]"
123+
comptime assert ring.flat_rank == 3, "ring must be [slots, R, channels]"
124+
comptime width = weight.static_shape[1]
125+
comptime assert width != -1, "Need static shape for weight"
126+
comptime accum_dtype = get_accum_type[dtype]()
127+
128+
var total_seq_len = Int(x.dim[0]())
129+
var channels = Int(x.dim[1]())
130+
if total_seq_len == 0:
131+
return
132+
# Kernel captures must be fixed-width; Int is not device-passable.
133+
var channel_blocks = Int32(ceildiv(channels, _BLOCK))
134+
var channels_dev = Int32(channels)
135+
136+
def kernel() {
137+
var x,
138+
var weight,
139+
var ring,
140+
var input_row_offsets,
141+
var positions,
142+
var conv_row,
143+
var output,
144+
var channel_blocks,
145+
var channels_dev,
146+
}:
147+
var token_idx = Int(block_idx.x) // Int(channel_blocks)
148+
var channel = (Int(block_idx.x) % Int(channel_blocks)) * Int(
149+
block_dim.x
150+
) + Int(thread_idx.x)
151+
if channel >= Int(channels_dev):
152+
return
153+
var batch_idx, idx_in_seq = get_batch_and_token_idx_from_row_offsets(
154+
input_row_offsets, token_idx
155+
)
156+
var position = Int(positions.load[width=1](Coord(token_idx))[0])
157+
var value = _short_conv_ring_step[width, accum_dtype](
158+
x,
159+
weight,
160+
ring,
161+
token_idx,
162+
channel,
163+
Int(conv_row.load[width=1](Coord(batch_idx))[0]),
164+
position,
165+
position - idx_in_seq,
166+
)
167+
output.store[width=1](Coord(token_idx, channel), value.cast[dtype]())
168+
169+
with Trace[TraceLevel.OP,>
170+
"short_conv_ring_fwd.channels_" + String(channels),
171+
task_id=get_safe_task_id(context),
172+
):
173+
context.enqueue_function(
174+
kernel,
175+
grid_dim=total_seq_len * Int(channel_blocks),
176+
block_dim=_BLOCK,
177+
)
178+
179+
180+
@inline(.always)
181+
def _commit_sequence_tail[
182+
x_dtype: DType, //
183+
](
184+
x: TileTensor[mut=False, x_dtype, ...],
185+
ring: TileTensor[mut=True, ...],
186+
input_row_offsets: TileTensor[mut=False, .uint32, ...],
187+
positions: TileTensor[mut=False, .uint32, ...],
188+
conv_row: TileTensor[mut=False, .uint32, ...],
189+
batch_idx: Int,
190+
channel: Int,
191+
):
192+
"""Writes one channel of a sequence's last `R` inputs into its slot.
193+
194+
Earlier inputs share entries with later ones, so only the tail is
195+
written and every entry has one writer.
196+
"""
197+
comptime ring_len = ring.static_shape[1]
198+
comptime assert ring_len != -1, "the conv ring needs a static length"
199+
200+
var start = Int(input_row_offsets.load[width=1](Coord(batch_idx))[0])
201+
var end = Int(input_row_offsets.load[width=1](Coord(batch_idx + 1))[0])
202+
var slot = Int(conv_row.load[width=1](Coord(batch_idx))[0])
203+
for token_idx in range(max(start, end - ring_len), end):
204+
var position = Int(positions.load[width=1](Coord(token_idx))[0])
205+
var value = x.load[width=1](Coord(token_idx, channel))
206+
ring.store[width=1](
207+
Coord(slot, position % ring_len, channel),
208+
value.cast[ring.dtype](),
209+
)
210+
211+
212+
def short_conv_ring_commit[
213+
dtype: DType,
214+
//,
215+
target: StaticString,
216+
](
217+
x: TileTensor[mut=False, dtype, ...],
218+
ring: TileTensor[mut=True, ...],
219+
input_row_offsets: TileTensor[mut=False, .uint32, ...],
220+
positions: TileTensor[mut=False, .uint32, ...],
221+
conv_row: TileTensor[mut=False, .uint32, ...],
222+
context: DeviceContext,
223+
) raises:
224+
"""Writes each sequence's last `R` rows of `x` into its ring slot.
225+
226+
`x`: `[total_seq_len, channels]`. `ring`: `[slots, R, channels]`.
227+
Launch after every reader of the ring in the same forward.
228+
"""
229+
comptime assert is_gpu[target](), "short_conv_ring_commit is GPU-only"
230+
comptime assert x.flat_rank == 2, "x must be [total_seq_len, channels]"
231+
comptime assert ring.flat_rank == 3, "ring must be [slots, R, channels]"
232+
comptime channels = ring.static_shape[2]
233+
comptime assert channels != -1, "the conv ring needs a static width"
234+
235+
var batch_size = Int(input_row_offsets.dim[0]()) - 1
236+
if batch_size <= 0 or x.dim[0]() == 0:
237+
return
238+
239+
def kernel() {
240+
var x,
241+
var ring,
242+
var input_row_offsets,
243+
var positions,
244+
var conv_row,
245+
}:
246+
var channel = Int(block_idx.x) * Int(block_dim.x) + Int(thread_idx.x)
247+
if channel >= channels:
248+
return
249+
_commit_sequence_tail(
250+
x,
251+
ring,
252+
input_row_offsets,
253+
positions,
254+
conv_row,
255+
Int(block_idx.y),
256+
channel,
257+
)
258+
259+
with Trace[TraceLevel.OP,>
260+
"short_conv_ring_commit.channels_" + String(channels),
261+
task_id=get_safe_task_id(context),
262+
):
263+
context.enqueue_function(
264+
kernel,
265+
grid_dim=(ceildiv(channels, _BLOCK), batch_size),
266+
block_dim=_BLOCK,
267+
)

‎max/python/docs/nn.state_space.rst‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,3 +34,5 @@ Functions
3434
kda_decode
3535
mamba2_ssd_chunk_scan_varlen_fwd
3636
mamba2_ssd_chunk_scan_varlen_fwd_inplace
37+
short_conv_ring_commit
38+
short_conv_ring_fwd

‎max/python/max/nn/state_space/__init__.py‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
mamba2_ssd_chunk_scan_varlen_fwd,
2727
mamba2_ssd_chunk_scan_varlen_fwd_inplace,
2828
)
29+
from .short_conv_ring import short_conv_ring_commit, short_conv_ring_fwd
2930
from .varlen_causal_conv1d import causal_conv1d_varlen_fwd
3031

3132
__all__ = [
@@ -36,4 +37,6 @@
3637
"kda_decode",
3738
"mamba2_ssd_chunk_scan_varlen_fwd",
3839
"mamba2_ssd_chunk_scan_varlen_fwd_inplace",
40+
"short_conv_ring_commit",
41+
"short_conv_ring_fwd",
3942
]

0 commit comments

Comments
 (0)