Function pad_to_tma_alignment

src/fused_mm_sampling/testing.py:75–91

Pad D to 16-byte alignment for TMA on SM 90+ (H100, B200, etc.). TMA requires the innermost tensor dimension to be 16-byte aligned. For bf16 (2 bytes), D must be a multiple of 8. After shift_logits_negative adds a bias column (D = hidden_size + 1), D=11 is not aligned. Zero-padding

(
    weights: torch.Tensor, hidden_states: torch.Tensor
)

def pad_to_tma_alignment(
    weights: torch.Tensor, hidden_states: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
    """Pad D to 16-byte alignment for TMA on SM 90+ (H100, B200, etc.).

    TMA requires the innermost tensor dimension to be 16-byte aligned.
    For bf16 (2 bytes), D must be a multiple of 8. After shift_logits_negative
    adds a bias column (D = hidden_size + 1), D=11 is not aligned. Zero-padding
    extra columns preserves logits (they contribute nothing to the matmul).
    """
    d = weights.shape[1]
    aligned_d = (d + 7) & ~7  # next multiple of 8 bf16 elements = 16 bytes
    if aligned_d > d:
        pad = aligned_d - d
        weights = torch.nn.functional.pad(weights, (0, pad))
        hidden_states = torch.nn.functional.pad(hidden_states, (0, pad))
    return weights, hidden_states


def shard_weights(weights: torch.Tensor, tp: TPInfo) -> torch.Tensor: