Where it bites
oxmega's stacked models produce [R, B, H, h] intermediates (replicas × batch × history × hidden) and must reshape to rank 3 before every matmul (plan_matmul supports rank 2/3 only); the Attention pooling does four reshapes/permutes per forward for that reason alone, each a copy on the GPU path.
Proposal
Generalise plan_matmul to N leading batch dims with NumPy broadcasting (the zero-stride trick already used for batch 1 extends directly), and add a minimal einsum for the common contractions ('rbhd,rdo->rbho'). Roadmap already lists 'rank-4+ matmul broadcasting and einsum'; this is the concrete consumer.
Where it bites
oxmega's stacked models produce [R, B, H, h] intermediates (replicas × batch × history × hidden) and must reshape to rank 3 before every matmul (
plan_matmulsupports rank 2/3 only); the Attention pooling does four reshapes/permutes per forward for that reason alone, each a copy on the GPU path.Proposal
Generalise
plan_matmulto N leading batch dims with NumPy broadcasting (the zero-stride trick already used for batch 1 extends directly), and add a minimaleinsumfor the common contractions ('rbhd,rdo->rbho'). Roadmap already lists 'rank-4+ matmul broadcasting and einsum'; this is the concrete consumer.