Hey there:
1). To do BF16 would require Eight B200s, which no one could rent me at the time. Doing it against FP8 is fine as FP8 KLD divergence versus BF16 will literally be 00.001 or something to the effect. So doing it against BF16 is going to show all models unilaterally worse. At some point in the future, I WILL do BF16, but it's not needed. BF16 won't show any model to be better, but all models to be across the board worse, equally.
2). I load the full vocab on VRAM. Look at it again. I do exactly what Tubro does for EXL3. Full vocab loaded. This is why it takes soo much VRAM. We do Every position, Every Window, for FULL VOCAB.

Trust me, I spent 3 months doing this. It's rock solid. I'm happy to take all feedback and thank you for reviewing the code, but it does do full vocab.

To Triple check myself, I asked a colleague, and all three frontier models:
I reviewed my own code, my friend reviewed it, and we asked Opus, Gemini, and GPT5.4 and all of them agree with certainty, I load the full vocab.

GPU model running:
prompt_hidden_states = hidden_states[offset : offset + num_logits]

logits = self.model.compute_logits(prompt_hidden_states)

Compute_Logits calls this in qwen3.5:
def compute_logits(

self,

hidden_states: torch.Tensor,

) -> torch.Tensor | None:

return self.logits_processor(self.lm_head, hidden_states)

Logit processing is then done here:
def _get_logits(

self,

hidden_states: torch.Tensor,

lm_head: VocabParallelEmbedding,

embedding_bias: torch.Tensor | None,

) -> torch.Tensor | None:

logits = lm_head.quant_method.apply(lm_head, hidden_states, bias=embedding_bias)

logits = self._gather_logits(logits)

if logits is not None:

logits = logits[..., : self.org_vocab_size]

return logits

In phase 2, here is where again, full vocab is hammered home:
if is_kld_mode:

from safetensors.torch import safe_open

with safe_open(

request.reference_logits_path,

framework="pt",

device=str(self.device),

) as f:

ref_logits_full = f.get_tensor(request.reference_logits_key).to(

self.device

)

ref_logits = ref_logits_full[start_idx : start_idx + num_logits]

vs = min(logits.shape[-1], ref_logits.shape[-1])

log_probs_model = F.log_softmax(logits[..., :vs].float(), dim=-1)

log_probs_ref = F.log_softmax(ref_logits[..., :vs].float(), dim=-1)

kld_per_pos = F.kl_div(

log_probs_model,

log_probs_ref,

reduction="none",

log_target=True,

).sum(dim=-1)

kld_sum = kld_per_pos.sum().item()

kld_count = kld_per_pos.numel()

Every Position, Every Window, FULL VOCAB:
for start_idx in range(0, num_tokens - context_length + stride, stride):

end_idx = start_idx + context_length

if end_idx > num_tokens:

break

window_tokens = tokens[start_idx:end_idx]