Prefill throughput in batched rubric grading and memory footprint on long medical contexts
Hi Junying and FreedomAI team,
Releasing an open 8B medical rubric grader calibrated for OnePO policy optimization is a great contribution, especially formatting multi-criteria evaluation into a single structured pass.
Looking at the batch scoring workload in score.py and the RL rollout reward loop:
Severe prefill-to-decode asymmetry during rollout scoring:
Rubric grading on 10,000-character conversations creates sequences of 2,500 to 3,500 prompt tokens, while the target output is just a compact JSON boolean array of 10 to 30 tokens.
In an active RL or OnePO training loop, scoring hundreds of policy rollouts across large batches means the 8B grader is almost permanently locked in memory-heavy prefill. On standard dense attention, loading multi-head KV caches across 3k+ contexts at batch sizes of 8-16 quickly saturates GPU memory bandwidth.152k vocabulary tax on a grading task:
On a Qwen3-8B backbone (hidden dimension 4,096), a 152,000 token vocabulary consumes ~622M parameters in the embedding matrix.
For a dedicated rubric grader whose generation space is effectively constrained to JSON brackets and boolean values, spending over 600M parameters (nearly 8% of the entire 8B budget) on static token tables creates unnecessary parameter overhead. Decoupling the input embedding via low-rank projection reclaims hundreds of millions of weights to invest directly into deeper reasoning layers.Associative recall across extended clinical histories:
Checking negative criteria (such as diagnosing without patient-specific evidence) across long multi-turn transcripts requires strict relational binding across the full context.
In an open architecture project called Maba (101M reference model: https://huggingface.co/AndrewThompson1233/maba-v1-architecture), we handle long-context sequence evaluation using hybrid linear recurrence (75% GDN-2 / 25% GQA):
GDN-2 updates an associative state in fixed O(1) memory, keeping 75% of layer states constant regardless of transcript length.
This slashes the KV cache footprint by ~76%, allowing significantly larger batch sizes during batched reward scoring in score.py without running out of VRAM.
In your OnePO training runs, what batch sizes and throughput were you averaging during the rubric reward generation phase on 10k-character transcripts?
Best,
Andrew