perf: project only sampled rows through lm_head during prefill
- add logits_positions to AutoRegressiveLM.forward, gathering rows before the final norm so the lm_head GEMM covers only the positions prefill samples from - execute_prefill builds last_token_indices up front and passes them in, dropping the post-forward gather of a [tokens, vocab] tensor - prefill graph warmup passes a single index; decode stays untouched (every row is sampled) and prefill itself runs eager, so graph capture is unaffected - update the ragged-prefill fake to slice by the received index and add a packed-row exact-equality test Benchmark: NVIDIA L20 (idle), CUDA 12.8, torch 2.11.0+cu128, 1.2B bf16 checkpoint, 512-token prompts, greedy; prefill B=32: 368.1 -> 323.5 ms (44.5k -> 50.6k tok/s, +13.8%), B=8: 89.3 -> 78.9 ms (+13.2%), B=1: 12.3 -> 11.4 ms (+7.9%); decode step unchanged; full suite: 897 passed
This commit is contained in:
@@ -195,7 +195,11 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
|
||||
executor._workspace = MagicMock()
|
||||
executor._workspace.max_batch_size = 16 # Add max_batch_size for validation
|
||||
all_logits = torch.arange(42, dtype=torch.float32).reshape(6, 7)
|
||||
executor.model = MagicMock(return_value={"logits": all_logits})
|
||||
|
||||
def fake_model(ids, *, position_ids, kv_cache, fwd, logits_positions):
|
||||
return {"logits": all_logits[logits_positions]}
|
||||
|
||||
executor.model = MagicMock(side_effect=fake_model)
|
||||
executor._sample_logits = MagicMock(
|
||||
return_value=([101, 102], torch.tensor([101, 102]))
|
||||
)
|
||||
@@ -210,6 +214,7 @@ def test_execute_prefill_packs_ragged_prompts_and_selects_last_logits():
|
||||
model_args, model_kwargs = executor.model.call_args
|
||||
assert model_args[0].tolist() == [11, 12, 21, 22, 23, 24]
|
||||
assert model_kwargs["position_ids"].tolist() == [1, 2, 1, 2, 3, 4]
|
||||
assert model_kwargs["logits_positions"].tolist() == [1, 5]
|
||||
executor.task_cache.bind.assert_called_once_with(
|
||||
["a", "b"], executor._workspace, start_pos=1
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user