| """The packed-batch record the model reads (a copy of model/loader.py's Batch, without the training loader).""" | |
| from __future__ import annotations | |
| from dataclasses import dataclass | |
| from torch import Tensor | |
| class Batch: | |
| """One packed microbatch on the GPU: every sequence is its own segment, positions restart at 0.""" | |
| x: Tensor # (N,) int64 inputs | |
| y: Tensor # (N,) int64 targets (ignored by inference) | |
| pos: Tensor # (N,) int32 position inside its segment (RoPE) | |
| cu: Tensor # (S+1,) int32 segment offsets | |
| cu64: Tensor # same offsets, int64 (fla) | |
| cu_host: Tensor # same offsets, int64 on the CPU | |
| row: int # longest segment | |
| segs: int # segment count | |
| weight: Tensor | None = None # (N,) loss weight per target; None = plain mean | |
| def tokens(self) -> int: | |
| return self.x.numel() | |