Skip to content

Commit e168999

Browse files
apstenku123apstenku123
committed
fix: thread side channels through generation
1 parent 15763a8 commit e168999

3 files changed

Lines changed: 370 additions & 10 deletions

File tree

‎cppmega_mlx/inference/generation.py‎

Lines changed: 221 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
from __future__ import annotations
44

5-
from collections.abc import Callable, Iterator
5+
from collections.abc import Callable, Iterator, Mapping
66
from dataclasses import dataclass
77
from typing import Any, Literal, cast
88

@@ -20,6 +20,20 @@
2020

2121
GenerationFinishReason = Literal["eos", "length"]
2222

23+
_ZERO_GENERATED_MODEL_KWARGS = frozenset(
24+
{
25+
"structure_ids",
26+
"dep_levels",
27+
"ast_depth_ids",
28+
"sibling_index_ids",
29+
"node_type_ids",
30+
}
31+
)
32+
_REPEAT_GENERATED_MODEL_KWARGS = frozenset({"document_ids", "platform_ids"})
33+
_SEQUENCE_ALIGNED_MODEL_KWARGS = (
34+
_ZERO_GENERATED_MODEL_KWARGS | _REPEAT_GENERATED_MODEL_KWARGS
35+
)
36+
2337

2438
@dataclass(frozen=True)
2539
class GenerationChunk:
@@ -86,6 +100,7 @@ def generate_tokens(
86100
prompt_ids: mx.array,
87101
*,
88102
max_new_tokens: int,
103+
model_kwargs: Mapping[str, mx.array] | None = None,
89104
eos_token_id: int | None = None,
90105
temperature: float = 1.0,
91106
top_k: int | None = None,
@@ -116,7 +131,10 @@ def generate_tokens(
116131
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
117132
raise ValueError("generation would exceed model.config.max_seq_length")
118133

119-
step_logits = next_token_logits(model(tokens), tokens)
134+
step_logits = next_token_logits(
135+
model(tokens, **_model_kwargs_for_prefix(model_kwargs, tokens)),
136+
tokens,
137+
)
120138

121139
step_key = None
122140
if key is not None:
@@ -144,6 +162,7 @@ def generate_tokens_with_kv_cache(
144162
prompt_ids: mx.array,
145163
*,
146164
max_new_tokens: int,
165+
model_kwargs: Mapping[str, mx.array] | None = None,
147166
cache: ContiguousKVCache | None = None,
148167
cache_config: ContiguousKVCacheConfig | None = None,
149168
num_layers: int | None = None,
@@ -194,7 +213,14 @@ def generate_tokens_with_kv_cache(
194213
)
195214

196215
key = rng_key
197-
step_logits = next_token_logits(model(tokens, kv_cache=kv_cache), tokens)
216+
step_logits = next_token_logits(
217+
model(
218+
tokens,
219+
kv_cache=kv_cache,
220+
**_model_kwargs_for_prefix(model_kwargs, tokens),
221+
),
222+
tokens,
223+
)
198224
for step in range(max_new_tokens):
199225
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
200226
raise ValueError("generation would exceed model.config.max_seq_length")
@@ -222,7 +248,14 @@ def generate_tokens_with_kv_cache(
222248
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
223249
raise ValueError("generation would exceed model.config.max_seq_length")
224250

225-
step_logits = next_token_logits(model(next_token, kv_cache=kv_cache), next_token)
251+
step_logits = next_token_logits(
252+
model(
253+
next_token,
254+
kv_cache=kv_cache,
255+
**_model_kwargs_for_generated_step(model_kwargs, next_token),
256+
),
257+
next_token,
258+
)
226259

227260
return tokens
228261

@@ -231,6 +264,7 @@ def build_prompt_cache(
231264
model: Any,
232265
prompt_ids: mx.array,
233266
*,
267+
model_kwargs: Mapping[str, mx.array] | None = None,
234268
cache: ContiguousKVCache | None = None,
235269
cache_config: ContiguousKVCacheConfig | None = None,
236270
num_layers: int | None = None,
@@ -269,7 +303,14 @@ def build_prompt_cache(
269303
if kv_cache_position(kv_cache) != 0:
270304
raise RuntimeError("prompt cache build requires an empty contiguous KV cache")
271305

272-
next_logits = next_token_logits(model(prompt_ids, kv_cache=kv_cache), prompt_ids)
306+
next_logits = next_token_logits(
307+
model(
308+
prompt_ids,
309+
kv_cache=kv_cache,
310+
**_model_kwargs_for_prefix(model_kwargs, prompt_ids),
311+
),
312+
prompt_ids,
313+
)
273314
return PromptCacheEntry(
274315
prompt_ids=mx.array(prompt_ids),
275316
cache=kv_cache,
@@ -283,6 +324,7 @@ def generate_tokens_with_prompt_cache(
283324
*,
284325
prompt_cache: PromptCacheEntry,
285326
max_new_tokens: int,
327+
model_kwargs: Mapping[str, mx.array] | None = None,
286328
eos_token_id: int | None = None,
287329
temperature: float = 1.0,
288330
top_k: int | None = None,
@@ -316,7 +358,18 @@ def generate_tokens_with_prompt_cache(
316358
if suffix.shape[1] == 0:
317359
step_logits = prompt_cache.next_logits
318360
else:
319-
step_logits = next_token_logits(model(suffix, kv_cache=kv_cache), suffix)
361+
step_logits = next_token_logits(
362+
model(
363+
suffix,
364+
kv_cache=kv_cache,
365+
**_model_kwargs_for_slice(
366+
model_kwargs,
367+
start=prefix_length,
368+
tokens=suffix,
369+
),
370+
),
371+
suffix,
372+
)
320373

321374
key = rng_key
322375
for step in range(max_new_tokens):
@@ -343,7 +396,14 @@ def generate_tokens_with_prompt_cache(
343396
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
344397
raise ValueError("generation would exceed model.config.max_seq_length")
345398

346-
step_logits = next_token_logits(model(next_token, kv_cache=kv_cache), next_token)
399+
step_logits = next_token_logits(
400+
model(
401+
next_token,
402+
kv_cache=kv_cache,
403+
**_model_kwargs_for_generated_step(model_kwargs, next_token),
404+
),
405+
next_token,
406+
)
347407

348408
return tokens
349409

@@ -353,6 +413,7 @@ def stream_generate_tokens(
353413
prompt_ids: mx.array,
354414
*,
355415
max_new_tokens: int,
416+
model_kwargs: Mapping[str, mx.array] | None = None,
356417
eos_token_id: int | None = None,
357418
temperature: float = 1.0,
358419
top_k: int | None = None,
@@ -394,6 +455,7 @@ def stream_generate_tokens(
394455
model,
395456
prompt_ids,
396457
max_new_tokens=max_new_tokens,
458+
model_kwargs=model_kwargs,
397459
eos_token_id=eos_token_id,
398460
temperature=temperature,
399461
top_k=top_k,
@@ -427,7 +489,10 @@ def stream_generate_tokens(
427489
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
428490
raise ValueError("generation would exceed model.config.max_seq_length")
429491

430-
step_logits = next_token_logits(model(tokens), tokens)
492+
step_logits = next_token_logits(
493+
model(tokens, **_model_kwargs_for_prefix(model_kwargs, tokens)),
494+
tokens,
495+
)
431496
key, step_key = _split_generation_key(key)
432497
next_token = sample_next_token(
433498
step_logits,
@@ -492,6 +557,137 @@ def _resolve_kv_cache(
492557
)
493558

494559

560+
def _model_kwargs_for_prefix(
561+
model_kwargs: Mapping[str, mx.array] | None,
562+
tokens: mx.array,
563+
) -> dict[str, mx.array]:
564+
if not model_kwargs:
565+
return {}
566+
return {
567+
name: _align_model_kwarg_for_prefix(name, value, tokens)
568+
for name, value in model_kwargs.items()
569+
}
570+
571+
572+
def _model_kwargs_for_slice(
573+
model_kwargs: Mapping[str, mx.array] | None,
574+
*,
575+
start: int,
576+
tokens: mx.array,
577+
) -> dict[str, mx.array]:
578+
if not model_kwargs:
579+
return {}
580+
batch_size, sequence_length = _batch_sequence(tokens)
581+
out: dict[str, mx.array] = {}
582+
for name, value in model_kwargs.items():
583+
_validate_model_kwarg_tensor(name, value)
584+
if not _sequence_aligned_model_kwarg(name, value, batch_size):
585+
out[name] = value
586+
continue
587+
available = max(0, int(value.shape[1]) - start)
588+
if available >= sequence_length:
589+
out[name] = value[:, start : start + sequence_length, ...]
590+
continue
591+
chunks = []
592+
if available > 0:
593+
chunks.append(value[:, start:, ...])
594+
chunks.append(
595+
_generated_model_kwarg_tail(
596+
name,
597+
value,
598+
batch_size=batch_size,
599+
sequence_length=sequence_length - available,
600+
)
601+
)
602+
out[name] = mx.concatenate(chunks, axis=1) if len(chunks) > 1 else chunks[0]
603+
return out
604+
605+
606+
def _model_kwargs_for_generated_step(
607+
model_kwargs: Mapping[str, mx.array] | None,
608+
tokens: mx.array,
609+
) -> dict[str, mx.array]:
610+
if not model_kwargs:
611+
return {}
612+
batch_size, sequence_length = _batch_sequence(tokens)
613+
out: dict[str, mx.array] = {}
614+
for name, value in model_kwargs.items():
615+
_validate_model_kwarg_tensor(name, value)
616+
if _sequence_aligned_model_kwarg(name, value, batch_size):
617+
out[name] = _generated_model_kwarg_tail(
618+
name,
619+
value,
620+
batch_size=batch_size,
621+
sequence_length=sequence_length,
622+
)
623+
else:
624+
out[name] = value
625+
return out
626+
627+
628+
def _align_model_kwarg_for_prefix(
629+
name: str,
630+
value: mx.array,
631+
tokens: mx.array,
632+
) -> mx.array:
633+
_validate_model_kwarg_tensor(name, value)
634+
batch_size, sequence_length = _batch_sequence(tokens)
635+
if not _sequence_aligned_model_kwarg(name, value, batch_size):
636+
return value
637+
638+
current_length = int(value.shape[1])
639+
if current_length == sequence_length:
640+
return value
641+
if current_length > sequence_length:
642+
return value[:, :sequence_length, ...]
643+
tail = _generated_model_kwarg_tail(
644+
name,
645+
value,
646+
batch_size=batch_size,
647+
sequence_length=sequence_length - current_length,
648+
)
649+
return mx.concatenate([value, tail], axis=1)
650+
651+
652+
def _validate_model_kwarg_tensor(name: str, value: mx.array) -> None:
653+
if not isinstance(value, mx.array):
654+
raise TypeError(f"model_kwargs[{name!r}] must be an mlx.core.array")
655+
656+
657+
def _sequence_aligned_model_kwarg(
658+
name: str,
659+
value: mx.array,
660+
batch_size: int,
661+
) -> bool:
662+
shape = value.shape
663+
if not shape or int(shape[0]) != batch_size:
664+
return False
665+
if name == "platform_ids":
666+
return len(shape) == 3
667+
return name in _SEQUENCE_ALIGNED_MODEL_KWARGS and len(shape) >= 2
668+
669+
670+
def _generated_model_kwarg_tail(
671+
name: str,
672+
value: mx.array,
673+
*,
674+
batch_size: int,
675+
sequence_length: int,
676+
) -> mx.array:
677+
if sequence_length <= 0:
678+
return value[:, :0, ...]
679+
tail_shape = (batch_size, sequence_length, *tuple(value.shape[2:]))
680+
if name in _REPEAT_GENERATED_MODEL_KWARGS and int(value.shape[1]) > 0:
681+
return mx.broadcast_to(value[:, -1:, ...], tail_shape)
682+
return mx.zeros(tail_shape, dtype=value.dtype)
683+
684+
685+
def _batch_sequence(tokens: mx.array) -> tuple[int, int]:
686+
if len(tokens.shape) != 2:
687+
raise ValueError("tokens must have shape (batch, sequence)")
688+
return int(tokens.shape[0]), int(tokens.shape[1])
689+
690+
495691
def _validate_prompt_cache_prefix(
496692
prompt_ids: mx.array,
497693
prompt_cache: PromptCacheEntry,
@@ -513,6 +709,7 @@ def _stream_generate_tokens_with_kv_cache(
513709
prompt_ids: mx.array,
514710
*,
515711
max_new_tokens: int,
712+
model_kwargs: Mapping[str, mx.array] | None,
516713
eos_token_id: int | None,
517714
temperature: float,
518715
top_k: int | None,
@@ -547,7 +744,14 @@ def _stream_generate_tokens_with_kv_cache(
547744
)
548745

549746
key = rng_key
550-
step_logits = next_token_logits(model(tokens, kv_cache=kv_cache), tokens)
747+
step_logits = next_token_logits(
748+
model(
749+
tokens,
750+
kv_cache=kv_cache,
751+
**_model_kwargs_for_prefix(model_kwargs, tokens),
752+
),
753+
tokens,
754+
)
551755
for step in range(max_new_tokens):
552756
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
553757
raise ValueError("generation would exceed model.config.max_seq_length")
@@ -574,7 +778,14 @@ def _stream_generate_tokens_with_kv_cache(
574778
break
575779
if max_seq_length is not None and tokens.shape[1] >= max_seq_length:
576780
raise ValueError("generation would exceed model.config.max_seq_length")
577-
step_logits = next_token_logits(model(next_token, kv_cache=kv_cache), next_token)
781+
step_logits = next_token_logits(
782+
model(
783+
next_token,
784+
kv_cache=kv_cache,
785+
**_model_kwargs_for_generated_step(model_kwargs, next_token),
786+
),
787+
next_token,
788+
)
578789

579790

580791
def _split_generation_key(key: Any | None) -> tuple[Any | None, Any | None]:

0 commit comments

Comments
 (0)