[Kernel][HY V4] Add Triton iHC pre/post fallback (#55059)
Signed-off-by: linitra24 <renshuang.zhou@daocloud.io> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
This commit is contained in:
linitra24
@@ -0,0 +1,199 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Benchmark the HY V4 Triton iHC pre/post kernels against eager PyTorch."""
|
||||
|
||||
import os
|
||||
import subprocess
|
||||
from functools import partial
|
||||
from importlib.metadata import version
|
||||
from statistics import median
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
import vllm
|
||||
from vllm.models.hy_v4.nvidia.triton_ihc import (
|
||||
triton_ihc_post,
|
||||
triton_ihc_pre,
|
||||
)
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
|
||||
def eager_pre(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
base: torch.Tensor,
|
||||
magnitude: float,
|
||||
hc_eps: float,
|
||||
norm_eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
num_tokens, hc_mult, hidden_size = x.shape
|
||||
x_flat = x.flatten(1).float()
|
||||
reciprocal_rms = torch.rsqrt(x_flat.square().mean(-1, keepdim=True) + norm_eps)
|
||||
mixes = F.linear(x_flat, weight) * reciprocal_rms
|
||||
pre = torch.sigmoid(mixes[:, :hc_mult] * scale[0] + base[:hc_mult])
|
||||
post = torch.sigmoid(mixes[:, hc_mult:] * scale[1] + base[hc_mult:])
|
||||
pre = pre + hc_eps
|
||||
post = magnitude * post + hc_eps
|
||||
output = torch.sum(pre.unsqueeze(-1) * x.float(), dim=1)
|
||||
return output.to(x.dtype).reshape(num_tokens, hidden_size), post
|
||||
|
||||
|
||||
def eager_post(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return (post.float().unsqueeze(-1) * x.float().unsqueeze(-2) + residual.float()).to(
|
||||
x.dtype
|
||||
)
|
||||
|
||||
|
||||
def _timer(method: str):
|
||||
if method == "cupti":
|
||||
from flashinfer.testing import bench_gpu_time_with_cupti
|
||||
|
||||
return partial(
|
||||
bench_gpu_time_with_cupti,
|
||||
use_cuda_graph=True,
|
||||
cold_l2_cache=True,
|
||||
)
|
||||
if method == "cudagraph":
|
||||
from flashinfer.testing import bench_gpu_time_with_cudagraph
|
||||
|
||||
return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True)
|
||||
raise ValueError(f"unknown timing method: {method}")
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def run_benchmark(
|
||||
token_counts: list[int],
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
method: str,
|
||||
) -> None:
|
||||
hc_mult = 4
|
||||
device = torch.device("cuda")
|
||||
torch.manual_seed(0)
|
||||
weight = torch.randn(
|
||||
2 * hc_mult,
|
||||
hc_mult * hidden_size,
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
scale = torch.randn(2, device=device, dtype=torch.float32) * 0.01
|
||||
base = torch.randn(2 * hc_mult, device=device, dtype=torch.float32)
|
||||
timer = _timer(method)
|
||||
|
||||
properties = torch.cuda.get_device_properties(device)
|
||||
git_branch = subprocess.run(
|
||||
["git", "branch", "--show-current"],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stdout.strip()
|
||||
git_commit = subprocess.run(
|
||||
["git", "rev-parse", "--short", "HEAD"],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
).stdout.strip()
|
||||
print(f"device: {properties.name}")
|
||||
print(f"branch: {git_branch}; commit: {git_commit}")
|
||||
print(
|
||||
f"vllm: {vllm.__version__}; torch: {torch.__version__}; "
|
||||
f"CUDA: {torch.version.cuda}"
|
||||
)
|
||||
print(
|
||||
f"triton: {triton.__version__}; flashinfer: {version('flashinfer-python')}; "
|
||||
f"dtype: {dtype}; method: {method}; cache: cold L2"
|
||||
)
|
||||
print(
|
||||
"env: VLLM_ENABLE_HPC_OPS="
|
||||
f"{os.getenv('VLLM_ENABLE_HPC_OPS', '<unset>')}; "
|
||||
"VLLM_BATCH_INVARIANT="
|
||||
f"{os.getenv('VLLM_BATCH_INVARIANT', '<unset>')}"
|
||||
)
|
||||
print(f"hidden_size: {hidden_size}; hc_mult: {hc_mult}")
|
||||
print(
|
||||
f"{'tokens':>8} {'op':>6} {'eager (us)':>12} "
|
||||
f"{'triton (us)':>12} {'speedup':>9} {'GiB':>8} "
|
||||
f"{'eager GB/s':>12} {'triton GB/s':>13}"
|
||||
)
|
||||
|
||||
for num_tokens in token_counts:
|
||||
x = torch.randn(
|
||||
num_tokens,
|
||||
hc_mult,
|
||||
hidden_size,
|
||||
device=device,
|
||||
dtype=dtype,
|
||||
)
|
||||
block_output = torch.randn(num_tokens, hidden_size, device=device, dtype=dtype)
|
||||
residual = torch.randn_like(x)
|
||||
eager_output, post = eager_pre(x, weight, scale, base, 2.0, 1e-6, 1e-5)
|
||||
triton_output, triton_post = triton_ihc_pre(
|
||||
x, weight, scale, base, 2.0, 1e-6, 1e-5
|
||||
)
|
||||
torch.testing.assert_close(triton_output, eager_output, atol=2e-2, rtol=1e-2)
|
||||
torch.testing.assert_close(triton_post, post, atol=2e-5, rtol=1e-5)
|
||||
torch.testing.assert_close(
|
||||
triton_ihc_post(block_output, residual, post),
|
||||
eager_post(block_output, residual, post),
|
||||
atol=0,
|
||||
rtol=0,
|
||||
)
|
||||
|
||||
benchmarks = (
|
||||
(
|
||||
"pre",
|
||||
partial(eager_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
|
||||
partial(triton_ihc_pre, x, weight, scale, base, 2.0, 1e-6, 1e-5),
|
||||
x.nbytes
|
||||
+ weight.nbytes
|
||||
+ scale.nbytes
|
||||
+ base.nbytes
|
||||
+ eager_output.nbytes
|
||||
+ post.nbytes,
|
||||
),
|
||||
(
|
||||
"post",
|
||||
partial(eager_post, block_output, residual, post),
|
||||
partial(triton_ihc_post, block_output, residual, post),
|
||||
block_output.nbytes + residual.nbytes + post.nbytes + residual.nbytes,
|
||||
),
|
||||
)
|
||||
for op_name, eager_fn, triton_fn, logical_bytes in benchmarks:
|
||||
eager_us = median(timer(eager_fn)) * 1e3
|
||||
triton_us = median(timer(triton_fn)) * 1e3
|
||||
speedup = eager_us / triton_us
|
||||
logical_gib = logical_bytes / 2**30
|
||||
eager_gbps = logical_bytes / eager_us / 1e3
|
||||
triton_gbps = logical_bytes / triton_us / 1e3
|
||||
print(
|
||||
f"{num_tokens:>8} {op_name:>6} {eager_us:>12.1f} "
|
||||
f"{triton_us:>12.1f} {speedup:>8.2f}x {logical_gib:>8.3f} "
|
||||
f"{eager_gbps:>12.1f} {triton_gbps:>13.1f}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = FlexibleArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"--token-counts",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[1, 2, 4, 8, 16, 64, 256, 1024, 4096, 8192],
|
||||
)
|
||||
parser.add_argument("--hidden-size", type=int, choices=[4096, 6144], default=6144)
|
||||
parser.add_argument("--dtype", choices=["float16", "bfloat16"], default="bfloat16")
|
||||
parser.add_argument("--method", choices=["cupti", "cudagraph"], default="cupti")
|
||||
args = parser.parse_args()
|
||||
run_benchmark(
|
||||
args.token_counts,
|
||||
args.hidden_size,
|
||||
getattr(torch, args.dtype),
|
||||
args.method,
|
||||
)
|
||||
@@ -9,9 +9,8 @@ scatters the result back over the channels (``HYV4HCPostLayer``). The final
|
||||
``HYV4HCHeadLayer`` merges the channels before the model's output norm.
|
||||
|
||||
NOTE: Each of the three steps has an optional single-kernel HPC replacement
|
||||
(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). They are only constructed
|
||||
when the hpc package is installed, ``VLLM_ENABLE_HPC_OPS=1`` and the shape /
|
||||
device constraints hold; otherwise the eager path below runs unchanged.
|
||||
(``HpcIHCPre`` / ``HpcIHCPost`` / ``HpcIHCHead``). Pre and post fall back to
|
||||
in-tree Triton kernels on CUDA when HPC is unavailable, then to the eager path.
|
||||
TODO: port the cross-layer post+pre fusion (``HpcIHCPostPre``) as well; it
|
||||
requires restructuring the decoder-layer forward scheduling.
|
||||
"""
|
||||
@@ -22,6 +21,11 @@ from transformers import PretrainedConfig
|
||||
|
||||
from vllm.model_executor.layers.hpc import HpcIHCHead, HpcIHCPost, HpcIHCPre
|
||||
from vllm.model_executor.layers.linear import ReplicatedLinear
|
||||
from vllm.models.hy_v4.nvidia.triton_ihc import (
|
||||
triton_ihc_post,
|
||||
triton_ihc_pre,
|
||||
triton_ihc_supported,
|
||||
)
|
||||
|
||||
|
||||
class HYV4HCPreLayer(nn.Module):
|
||||
@@ -106,6 +110,16 @@ class HYV4HCPreLayer(nn.Module):
|
||||
"""
|
||||
if self.hpc_op is not None:
|
||||
return self.hpc_op(x)
|
||||
if triton_ihc_supported(x):
|
||||
return triton_ihc_pre(
|
||||
x,
|
||||
self.hc_fn.weight,
|
||||
self.hc_scale,
|
||||
self.hc_base,
|
||||
self.magnitude,
|
||||
self.hc_eps,
|
||||
self.layernorm_epsilon,
|
||||
)
|
||||
|
||||
shape = x.size() # [num_tokens, hc, d]
|
||||
hc = self.hc_mult
|
||||
@@ -175,6 +189,8 @@ class HYV4HCPostLayer(nn.Module):
|
||||
"""
|
||||
if self.hpc_op is not None:
|
||||
return self.hpc_op(x, residual, post)
|
||||
if triton_ihc_supported(x):
|
||||
return triton_ihc_post(x, residual, post)
|
||||
|
||||
dtype = x.dtype
|
||||
x = x.float()
|
||||
|
||||
@@ -0,0 +1,313 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Triton iHC pre/post kernels for HY V4.
|
||||
|
||||
Adapted from the SGLang HY V4 implementation:
|
||||
https://github.com/sgl-project/sglang/pull/36805
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from vllm import envs
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import HAS_TRITON, tl, triton
|
||||
|
||||
_BLOCK_K = 1024
|
||||
_BLOCK_D = 1024
|
||||
|
||||
|
||||
def triton_ihc_supported(x: torch.Tensor) -> bool:
|
||||
"""Return whether the in-tree Triton path can run for this input."""
|
||||
return (
|
||||
HAS_TRITON
|
||||
and current_platform.is_cuda()
|
||||
and x.is_cuda
|
||||
and x.dtype in (torch.float16, torch.bfloat16)
|
||||
and not envs.VLLM_BATCH_INVARIANT
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _ihc_pre_stage1(
|
||||
x_ptr,
|
||||
weight_ptr,
|
||||
partial_ptr,
|
||||
K_TOTAL: tl.constexpr,
|
||||
HC_MULT: tl.constexpr,
|
||||
HC_POW2: tl.constexpr,
|
||||
NUM_SPLITS: tl.constexpr,
|
||||
BLOCK_K: tl.constexpr,
|
||||
PARTIAL_STRIDE: tl.constexpr,
|
||||
):
|
||||
token_idx = tl.program_id(0).to(tl.int64)
|
||||
split_idx = tl.program_id(1)
|
||||
hc_idx = tl.arange(0, HC_POW2)
|
||||
hc_mask = hc_idx < HC_MULT
|
||||
|
||||
k_offsets = split_idx * BLOCK_K + tl.arange(0, BLOCK_K)
|
||||
k_mask = k_offsets < K_TOTAL
|
||||
x = tl.load(
|
||||
x_ptr + token_idx * K_TOTAL + k_offsets,
|
||||
mask=k_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
sum_squares = tl.sum(x * x, axis=0)
|
||||
|
||||
weight_offsets = hc_idx[:, None] * K_TOTAL + k_offsets[None, :]
|
||||
weight_mask = hc_mask[:, None] & k_mask[None, :]
|
||||
pre_mix = tl.sum(
|
||||
tl.load(weight_ptr + weight_offsets, mask=weight_mask, other=0.0) * x[None, :],
|
||||
axis=1,
|
||||
)
|
||||
post_mix = tl.sum(
|
||||
tl.load(
|
||||
weight_ptr + HC_MULT * K_TOTAL + weight_offsets,
|
||||
mask=weight_mask,
|
||||
other=0.0,
|
||||
)
|
||||
* x[None, :],
|
||||
axis=1,
|
||||
)
|
||||
|
||||
partial = partial_ptr + (token_idx * NUM_SPLITS + split_idx) * PARTIAL_STRIDE
|
||||
tl.store(partial, sum_squares)
|
||||
tl.store(partial + 1 + hc_idx, pre_mix, mask=hc_mask)
|
||||
tl.store(partial + 1 + HC_POW2 + hc_idx, post_mix, mask=hc_mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _ihc_pre_stage2(
|
||||
x_ptr,
|
||||
partial_ptr,
|
||||
scale_ptr,
|
||||
base_ptr,
|
||||
output_ptr,
|
||||
post_ptr,
|
||||
HIDDEN_SIZE: tl.constexpr,
|
||||
K_TOTAL: tl.constexpr,
|
||||
HC_MULT: tl.constexpr,
|
||||
HC_POW2: tl.constexpr,
|
||||
NUM_SPLITS: tl.constexpr,
|
||||
PARTIAL_STRIDE: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
MAGNITUDE: tl.constexpr,
|
||||
NORM_EPS: tl.constexpr,
|
||||
HC_EPS: tl.constexpr,
|
||||
):
|
||||
token_idx = tl.program_id(0).to(tl.int64)
|
||||
hidden_block_idx = tl.program_id(1)
|
||||
hc_idx = tl.arange(0, HC_POW2)
|
||||
hc_mask = hc_idx < HC_MULT
|
||||
|
||||
partial_row = partial_ptr + token_idx * NUM_SPLITS * PARTIAL_STRIDE
|
||||
sum_squares = tl.zeros((), dtype=tl.float32)
|
||||
pre_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
|
||||
post_mix = tl.zeros((HC_POW2,), dtype=tl.float32)
|
||||
for split_idx in tl.static_range(NUM_SPLITS):
|
||||
partial = partial_row + split_idx * PARTIAL_STRIDE
|
||||
sum_squares += tl.load(partial)
|
||||
pre_mix += tl.load(partial + 1 + hc_idx, mask=hc_mask, other=0.0)
|
||||
post_mix += tl.load(
|
||||
partial + 1 + HC_POW2 + hc_idx,
|
||||
mask=hc_mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
reciprocal_rms = tl.rsqrt(sum_squares / K_TOTAL + NORM_EPS)
|
||||
pre_scale = tl.load(scale_ptr)
|
||||
post_scale = tl.load(scale_ptr + 1)
|
||||
pre_base = tl.load(base_ptr + hc_idx, mask=hc_mask, other=0.0)
|
||||
post_base = tl.load(base_ptr + HC_MULT + hc_idx, mask=hc_mask, other=0.0)
|
||||
pre = tl.sigmoid(pre_mix * reciprocal_rms * pre_scale + pre_base) + HC_EPS
|
||||
|
||||
if hidden_block_idx == 0:
|
||||
post = (
|
||||
MAGNITUDE * tl.sigmoid(post_mix * reciprocal_rms * post_scale + post_base)
|
||||
+ HC_EPS
|
||||
)
|
||||
tl.store(
|
||||
post_ptr + token_idx * HC_MULT + hc_idx,
|
||||
post,
|
||||
mask=hc_mask,
|
||||
)
|
||||
|
||||
hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
hidden_mask = hidden_offsets < HIDDEN_SIZE
|
||||
x_row = x_ptr + token_idx * K_TOTAL
|
||||
output = tl.zeros((BLOCK_D,), dtype=tl.float32)
|
||||
for channel_idx in tl.static_range(HC_MULT):
|
||||
channel = tl.load(
|
||||
x_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
|
||||
mask=hidden_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
channel_gate = tl.sum(tl.where(hc_idx == channel_idx, pre, 0.0), axis=0)
|
||||
output += channel_gate * channel
|
||||
|
||||
tl.store(
|
||||
output_ptr + token_idx * HIDDEN_SIZE + hidden_offsets,
|
||||
output.to(output_ptr.dtype.element_ty),
|
||||
mask=hidden_mask,
|
||||
)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _ihc_post_kernel(
|
||||
x_ptr,
|
||||
residual_ptr,
|
||||
post_ptr,
|
||||
output_ptr,
|
||||
HIDDEN_SIZE: tl.constexpr,
|
||||
HC_MULT: tl.constexpr,
|
||||
HC_POW2: tl.constexpr,
|
||||
BLOCK_D: tl.constexpr,
|
||||
):
|
||||
token_idx = tl.program_id(0).to(tl.int64)
|
||||
hidden_block_idx = tl.program_id(1)
|
||||
hc_idx = tl.arange(0, HC_POW2)
|
||||
hc_mask = hc_idx < HC_MULT
|
||||
post = tl.load(
|
||||
post_ptr + token_idx * HC_MULT + hc_idx,
|
||||
mask=hc_mask,
|
||||
other=0.0,
|
||||
)
|
||||
|
||||
hidden_offsets = hidden_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
|
||||
hidden_mask = hidden_offsets < HIDDEN_SIZE
|
||||
x = tl.load(
|
||||
x_ptr + token_idx * HIDDEN_SIZE + hidden_offsets,
|
||||
mask=hidden_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
residual_row = residual_ptr + token_idx * HC_MULT * HIDDEN_SIZE
|
||||
output_row = output_ptr + token_idx * HC_MULT * HIDDEN_SIZE
|
||||
for channel_idx in tl.static_range(HC_MULT):
|
||||
residual = tl.load(
|
||||
residual_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
|
||||
mask=hidden_mask,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
channel_gate = tl.sum(tl.where(hc_idx == channel_idx, post, 0.0), axis=0)
|
||||
output = channel_gate * x + residual
|
||||
tl.store(
|
||||
output_row + channel_idx * HIDDEN_SIZE + hidden_offsets,
|
||||
output.to(output_ptr.dtype.element_ty),
|
||||
mask=hidden_mask,
|
||||
)
|
||||
|
||||
|
||||
def triton_ihc_pre(
|
||||
x: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
scale: torch.Tensor,
|
||||
base: torch.Tensor,
|
||||
magnitude: float,
|
||||
hc_eps: float,
|
||||
norm_eps: float,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Reduce iHC channels and produce the post gates."""
|
||||
assert x.dim() == 3
|
||||
assert x.is_cuda and weight.is_cuda and scale.is_cuda and base.is_cuda
|
||||
assert weight.dtype == torch.float32
|
||||
assert scale.dtype == torch.float32 and base.dtype == torch.float32
|
||||
|
||||
x = x.contiguous()
|
||||
assert weight.is_contiguous()
|
||||
scale = scale.contiguous()
|
||||
base = base.contiguous()
|
||||
num_tokens, hc_mult, hidden_size = x.shape
|
||||
k_total = hc_mult * hidden_size
|
||||
assert weight.shape == (2 * hc_mult, k_total)
|
||||
assert scale.shape == (2,)
|
||||
assert base.shape == (2 * hc_mult,)
|
||||
|
||||
if num_tokens == 0:
|
||||
return (
|
||||
torch.empty((0, hidden_size), dtype=x.dtype, device=x.device),
|
||||
torch.empty((0, hc_mult), dtype=torch.float32, device=x.device),
|
||||
)
|
||||
|
||||
output = torch.empty((num_tokens, hidden_size), dtype=x.dtype, device=x.device)
|
||||
post = torch.empty((num_tokens, hc_mult), dtype=torch.float32, device=x.device)
|
||||
hc_pow2 = triton.next_power_of_2(hc_mult)
|
||||
num_splits = triton.cdiv(k_total, _BLOCK_K)
|
||||
partial_stride = 1 + 2 * hc_pow2
|
||||
partial = torch.empty(
|
||||
(num_tokens, num_splits, partial_stride),
|
||||
dtype=torch.float32,
|
||||
device=x.device,
|
||||
)
|
||||
|
||||
_ihc_pre_stage1[(num_tokens, num_splits)](
|
||||
x,
|
||||
weight,
|
||||
partial,
|
||||
K_TOTAL=k_total,
|
||||
HC_MULT=hc_mult,
|
||||
HC_POW2=hc_pow2,
|
||||
NUM_SPLITS=num_splits,
|
||||
BLOCK_K=_BLOCK_K,
|
||||
PARTIAL_STRIDE=partial_stride,
|
||||
num_warps=8,
|
||||
enable_fp_fusion=False,
|
||||
)
|
||||
_ihc_pre_stage2[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
|
||||
x,
|
||||
partial,
|
||||
scale,
|
||||
base,
|
||||
output,
|
||||
post,
|
||||
HIDDEN_SIZE=hidden_size,
|
||||
K_TOTAL=k_total,
|
||||
HC_MULT=hc_mult,
|
||||
HC_POW2=hc_pow2,
|
||||
NUM_SPLITS=num_splits,
|
||||
PARTIAL_STRIDE=partial_stride,
|
||||
BLOCK_D=_BLOCK_D,
|
||||
MAGNITUDE=magnitude,
|
||||
NORM_EPS=norm_eps,
|
||||
HC_EPS=hc_eps,
|
||||
num_warps=4,
|
||||
enable_fp_fusion=False,
|
||||
)
|
||||
return output, post
|
||||
|
||||
|
||||
def triton_ihc_post(
|
||||
x: torch.Tensor,
|
||||
residual: torch.Tensor,
|
||||
post: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Scatter a sub-block output back over the iHC residual channels."""
|
||||
assert x.dim() == 2
|
||||
assert x.is_cuda and residual.is_cuda and post.is_cuda
|
||||
assert post.dtype == torch.float32
|
||||
|
||||
x = x.contiguous()
|
||||
residual = residual.contiguous()
|
||||
post = post.contiguous()
|
||||
num_tokens, hidden_size = x.shape
|
||||
hc_mult = post.shape[-1]
|
||||
assert residual.shape == (num_tokens, hc_mult, hidden_size)
|
||||
assert post.shape == (num_tokens, hc_mult)
|
||||
|
||||
if num_tokens == 0:
|
||||
return torch.empty((0, hc_mult, hidden_size), dtype=x.dtype, device=x.device)
|
||||
|
||||
output = torch.empty(
|
||||
(num_tokens, hc_mult, hidden_size), dtype=x.dtype, device=x.device
|
||||
)
|
||||
_ihc_post_kernel[(num_tokens, triton.cdiv(hidden_size, _BLOCK_D))](
|
||||
x,
|
||||
residual,
|
||||
post,
|
||||
output,
|
||||
HIDDEN_SIZE=hidden_size,
|
||||
HC_MULT=hc_mult,
|
||||
HC_POW2=triton.next_power_of_2(hc_mult),
|
||||
BLOCK_D=_BLOCK_D,
|
||||
num_warps=4,
|
||||
enable_fp_fusion=False,
|
||||
)
|
||||
return output
|
||||
Reference in New Issue
Block a user