feat: Reduce reference model memory with with parallel logprob computation#608

Open

gitlost-murali wants to merge 12 commits intometa-pytorch:mainfrom

gitlost-murali:optimize-ref-model-usage

Commits on Dec 3, 2025

feat: Add parallel log-prob computation for vocab-sharded tensors in reference model
gitlost-murali
committed
refactor: use util function
gitlost-murali
committed
refactor: remove redundant code and reuse compute_logprobs function
gitlost-murali
committed
test: add test to verify correctness with compute_logprobs
gitlost-murali
committed
chore: make routing explicit
gitlost-murali
committed
style: clean up inline comments in parallel_logprobs
gitlost-murali
committed
refactor: move compute_logprobs_parallel into ops alongside compute_logprobs
gitlost-murali
committed
refactor: simplify and refactor logprobs parallel method
gitlost-murali
committed
test: move parallel logprobs test to test_ops.py reflecting folder structure
gitlost-murali
committed
chore: merge single use declarations
gitlost-murali
committed
fix: safely handle edgecase of uneven vocab shards in tensor-parallel logprobs
gitlost-murali
committed
fix: fix compute_logprobs_parallel import
gitlost-murali
committed