Skip to content

perf: streamline greedy sampling and skip identity penalties - #1643

Draft
sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:sampling-optim
Draft

sufubao wants to merge 1 commit into
ModelTC:mainfrom
sufubao:sampling-optim

Conversation

@sufubao

@sufubao sufubao commented Oct 20, 2026

Copy link
Copy Markdown
Collaborator

采样参数未启用 presence、frequency、repetition 惩罚时,跳过不改变 logits 的全词表惩罚计算;EOS 屏蔽和长度惩罚仍正常执行。

为无输出约束的 FP32 greedy 采样增加分块归约,直接计算 token ID 和 logprob,避免生成完整 softmax 缓冲;完整词表的 draft argmax 使用独立的两阶段归约。保留其他输入和采样模式的原有路径。

这些改动从 ReplaySSM 分支独立拆出,不依赖线性注意力或 SSM 状态管理。

验证:test_greedy_sampling.py 10 项通过,覆盖 NaN/Inf、并列最大值、屏蔽词表分块、EOS、惩罚和随机采样;git diff --check 通过。未测量性能收益。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant