qknorm_rope ¶

Fused per-head RMSNorm and partial rotary embedding for MiniMax H3.

Functions:¶

fastvideo.models.dits.minimax_h3_fusions.qknorm_rope.fused_qknorm_rope ¶

fused_qknorm_rope(x: Tensor, weight: Tensor, cos: Tensor, sin: Tensor, eps: float) -> Tensor

Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.

RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the final store. Triton's reduction order and the absence of eager's BF16 intermediate materializations can produce small, expected rounding drift.

Row offsets are computed in int64, so inputs beyond 2**31 total elements (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address correctly.

Source code in fastvideo/models/dits/minimax_h3_fusions/qknorm_rope.py

def fused_qknorm_rope(
    x: torch.Tensor,
    weight: torch.Tensor,
    cos: torch.Tensor,
    sin: torch.Tensor,
    eps: float,
) -> torch.Tensor:
    """Run per-head RMSNorm and partial RoPE in one Sol-Engine-style kernel.

    RMSNorm reduction and RoPE arithmetic stay in FP32 registers until the
    final store. Triton's reduction order and the absence of eager's BF16
    intermediate materializations can produce small, expected rounding drift.

    Row offsets are computed in int64, so inputs beyond 2**31 total elements
    (about 300k tokens per rank at H3's 56 heads x 128 head_dim) address
    correctly.
    """
    if torch.is_grad_enabled() and any(
            isinstance(tensor, torch.Tensor) and tensor.requires_grad for tensor in (x, weight, cos, sin)):
        raise RuntimeError("fused_qknorm_rope is inference-only and does not implement autograd")
    if torch.compiler.is_compiling():
        return torch.ops.fastvideo._minimax_h3_qknorm_rope(x, weight, cos, sin, eps)
    return _fused_qknorm_rope_impl(x, weight, cos, sin, eps)