370 lines
11 KiB
Python
370 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Fused indexed modulation with FP32 accumulation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import torch
|
|
from vllm.triton_utils import tl, triton
|
|
|
|
# Ascend CANN limits the launch coreDim to uint16.
|
|
_MAX_1D_GRID_SIZE = 65535
|
|
|
|
|
|
def _iter_row_chunks(rows: int):
|
|
for row_offset in range(0, rows, _MAX_1D_GRID_SIZE):
|
|
yield row_offset, min(rows - row_offset, _MAX_1D_GRID_SIZE)
|
|
|
|
|
|
def _launch_row_chunks(kernel, rows: int, device_type: str, *args, **kwargs) -> None:
|
|
if device_type != "npu":
|
|
kernel[(rows,)](*args, 0, **kwargs)
|
|
return
|
|
for row_offset, chunk_rows in _iter_row_chunks(rows):
|
|
kernel[(chunk_rows,)](*args, row_offset, **kwargs)
|
|
|
|
|
|
def _use_hopper_bf16_affine_semantics(x: torch.Tensor) -> bool:
|
|
"""Use the unfused AdaLN dtype boundaries on Hopper.
|
|
|
|
The fused FP32 modulation path is stable on Blackwell, but the SM90
|
|
kernel can accumulate enough extra precision in the repeated RMSNorm and
|
|
AdaLN operations to move the 50-step I2VA result below its CI threshold.
|
|
Keep the existing path on other architectures until their numerical
|
|
baselines are independently validated.
|
|
"""
|
|
return x.device.type == "cuda" and torch.cuda.get_device_capability(x.device)[0] == 9
|
|
|
|
|
|
@triton.jit
|
|
def _indexed_scale_shift_kernel(
|
|
output_ptr,
|
|
x_ptr,
|
|
shift_ptr,
|
|
scale_ptr,
|
|
indices_ptr,
|
|
hidden_size,
|
|
stride_x_row,
|
|
stride_shift_row,
|
|
stride_scale_row,
|
|
stride_indices,
|
|
row_offset,
|
|
block_n: tl.constexpr,
|
|
):
|
|
row = tl.program_id(0) + row_offset
|
|
columns = tl.arange(0, block_n)
|
|
mask = columns < hidden_size
|
|
index = tl.load(indices_ptr + row * stride_indices)
|
|
|
|
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
tl.store(
|
|
output_ptr + row * stride_x_row + columns,
|
|
x * (1.0 + scale) + shift,
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _indexed_gate_kernel(
|
|
output_ptr,
|
|
x_ptr,
|
|
gate_ptr,
|
|
other_ptr,
|
|
indices_ptr,
|
|
hidden_size,
|
|
stride_output_row,
|
|
stride_x_row,
|
|
stride_gate_row,
|
|
stride_other_row,
|
|
stride_indices,
|
|
row_offset,
|
|
block_n: tl.constexpr,
|
|
):
|
|
row = tl.program_id(0) + row_offset
|
|
columns = tl.arange(0, block_n)
|
|
mask = columns < hidden_size
|
|
index = tl.load(indices_ptr + row * stride_indices)
|
|
|
|
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
other = tl.load(other_ptr + row * stride_other_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
|
|
tl.store(
|
|
output_ptr + row * stride_output_row + columns,
|
|
x + gate * other,
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _rms_norm_indexed_scale_shift_kernel(
|
|
output_ptr,
|
|
x_ptr,
|
|
weight_ptr,
|
|
shift_ptr,
|
|
scale_ptr,
|
|
indices_ptr,
|
|
hidden_size: tl.constexpr,
|
|
eps: tl.constexpr,
|
|
stride_x_row,
|
|
stride_shift_row,
|
|
stride_scale_row,
|
|
stride_indices,
|
|
row_offset,
|
|
use_bf16_affine: tl.constexpr,
|
|
block_n: tl.constexpr,
|
|
):
|
|
row = tl.program_id(0) + row_offset
|
|
columns = tl.arange(0, block_n)
|
|
mask = columns < hidden_size
|
|
index = tl.load(indices_ptr + row * stride_indices)
|
|
|
|
x = tl.load(x_ptr + row * stride_x_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
|
|
variance = tl.sum(x * x, axis=0) / hidden_size
|
|
if use_bf16_affine:
|
|
# Match the unfused RMSNorm output consumed by AdaLN on Hopper.
|
|
normalized = (x * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
|
|
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
|
|
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
|
|
output = normalized * (scale + 1.0) + shift
|
|
else:
|
|
normalized = x * tl.rsqrt(variance + eps) * weight
|
|
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
output = normalized * (1.0 + scale) + shift
|
|
tl.store(
|
|
output_ptr + row * hidden_size + columns,
|
|
output,
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
@triton.jit
|
|
def _indexed_gate_rms_norm_scale_shift_kernel(
|
|
residual_out_ptr,
|
|
modulated_out_ptr,
|
|
residual_ptr,
|
|
gate_ptr,
|
|
branch_ptr,
|
|
weight_ptr,
|
|
shift_ptr,
|
|
scale_ptr,
|
|
indices_ptr,
|
|
hidden_size: tl.constexpr,
|
|
eps: tl.constexpr,
|
|
stride_residual_row,
|
|
stride_gate_row,
|
|
stride_branch_row,
|
|
stride_shift_row,
|
|
stride_scale_row,
|
|
stride_indices,
|
|
row_offset,
|
|
use_bf16_affine: tl.constexpr,
|
|
block_n: tl.constexpr,
|
|
):
|
|
row = tl.program_id(0) + row_offset
|
|
columns = tl.arange(0, block_n)
|
|
mask = columns < hidden_size
|
|
index = tl.load(indices_ptr + row * stride_indices)
|
|
|
|
residual = tl.load(residual_ptr + row * stride_residual_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
gate = tl.load(gate_ptr + index * stride_gate_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
branch = tl.load(branch_ptr + row * stride_branch_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
# Match the BF16 residual value consumed by the unfused RMSNorm path.
|
|
updated = (residual + gate * branch).to(tl.bfloat16).to(tl.float32)
|
|
tl.store(residual_out_ptr + row * hidden_size + columns, updated, mask=mask)
|
|
|
|
weight = tl.load(weight_ptr + columns, mask=mask, other=0.0).to(tl.float32)
|
|
variance = tl.sum(updated * updated, axis=0) / hidden_size
|
|
if use_bf16_affine:
|
|
# Match the unfused RMSNorm output consumed by AdaLN on Hopper.
|
|
normalized = (updated * tl.rsqrt(variance + eps) * weight).to(tl.bfloat16)
|
|
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0)
|
|
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0)
|
|
output = normalized * (scale + 1.0) + shift
|
|
else:
|
|
normalized = updated * tl.rsqrt(variance + eps) * weight
|
|
shift = tl.load(shift_ptr + index * stride_shift_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
scale = tl.load(scale_ptr + index * stride_scale_row + columns, mask=mask, other=0.0).to(tl.float32)
|
|
output = normalized * (1.0 + scale) + shift
|
|
tl.store(
|
|
modulated_out_ptr + row * hidden_size + columns,
|
|
output,
|
|
mask=mask,
|
|
)
|
|
|
|
|
|
def indexed_scale_shift_(
|
|
x: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
indices: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Apply indexed scale/shift in-place to a disposable contiguous input."""
|
|
if x.is_cpu:
|
|
x.copy_((x * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(x.dtype))
|
|
return x
|
|
rows, hidden_size = x.shape
|
|
if rows == 0:
|
|
return x
|
|
_launch_row_chunks(
|
|
_indexed_scale_shift_kernel,
|
|
rows,
|
|
x.device.type,
|
|
x,
|
|
x,
|
|
shift,
|
|
scale,
|
|
indices,
|
|
hidden_size,
|
|
x.stride(0),
|
|
shift.stride(0),
|
|
scale.stride(0),
|
|
indices.stride(0),
|
|
block_n=triton.next_power_of_2(hidden_size),
|
|
num_warps=8,
|
|
)
|
|
return x
|
|
|
|
|
|
def indexed_gate(
|
|
x: torch.Tensor,
|
|
gate: torch.Tensor,
|
|
other: torch.Tensor,
|
|
indices: torch.Tensor,
|
|
) -> torch.Tensor:
|
|
"""Return ``x + gate[indices] * other`` without indexed temporaries."""
|
|
if x.is_cpu:
|
|
return (x + gate.index_select(0, indices) * other).to(x.dtype)
|
|
output = torch.empty_like(x)
|
|
rows, hidden_size = x.shape
|
|
if rows == 0:
|
|
return output
|
|
_launch_row_chunks(
|
|
_indexed_gate_kernel,
|
|
rows,
|
|
x.device.type,
|
|
output,
|
|
x,
|
|
gate,
|
|
other,
|
|
indices,
|
|
hidden_size,
|
|
output.stride(0),
|
|
x.stride(0),
|
|
gate.stride(0),
|
|
other.stride(0),
|
|
indices.stride(0),
|
|
block_n=triton.next_power_of_2(hidden_size),
|
|
num_warps=8,
|
|
)
|
|
return output
|
|
|
|
|
|
def rms_norm_indexed_scale_shift(
|
|
x: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
indices: torch.Tensor,
|
|
eps: float,
|
|
) -> torch.Tensor:
|
|
"""Fuse RMSNorm with an indexed affine transform."""
|
|
if x.is_cpu:
|
|
input_dtype = x.dtype
|
|
normalized = x.float()
|
|
variance = normalized.pow(2).mean(-1, keepdim=True)
|
|
normalized = normalized * torch.rsqrt(variance + eps)
|
|
normalized = (weight.float() * normalized).to(input_dtype)
|
|
return (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(input_dtype)
|
|
output = torch.empty_like(x)
|
|
rows, hidden_size = x.shape
|
|
if rows:
|
|
_launch_row_chunks(
|
|
_rms_norm_indexed_scale_shift_kernel,
|
|
rows,
|
|
x.device.type,
|
|
output,
|
|
x,
|
|
weight,
|
|
shift,
|
|
scale,
|
|
indices,
|
|
hidden_size,
|
|
eps,
|
|
x.stride(0),
|
|
shift.stride(0),
|
|
scale.stride(0),
|
|
indices.stride(0),
|
|
use_bf16_affine=_use_hopper_bf16_affine_semantics(x),
|
|
block_n=triton.next_power_of_2(hidden_size),
|
|
num_warps=8,
|
|
)
|
|
return output
|
|
|
|
|
|
def indexed_gate_rms_norm_scale_shift(
|
|
residual: torch.Tensor,
|
|
gate: torch.Tensor,
|
|
branch: torch.Tensor,
|
|
weight: torch.Tensor,
|
|
shift: torch.Tensor,
|
|
scale: torch.Tensor,
|
|
indices: torch.Tensor,
|
|
eps: float,
|
|
) -> tuple[torch.Tensor, torch.Tensor]:
|
|
"""Fuse gated residual, the following RMSNorm, and AdaLN affine."""
|
|
if residual.is_cpu:
|
|
input_dtype = residual.dtype
|
|
residual_out = (residual + gate.index_select(0, indices) * branch).to(input_dtype)
|
|
normalized = residual_out.float()
|
|
variance = normalized.pow(2).mean(-1, keepdim=True)
|
|
normalized = normalized * torch.rsqrt(variance + eps)
|
|
normalized = (weight.float() * normalized).to(input_dtype)
|
|
modulated_out = (normalized * (1.0 + scale.index_select(0, indices)) + shift.index_select(0, indices)).to(
|
|
input_dtype
|
|
)
|
|
return residual_out, modulated_out
|
|
residual_out = torch.empty_like(residual)
|
|
modulated_out = torch.empty_like(residual)
|
|
rows, hidden_size = residual.shape
|
|
if rows:
|
|
_launch_row_chunks(
|
|
_indexed_gate_rms_norm_scale_shift_kernel,
|
|
rows,
|
|
residual.device.type,
|
|
residual_out,
|
|
modulated_out,
|
|
residual,
|
|
gate,
|
|
branch,
|
|
weight,
|
|
shift,
|
|
scale,
|
|
indices,
|
|
hidden_size,
|
|
eps,
|
|
residual.stride(0),
|
|
gate.stride(0),
|
|
branch.stride(0),
|
|
shift.stride(0),
|
|
scale.stride(0),
|
|
indices.stride(0),
|
|
use_bf16_affine=_use_hopper_bf16_affine_semantics(residual),
|
|
block_n=triton.next_power_of_2(hidden_size),
|
|
num_warps=8,
|
|
)
|
|
return residual_out, modulated_out
|
|
|
|
|
|
__all__ = [
|
|
"indexed_gate",
|
|
"indexed_gate_rms_norm_scale_shift",
|
|
"indexed_scale_shift_",
|
|
"rms_norm_indexed_scale_shift",
|
|
]
|