h3-blackwell-runtime/src/h3_blackwell_runtime/h3_fusion.py

"""Triton fusion for H3's segmented BF16 modulation and residual gates."""

from __future__ import annotations

import torch
import triton
import triton.language as tl


_TILE = 1024
_segment_cache_key = None
_segment_cache_value = None


@triton.jit
def _round_bf16_fp32(value):
    """Apply round-to-nearest-even BF16 precision while retaining FP32."""
    bits = value.to(tl.int32, bitcast=True)
    bits = bits + 0x7FFF + ((bits >> 16) & 1)
    return (bits & -65536).to(tl.float32, bitcast=True)


@triton.jit
def _modulate_kernel(
    x_ptr,
    scale_ptr,
    shift_ptr,
    row_index_ptr,
    tokens,
    width,
    sx_t,
    sx_d,
    ss_t,
    ss_d,
    sh_t,
    sh_d,
    BLOCK: tl.constexpr,
):
    row = tl.program_id(0)
    columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
    valid = (row < tokens) & (columns < width)
    table_row = tl.load(row_index_ptr + row)
    x = tl.load(x_ptr + row * sx_t + columns * sx_d, mask=valid).to(tl.float32)
    scale = tl.load(
        scale_ptr + table_row * ss_t + columns * ss_d, mask=valid,
    ).to(tl.float32)
    shift = tl.load(
        shift_ptr + table_row * sh_t + columns * sh_d, mask=valid,
    ).to(tl.float32)

    scale = _round_bf16_fp32(scale)
    shift = _round_bf16_fp32(shift)
    multiplied = _round_bf16_fp32(x * _round_bf16_fp32(1.0 + scale))
    result = _round_bf16_fp32(multiplied + shift)
    tl.store(x_ptr + row * sx_t + columns * sx_d, result.to(tl.bfloat16), mask=valid)


@triton.jit
def _gate_add_kernel(
    residual_ptr,
    gate_ptr,
    update_ptr,
    row_index_ptr,
    tokens,
    width,
    sr_t,
    sr_d,
    sg_t,
    sg_d,
    su_t,
    su_d,
    BLOCK: tl.constexpr,
):
    row = tl.program_id(0)
    columns = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)
    valid = (row < tokens) & (columns < width)
    table_row = tl.load(row_index_ptr + row)
    residual = tl.load(
        residual_ptr + row * sr_t + columns * sr_d, mask=valid,
    ).to(tl.float32)
    update = tl.load(
        update_ptr + row * su_t + columns * su_d, mask=valid,
    ).to(tl.float32)
    gate = tl.load(
        gate_ptr + table_row * sg_t + columns * sg_d, mask=valid,
    ).to(tl.float32)

    result = residual + update * _round_bf16_fp32(gate)
    tl.store(
        residual_ptr + row * sr_t + columns * sr_d,
        result.to(tl.bfloat16),
        mask=valid,
    )


def segment_index(
    tokens: int,
    segments: list[tuple[int, int, int]],
    device: torch.device,
) -> torch.Tensor:
    """Return one cached row-to-modulation-table lookup for a packed layout."""
    global _segment_cache_key, _segment_cache_value
    normalized = tuple((int(start), int(stop), int(row)) for start, stop, row in segments)
    key = (str(device), int(tokens), normalized)
    if key == _segment_cache_key:
        return _segment_cache_value
    host = torch.empty(tokens, dtype=torch.int32)
    cursor = 0
    for start, stop, row in normalized:
        if start != cursor or stop < start or stop > tokens or row < 0:
            raise ValueError("segments must be ordered, contiguous, and in range")
        host[start:stop] = row
        cursor = stop
    if cursor != tokens:
        raise ValueError("segments must cover every packed token")
    _segment_cache_key = key
    _segment_cache_value = host.to(device=device)
    return _segment_cache_value


def _validate_activation(tensor: torch.Tensor, name: str) -> None:
    if tensor.ndim != 2:
        raise ValueError(f"{name} must have shape [tokens, hidden]")
    if tensor.device.type != "cuda" or tensor.dtype != torch.bfloat16:
        raise TypeError(f"{name} must be a CUDA BF16 tensor")
    if tensor.stride(1) != 1:
        raise ValueError(f"{name}'s hidden dimension must be contiguous")


def _validate_table(table: torch.Tensor, activation: torch.Tensor, name: str) -> None:
    if table.ndim != 2 or table.shape[1] != activation.shape[1]:
        raise ValueError(f"{name} must have shape [rows, {activation.shape[1]}]")
    if table.device != activation.device or table.dtype not in {torch.bfloat16, torch.float32}:
        raise TypeError(f"{name} must be CUDA BF16/FP32 on the activation device")
    if table.stride(1) != 1:
        raise ValueError(f"{name}'s hidden dimension must be contiguous")


def fused_modulate_(
    activation: torch.Tensor,
    shift: torch.Tensor,
    scale: torch.Tensor,
    row_index: torch.Tensor,
) -> torch.Tensor:
    """Apply segmented scale/shift in place with eager-equivalent BF16 rounding."""
    _validate_activation(activation, "activation")
    _validate_table(shift, activation, "shift")
    _validate_table(scale, activation, "scale")
    if row_index.shape != (activation.shape[0],) or row_index.dtype != torch.int32:
        raise TypeError("row_index must be int32 with one entry per token")
    if row_index.device != activation.device or not row_index.is_contiguous():
        raise TypeError("row_index must be contiguous on the activation device")
    grid = (activation.shape[0], triton.cdiv(activation.shape[1], _TILE))
    _modulate_kernel[grid](
        activation,
        scale,
        shift,
        row_index,
        activation.shape[0],
        activation.shape[1],
        activation.stride(0),
        activation.stride(1),
        scale.stride(0),
        scale.stride(1),
        shift.stride(0),
        shift.stride(1),
        BLOCK=_TILE,
        num_warps=4,
    )
    return activation


def fused_gate_add_(
    residual: torch.Tensor,
    update: torch.Tensor,
    gate: torch.Tensor,
    row_index: torch.Tensor,
) -> torch.Tensor:
    """Apply the segmented residual gate in place with addcmul-equivalent math."""
    _validate_activation(residual, "residual")
    _validate_activation(update, "update")
    if residual.shape != update.shape or residual.device != update.device:
        raise ValueError("residual and update must share shape and device")
    _validate_table(gate, residual, "gate")
    if row_index.shape != (residual.shape[0],) or row_index.dtype != torch.int32:
        raise TypeError("row_index must be int32 with one entry per token")
    if row_index.device != residual.device or not row_index.is_contiguous():
        raise TypeError("row_index must be contiguous on the residual device")
    grid = (residual.shape[0], triton.cdiv(residual.shape[1], _TILE))
    _gate_add_kernel[grid](
        residual,
        gate,
        update,
        row_index,
        residual.shape[0],
        residual.shape[1],
        residual.stride(0),
        residual.stride(1),
        gate.stride(0),
        gate.stride(1),
        update.stride(0),
        update.stride(1),
        BLOCK=_TILE,
        num_warps=4,
    )
    return residual