def chunk_forward(input_chunk, weight, bias=None, temperature=1.0):
        """Forward pass computation for a single chunk without explicit reshaping."""
        # Directly compute logits via batched matrix multiplication: [B, T, H] @ [H, V] -> [B, T, V]
        logits = torch.matmul(input_chunk, weight.t())
        if bias is not None:
            logits = logits + bias  # Broadcasts bias to [B, T, V]
        if temperature != 1.0:
            logits = logits / temperature

        # Compute log probabilities using softmax over the last dimension
        log_probs = F.log_softmax(logits.float(), dim=-1)

        return log_probs, logits
Forward pass computation for a single chunk without explicit reshaping.
src/liger_kernel/chunked_loss/fused_linear_ppo.py