dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

feat: import CUDA kernels from xllm/CCCL/FLA upstream repos

Browse Source 
Sources cloned and tree'd (no --depth):
  - jd-opensource/xllm: ILU kernels, CUDA kernels, MoE kernels
  - NVIDIA/cccl: CUB tuning/dispatch headers (block-level primitives)
  - fla-org/flash-linear-attention: Triton GDN kernels
  - NVIDIA/cutlass: grouped GEMM reference (read, not copied)
  - Dao-AILab/flash-attention: attention kernel reference (SM80+, read only)

New CUDA kernels (from xllm, SM-agnostic, portable to BI-V100):
  ex_engine/xllm_kernels/cuda/activation.cu    (188 lines) — silu_and_mul, gelu
  ex_engine/xllm_kernels/cuda/norm.cu          (600 lines) — rms_norm, fused_add_rms_norm
  ex_engine/xllm_kernels/cuda/rope.cu          (258 lines) — rotary_embedding
  ex_engine/xllm_kernels/cuda/block_copy.cu    (209 lines) — copy_blocks, swap_blocks
  ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu (101 lines) — KV cache ops
  ex_engine/xllm_kernels/cuda/headers/         (5 headers for compilation)

ILU bridge kernel sources (from xllm, verified SAME as upstream):
  ex_engine/xllm_kernels/ilu/    (10 files, 925 lines total)
  — activation.cpp, attention.cpp, fused_moe.cpp, group_gemm.cpp,
    matmul.cpp, norm.cpp, rope.cpp, ilu_ops_api.h, ixformer.h, utils.h

FLA Triton GDN kernels (for GatedDeltaNet without SM90+ FlashQLA):
  ex_engine/fla_kernels/gated_delta_rule/  (7 files, 2370 lines)
  — chunk_fwd.py (428), chunk.py (487), wy_fast.py (409),
    fused_recurrent.py (392), naive.py (161), gate.py (380)

CCCL sync (12 tuning + 14 dispatch headers updated from NVIDIA/cccl):
  cccl_upstream/cub/cub/device/dispatch/tuning/ — 12 changed files synced
  cccl_upstream/cub/cub/device/dispatch/ — 14 changed dispatch files synced

Compilation targets for real machine (ivcore10):
  1. CUDA kernels: --cuda-gpu-arch=ivcore10 via corex clang/16
  2. ILU bridges: torch.utils.cpp_extension linking ixformer .so
  3. FLA kernels: Triton JIT (if Triton works on BI-V100)
This commit is contained in:
claude
2026-08-14 07:48:52 +00:00
parent 051b02d3cd
commit 8d75652949
66 changed files with 12834 additions and 5458 deletions
Download Patch File Download Diff File Expand all files Collapse all files

17
ex_engine/fla_kernels/gated_delta_rule/__init__.py Normal file
View File

@@ -0,0 +1,17 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from .chunk import chunk_gated_delta_rule, chunk_gdn
from .fused_recurrent import fused_recurrent_gated_delta_rule, fused_recurrent_gdn
from .naive import naive_chunk_gated_delta_rule, naive_recurrent_gated_delta_rule
__all__ = [
"chunk_gated_delta_rule", "chunk_gdn",
"fused_recurrent_gated_delta_rule", "fused_recurrent_gdn",
"naive_chunk_gated_delta_rule",
"naive_recurrent_gated_delta_rule",
]

591
ex_engine/fla_kernels/gated_delta_rule/chunk.py Normal file
View File

@@ -0,0 +1,591 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import warnings
import torch
from fla.modules.l2norm import l2norm_bwd, l2norm_fwd
from fla.ops.backends import dispatch
from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_bwd_dhu, chunk_gated_delta_rule_fwd_h
from fla.ops.common.chunk_o import chunk_bwd_dqkwg, chunk_bwd_dv_local, chunk_fwd_o
from fla.ops.common.gate import fused_beta_sigmoid, fused_beta_sigmoid_bwd
from fla.ops.cp import FLACPContext
from fla.ops.cp.chunk_delta_h import (
chunk_gated_delta_rule_bwd_dhu_pre_process,
chunk_gated_delta_rule_fwd_h_pre_process,
compress_h0,
expand_h0,
)
from fla.ops.gated_delta_rule.chunk_fwd import chunk_gated_delta_rule_fwd_intra
from fla.ops.gated_delta_rule.gate import gdn_gate_bwd, gdn_gate_chunk_cumsum
from fla.ops.gated_delta_rule.wy_fast import prepare_wy_repr_bwd, recompute_w_u_fwd
from fla.ops.utils import chunk_local_cumsum
from fla.ops.utils.constant import RCP_LN2
from fla.ops.utils.index import prepare_chunk_indices
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
def chunk_gated_delta_rule_fwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
output_final_state: bool,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
chunk_size: int = 64,
):
g_input = g if use_gate_in_kernel else None
if use_gate_in_kernel:
g = gdn_gate_chunk_cumsum(
g=g,
A_log=A_log,
chunk_size=chunk_size,
scale=RCP_LN2,
dt_bias=dt_bias,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
else:
g = chunk_local_cumsum(
g,
chunk_size=chunk_size,
scale=RCP_LN2,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
# obtain WY representation. u is actually the new v.
# fused kkt + solve_tril + recompute_w_u
w, u, A = chunk_gated_delta_rule_fwd_intra(
k=k,
v=v,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=chunk_size,
)
if cp_context is not None:
initial_state = chunk_gated_delta_rule_fwd_h_pre_process(
k=k,
w=w,
u=u,
g=g,
cu_seqlens=cu_seqlens,
initial_state=initial_state,
context=cp_context,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
h, v_new, final_state = chunk_gated_delta_rule_fwd_h(
k=k,
w=w,
u=u,
g=g,
initial_state=initial_state,
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
if cp_context is not None:
initial_state = compress_h0(initial_state, context=cp_context)
o = chunk_fwd_o(
q=q,
k=k,
v=v_new,
h=h,
g=g,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
return g, o, A, final_state, initial_state, g_input
def chunk_gated_delta_rule_bwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
do: torch.Tensor,
dht: torch.Tensor,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
chunk_indices: torch.LongTensor | None = None,
use_gate_in_kernel: bool = False,
g_input: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
chunk_size: int = 64,
):
w, u = recompute_w_u_fwd(
k=k,
v=v,
beta=beta,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
if cp_context is not None:
initial_state = expand_h0(initial_state, context=cp_context)
h, v_new, _ = chunk_gated_delta_rule_fwd_h(
k=k,
w=w,
u=u,
g=g,
initial_state=initial_state,
output_final_state=False,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dv = chunk_bwd_dv_local(
q=q,
k=k,
g=g,
do=do,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=chunk_size,
)
if cp_context is not None:
# initial_state is None in the CP mode
# We only need to compute dht of current rank and pass it to the backward kernel
dht, initial_state = chunk_gated_delta_rule_bwd_dhu_pre_process(
q=q,
k=k,
w=w,
do=do,
dv=dv,
g=g,
scale=scale,
cu_seqlens=cu_seqlens,
dht=dht,
initial_state=initial_state,
context=cp_context,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dh, dh0, dv = chunk_gated_delta_rule_bwd_dhu(
q=q,
k=k,
w=w,
g=g,
h0=initial_state,
dht=dht,
do=do,
dv=dv,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dq, dk, dw, dg = chunk_bwd_dqkwg(
q=q,
k=k,
v=v_new,
w=w,
g=g,
h=h,
dv=dv,
do=do,
dh=dh,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
chunk_size=chunk_size,
)
dk2, dv, db, dg2 = prepare_wy_repr_bwd(
k=k,
v=v,
beta=beta,
g=g,
A=A,
dw=dw,
du=dv,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
dk.add_(dk2)
dg.add_(dg2)
dg = chunk_local_cumsum(dg, chunk_size=chunk_size, reverse=True, cu_seqlens=cu_seqlens, chunk_indices=chunk_indices)
dA_log, ddt_bias = None, None
if use_gate_in_kernel:
dg, dA_log, ddt_bias = gdn_gate_bwd(g=g_input, A_log=A_log, dt_bias=dt_bias, dyg=dg)
return dq, dk, dv, db, dg, dh0, dA_log, ddt_bias
class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
def forward(
ctx,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float,
initial_state: torch.Tensor,
output_final_state: bool,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
use_qk_l2norm_in_kernel: bool = False,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
cp_context: FLACPContext | None = None,
chunk_size: int = 64,
):
q_rstd, k_rstd = None, None
if use_qk_l2norm_in_kernel:
q, q_rstd = l2norm_fwd(q)
k, k_rstd = l2norm_fwd(k)
beta_raw = beta
if use_beta_sigmoid_in_kernel:
beta = fused_beta_sigmoid(beta_raw, scale=2.0 if allow_neg_eigval else 1.0)
chunk_indices = None
if cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=cu_seqlens_cpu)
g, o, A, final_state, initial_state, g_input = chunk_gated_delta_rule_fwd(
q=q,
k=k,
v=v,
g=g,
beta=beta,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
cu_seqlens=cu_seqlens,
cp_context=cp_context,
chunk_indices=chunk_indices,
state_v_first=state_v_first,
use_gate_in_kernel=use_gate_in_kernel,
A_log=A_log,
dt_bias=dt_bias,
chunk_size=chunk_size,
)
ctx.save_for_backward(
q,
q_rstd,
k,
k_rstd,
v,
g,
beta_raw,
beta,
A,
initial_state,
cu_seqlens,
chunk_indices,
g_input,
A_log,
dt_bias,
)
ctx.scale = scale
ctx.chunk_size = chunk_size
ctx.use_qk_l2norm_in_kernel = use_qk_l2norm_in_kernel
ctx.use_beta_sigmoid_in_kernel = use_beta_sigmoid_in_kernel
ctx.allow_neg_eigval = allow_neg_eigval
ctx.cp_context = cp_context
ctx.state_v_first = state_v_first
ctx.use_gate_in_kernel = use_gate_in_kernel
return o.to(q.dtype), final_state
@staticmethod
@input_guard
@autocast_custom_bwd
def backward(
ctx,
do: torch.Tensor,
dht: torch.Tensor,
):
(
q,
q_rstd,
k,
k_rstd,
v,
g,
beta_raw,
beta,
A,
initial_state,
cu_seqlens,
chunk_indices,
g_input,
A_log,
dt_bias,
) = ctx.saved_tensors
dq, dk, dv, db, dg, dh0, dA_log, ddt_bias = chunk_gated_delta_rule_bwd(
q=q,
k=k,
v=v,
g=g,
beta=beta,
A=A,
scale=ctx.scale,
initial_state=initial_state,
do=do,
dht=dht,
cu_seqlens=cu_seqlens,
cp_context=ctx.cp_context,
chunk_indices=chunk_indices,
state_v_first=ctx.state_v_first,
use_gate_in_kernel=ctx.use_gate_in_kernel,
g_input=g_input,
A_log=A_log,
dt_bias=dt_bias,
chunk_size=ctx.chunk_size,
)
if ctx.use_qk_l2norm_in_kernel:
dq = l2norm_bwd(q, q_rstd, dq)
dk = l2norm_bwd(k, k_rstd, dk)
if ctx.use_beta_sigmoid_in_kernel:
db = fused_beta_sigmoid_bwd(beta_raw, db, scale=2.0 if ctx.allow_neg_eigval else 1.0)
return (
dq.to(q), dk.to(k), dv.to(v), dg.to(g), db.to(beta_raw),
None, dh0, None, None, None, None, None, None, dA_log, ddt_bias,
None, None, None, None,
)
@dispatch('gated_delta_rule')
@torch.compiler.disable
def chunk_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
scale: float | None = None,
initial_state: torch.Tensor | None = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
cu_seqlens_cpu: torch.LongTensor | None = None,
cp_context: FLACPContext | None = None,
**kwargs,
):
r"""
Args:
q (torch.Tensor):
queries of shape `[B, T, H, K]`.
k (torch.Tensor):
keys of shape `[B, T, H, K]`.
v (torch.Tensor):
values of shape `[B, T, HV, V]`.
GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
g (torch.Tensor):
(forget) gating tensor of shape `[B, T, HV]`.
When `use_gate_in_kernel=False` (default), `g` should be in log space (pre-computed decay).
When `use_gate_in_kernel=True`, `g` is the raw input before gate activation;
the kernel fuses `-exp(A_log) * softplus(g + dt_bias)` + chunk cumsum internally.
beta (torch.Tensor):
betas of shape `[B, T, HV]`.
scale (Optional[float]):
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, HV, K, V]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
use_qk_l2norm_in_kernel (bool):
Whether to apply L2norm to the q/k tensor internally. Default: `False`.
use_gate_in_kernel (bool):
Whether to compute the log-space GDN decay internally.
When `True`, the passed `g` is the raw input, and `A_log` must be provided.
The kernel fuses gate activation + chunk cumsum in a single pass.
Default: `False`.
A_log (Optional[torch.Tensor]):
Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
dt_bias (Optional[torch.Tensor]):
Bias added to `g` before activation, of shape `[HV]`.
Only used when `use_gate_in_kernel=True`.
use_beta_sigmoid_in_kernel (bool):
Whether to apply `torch.sigmoid(beta)` before launching the chunk kernel.
- If `True`, the passed `beta` acts as the raw beta logits.
- If `False`, `beta` is expected to already be in post-sigmoid space.
Default: `False`.
allow_neg_eigval (bool):
Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
state_v_first (Optional[bool]):
Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
cp_context (Optional[FLACPContext]):
Context parallel context for distributed training across multiple devices.
When provided, `initial_state` and `output_final_state` are not supported,
and `cu_seqlens` will be overridden by the context. Default: `None`.
Returns:
o (torch.Tensor):
Outputs of shape `[B, T, HV, V]`.
final_state (torch.Tensor):
Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
Examples::
>>> import torch
>>> import torch.nn.functional as F
>>> from einops import rearrange
>>> from fla.ops.gated_delta_rule import chunk_gated_delta_rule
# inputs with equal lengths
>>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
>>> q = torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda')
>>> k = F.normalize(torch.randn(B, T, H, K, dtype=torch.bfloat16, device='cuda'), p=2, dim=-1)
>>> v = torch.randn(B, T, HV, V, dtype=torch.bfloat16, device='cuda')
>>> beta = torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda').sigmoid()
>>> g = F.logsigmoid(torch.rand(B, T, HV, dtype=torch.bfloat16, device='cuda'))
>>> h0 = torch.randn(B, HV, K, V, dtype=torch.bfloat16, device='cuda')
>>> o, ht = chunk_gated_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True
)
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
>>> q, k, v, beta, g = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, beta, g))
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
>>> o, ht = chunk_gated_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True,
cu_seqlens=cu_seqlens
)
"""
if 'transpose_state_layout' in kwargs:
if state_v_first:
raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
warnings.warn(
"`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
DeprecationWarning,
stacklevel=2,
)
state_v_first = kwargs.pop('transpose_state_layout')
# Validate head dimensions
if q.shape[2] != k.shape[2]:
raise ValueError(
f"q and k must have the same number of heads, "
f"but got q.shape[2]={q.shape[2]} and k.shape[2]={k.shape[2]}"
)
H, HV = q.shape[2], v.shape[2]
if HV % H != 0:
raise ValueError(
f"For GVA, num_v_heads (HV={HV}) must be evenly divisible by "
f"num_heads (H={H}), but got HV % H = {HV % H}"
)
if 'head_first' in kwargs:
raise DeprecationWarning(
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
)
chunk_size = kwargs.pop('chunk_size', 64)
if chunk_size not in (16, 32, 64):
raise ValueError(f"`chunk_size` must be 16, 32, or 64 for Gated Delta Rule, got {chunk_size}.")
if cp_context is not None:
assert initial_state is None, "Initial state is not supported for CP"
assert output_final_state is False, "Output final state is not supported for CP"
assert cp_context.cu_seqlens is not None, "cu_seqlens is required for CP"
cu_seqlens = cp_context.cu_seqlens
if cp_context.cu_seqlens_cpu is not None:
cu_seqlens_cpu = cp_context.cu_seqlens_cpu
if cu_seqlens is not None:
if q.shape[0] != 1:
raise ValueError(
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
f"Please flatten variable-length inputs before processing.",
)
if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
raise ValueError(
f"The number of initial states is expected to be equal to the number of input sequences, "
f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
)
use_gate_in_kernel = kwargs.get('use_gate_in_kernel', False)
A_log = kwargs.get('A_log')
dt_bias = kwargs.get('dt_bias')
if use_gate_in_kernel:
assert A_log is not None, "A_log must be provided when use_gate_in_kernel=True."
if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
if scale is None:
scale = k.shape[-1] ** -0.5
o, final_state = ChunkGatedDeltaRuleFunction.apply(
q,
k,
v,
g,
beta,
scale,
initial_state,
output_final_state,
state_v_first,
cu_seqlens,
cu_seqlens_cpu,
use_qk_l2norm_in_kernel,
use_gate_in_kernel,
A_log,
dt_bias,
use_beta_sigmoid_in_kernel,
allow_neg_eigval,
cp_context,
chunk_size,
)
return o, final_state
chunk_gdn = chunk_gated_delta_rule

428
ex_engine/fla_kernels/gated_delta_rule/chunk_fwd.py Normal file
View File

@@ -0,0 +1,428 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.common.chunk_scaled_dot_kkt import chunk_scaled_dot_kkt_fwd
from fla.ops.gated_delta_rule.wy_fast import recompute_w_u_fwd
from fla.ops.utils import prepare_chunk_indices, solve_tril
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.op import exp2
from fla.utils import IS_INTEL, IS_TF32_SUPPORTED, autotune_cache_kwargs
if IS_TF32_SUPPORTED:
SOLVE_TRIL_DOT_PRECISION = tl.constexpr('tf32')
else:
SOLVE_TRIL_DOT_PRECISION = tl.constexpr('ieee')
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({'BK': BK}, num_warps=num_warps)
for BK in [32, 64]
for num_warps in [1, 2, 4]
],
key=['H', 'HV', 'K', 'BC'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
k,
g,
beta,
A,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
BT: tl.constexpr,
BC: tl.constexpr,
BK: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
"""
Fused kernel: compute beta * K @ K^T (lower triangular) + solve_tril (I+A)^{-1} in one pass.
This kernel fuses chunk_scaled_dot_kkt_fwd and solve_tril into a single kernel,
avoiding the HBM round-trip for the intermediate A matrix.
Steps:
1. Compute all 10 lower-triangular [BC, BC] blocks of beta * K @ K^T in registers
2. Apply gate and beta scaling
3. Forward substitution on diagonal blocks
4. Block merge to get full (I+A)^{-1}
5. Write result to A (output)
"""
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
if i_t * BT >= T:
return
i_tc0 = i_t * BT
i_tc1 = i_t * BT + BC
i_tc2 = i_t * BT + 2 * BC
i_tc3 = i_t * BT + 3 * BC
k += (bos * H + i_h // (HV // H)) * K
A += (bos * HV + i_h) * BT
o_i = tl.arange(0, BC)
m_tc0 = (i_tc0 + o_i) < T
m_tc1 = (i_tc1 + o_i) < T
m_tc2 = (i_tc2 + o_i) < T
m_tc3 = (i_tc3 + o_i) < T
# load beta for each sub-chunk
p_b0 = beta + bos * HV + i_h + (i_tc0 + o_i) * HV
p_b1 = beta + bos * HV + i_h + (i_tc1 + o_i) * HV
p_b2 = beta + bos * HV + i_h + (i_tc2 + o_i) * HV
p_b3 = beta + bos * HV + i_h + (i_tc3 + o_i) * HV
b_b0 = tl.load(p_b0, mask=m_tc0, other=0.0).to(tl.float32)
b_b1 = tl.load(p_b1, mask=m_tc1, other=0.0).to(tl.float32)
b_b2 = tl.load(p_b2, mask=m_tc2, other=0.0).to(tl.float32)
b_b3 = tl.load(p_b3, mask=m_tc3, other=0.0).to(tl.float32)
# load gate if used
if USE_G:
p_g0 = g + bos * HV + i_h + (i_tc0 + o_i) * HV
p_g1 = g + bos * HV + i_h + (i_tc1 + o_i) * HV
p_g2 = g + bos * HV + i_h + (i_tc2 + o_i) * HV
p_g3 = g + bos * HV + i_h + (i_tc3 + o_i) * HV
b_g0 = tl.load(p_g0, mask=m_tc0, other=0.0).to(tl.float32)
b_g1 = tl.load(p_g1, mask=m_tc1, other=0.0).to(tl.float32)
b_g2 = tl.load(p_g2, mask=m_tc2, other=0.0).to(tl.float32)
b_g3 = tl.load(p_g3, mask=m_tc3, other=0.0).to(tl.float32)
############################################################################
# Step 1: compute all 10 lower-triangular [BC, BC] blocks of K @ K^T
############################################################################
# 4 diagonal blocks
b_A00 = tl.zeros([BC, BC], dtype=tl.float32)
b_A11 = tl.zeros([BC, BC], dtype=tl.float32)
b_A22 = tl.zeros([BC, BC], dtype=tl.float32)
b_A33 = tl.zeros([BC, BC], dtype=tl.float32)
# 6 off-diagonal blocks
b_A10 = tl.zeros([BC, BC], dtype=tl.float32)
b_A20 = tl.zeros([BC, BC], dtype=tl.float32)
b_A21 = tl.zeros([BC, BC], dtype=tl.float32)
b_A30 = tl.zeros([BC, BC], dtype=tl.float32)
b_A31 = tl.zeros([BC, BC], dtype=tl.float32)
b_A32 = tl.zeros([BC, BC], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
p_k0 = k + (i_tc0 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k0 = tl.load(p_k0, mask=m_tc0[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 0
b_A00 += tl.dot(b_k0, tl.trans(b_k0))
if i_tc1 < T:
p_k1 = k + (i_tc1 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k1 = tl.load(p_k1, mask=m_tc1[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 1
b_A11 += tl.dot(b_k1, tl.trans(b_k1))
# off-diagonal (1,0)
b_A10 += tl.dot(b_k1, tl.trans(b_k0))
if i_tc2 < T:
p_k2 = k + (i_tc2 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k2 = tl.load(p_k2, mask=m_tc2[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 2
b_A22 += tl.dot(b_k2, tl.trans(b_k2))
# off-diagonal (2,0), (2,1)
b_A20 += tl.dot(b_k2, tl.trans(b_k0))
b_A21 += tl.dot(b_k2, tl.trans(b_k1))
if i_tc3 < T:
p_k3 = k + (i_tc3 + o_i)[:, None] * (H*K) + o_k[None, :]
b_k3 = tl.load(p_k3, mask=m_tc3[:, None] & (o_k[None, :] < K), other=0.0)
# diagonal block 3
b_A33 += tl.dot(b_k3, tl.trans(b_k3))
# off-diagonal (3,0), (3,1), (3,2)
b_A30 += tl.dot(b_k3, tl.trans(b_k0))
b_A31 += tl.dot(b_k3, tl.trans(b_k1))
b_A32 += tl.dot(b_k3, tl.trans(b_k2))
############################################################################
# Step 2: apply gate and beta scaling
############################################################################
# apply gate, beta scaling, and masking
# m_d: strictly lower triangular mask for diagonal blocks
# m_tc: boundary mask to prevent NaN from 0 * inf (IEEE 754) when
# out-of-bounds g loads as 0 via boundary_check and exp2(0 - g_inbounds) overflows
m_d = o_i[:, None] > o_i[None, :]
m_I = o_i[:, None] == o_i[None, :]
if USE_G:
b_A00 *= tl.where(m_d & m_tc0[:, None] & m_tc0[None, :], exp2(b_g0[:, None] - b_g0[None, :]), 0.)
b_A11 *= tl.where(m_d & m_tc1[:, None] & m_tc1[None, :], exp2(b_g1[:, None] - b_g1[None, :]), 0.)
b_A22 *= tl.where(m_d & m_tc2[:, None] & m_tc2[None, :], exp2(b_g2[:, None] - b_g2[None, :]), 0.)
b_A33 *= tl.where(m_d & m_tc3[:, None] & m_tc3[None, :], exp2(b_g3[:, None] - b_g3[None, :]), 0.)
b_A10 *= tl.where(m_tc1[:, None] & m_tc0[None, :], exp2(b_g1[:, None] - b_g0[None, :]), 0.)
b_A20 *= tl.where(m_tc2[:, None] & m_tc0[None, :], exp2(b_g2[:, None] - b_g0[None, :]), 0.)
b_A21 *= tl.where(m_tc2[:, None] & m_tc1[None, :], exp2(b_g2[:, None] - b_g1[None, :]), 0.)
b_A30 *= tl.where(m_tc3[:, None] & m_tc0[None, :], exp2(b_g3[:, None] - b_g0[None, :]), 0.)
b_A31 *= tl.where(m_tc3[:, None] & m_tc1[None, :], exp2(b_g3[:, None] - b_g1[None, :]), 0.)
b_A32 *= tl.where(m_tc3[:, None] & m_tc2[None, :], exp2(b_g3[:, None] - b_g2[None, :]), 0.)
else:
b_A00 = tl.where(m_d, b_A00, 0.)
b_A11 = tl.where(m_d, b_A11, 0.)
b_A22 = tl.where(m_d, b_A22, 0.)
b_A33 = tl.where(m_d, b_A33, 0.)
# diagonal blocks: scaled by beta
b_A00 = b_A00 * b_b0[:, None]
b_A11 = b_A11 * b_b1[:, None]
b_A22 = b_A22 * b_b2[:, None]
b_A33 = b_A33 * b_b3[:, None]
# off-diagonal blocks: full block, scaled by beta
b_A10 = b_A10 * b_b1[:, None]
b_A20 = b_A20 * b_b2[:, None]
b_A21 = b_A21 * b_b2[:, None]
b_A30 = b_A30 * b_b3[:, None]
b_A31 = b_A31 * b_b3[:, None]
b_A32 = b_A32 * b_b3[:, None]
############################################################################
# Step 3: forward substitution on diagonal blocks -> (I + A_diag)^{-1}
#
# Same algorithm as solve_tril, but rows are extracted from in-register
# [BC, BC] tensor via tl.sum(tl.where(mask, tensor, 0), 0) instead of
# tl.load from HBM.
############################################################################
b_Ai00 = -b_A00
b_Ai11 = -b_A11
b_Ai22 = -b_A22
b_Ai33 = -b_A33
for i in range(2, min(BC, T - i_tc0)):
b_a00 = tl.sum(tl.where((o_i == i)[:, None], -b_A00, 0.), 0)
b_a00 = tl.where(o_i < i, b_a00, 0.)
b_a00 = b_a00 + tl.sum(b_a00[:, None] * b_Ai00, 0)
b_Ai00 = tl.where((o_i == i)[:, None], b_a00, b_Ai00)
for i in range(2, min(BC, T - i_tc1)):
b_a11 = tl.sum(tl.where((o_i == i)[:, None], -b_A11, 0.), 0)
b_a11 = tl.where(o_i < i, b_a11, 0.)
b_a11 = b_a11 + tl.sum(b_a11[:, None] * b_Ai11, 0)
b_Ai11 = tl.where((o_i == i)[:, None], b_a11, b_Ai11)
for i in range(2, min(BC, T - i_tc2)):
b_a22 = tl.sum(tl.where((o_i == i)[:, None], -b_A22, 0.), 0)
b_a22 = tl.where(o_i < i, b_a22, 0.)
b_a22 = b_a22 + tl.sum(b_a22[:, None] * b_Ai22, 0)
b_Ai22 = tl.where((o_i == i)[:, None], b_a22, b_Ai22)
for i in range(2, min(BC, T - i_tc3)):
b_a33 = tl.sum(tl.where((o_i == i)[:, None], -b_A33, 0.), 0)
b_a33 = tl.where(o_i < i, b_a33, 0.)
b_a33 = b_a33 + tl.sum(b_a33[:, None] * b_Ai33, 0)
b_Ai33 = tl.where((o_i == i)[:, None], b_a33, b_Ai33)
b_Ai00 += m_I
b_Ai11 += m_I
b_Ai22 += m_I
b_Ai33 += m_I
############################################################################
# Step 4: block merge -> full (I + A)^{-1}
############################################################################
b_Ai10 = -tl.dot(
tl.dot(b_Ai11, b_A10, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai00,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai21 = -tl.dot(
tl.dot(b_Ai22, b_A21, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai11,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai32 = -tl.dot(
tl.dot(b_Ai33, b_A32, input_precision=SOLVE_TRIL_DOT_PRECISION),
b_Ai22,
input_precision=SOLVE_TRIL_DOT_PRECISION
)
b_Ai20 = -tl.dot(
b_Ai22,
tl.dot(b_A20, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A21, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
b_Ai31 = -tl.dot(
b_Ai33,
tl.dot(b_A31, b_Ai11, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A32, b_Ai21, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
b_Ai30 = -tl.dot(
b_Ai33,
tl.dot(b_A30, b_Ai00, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A31, b_Ai10, input_precision=SOLVE_TRIL_DOT_PRECISION) +
tl.dot(b_A32, b_Ai20, input_precision=SOLVE_TRIL_DOT_PRECISION),
input_precision=SOLVE_TRIL_DOT_PRECISION,
)
############################################################################
# Step 5: store full (I + A)^{-1} to output A
############################################################################
p_A00 = A + (i_tc0 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A10 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A11 = A + (i_tc1 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A20 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A21 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A22 = A + (i_tc2 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
p_A30 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + o_i[None, :]
p_A31 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (BC + o_i)[None, :]
p_A32 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (2*BC + o_i)[None, :]
p_A33 = A + (i_tc3 + o_i)[:, None] * (HV*BT) + (3*BC + o_i)[None, :]
m_A0 = m_tc0[:, None] & (o_i[None, :] < BT)
m_A1 = m_tc1[:, None] & (o_i[None, :] < BT)
m_A2 = m_tc2[:, None] & (o_i[None, :] < BT)
m_A3 = m_tc3[:, None] & (o_i[None, :] < BT)
m_A11 = m_tc1[:, None] & ((BC + o_i)[None, :] < BT)
m_A21 = m_tc2[:, None] & ((BC + o_i)[None, :] < BT)
m_A22 = m_tc2[:, None] & ((2*BC + o_i)[None, :] < BT)
m_A31 = m_tc3[:, None] & ((BC + o_i)[None, :] < BT)
m_A32 = m_tc3[:, None] & ((2*BC + o_i)[None, :] < BT)
m_A33 = m_tc3[:, None] & ((3*BC + o_i)[None, :] < BT)
tl.store(p_A00, b_Ai00.to(A.dtype.element_ty), mask=m_A0)
tl.store(p_A10, b_Ai10.to(A.dtype.element_ty), mask=m_A1)
tl.store(p_A11, b_Ai11.to(A.dtype.element_ty), mask=m_A11)
tl.store(p_A20, b_Ai20.to(A.dtype.element_ty), mask=m_A2)
tl.store(p_A21, b_Ai21.to(A.dtype.element_ty), mask=m_A21)
tl.store(p_A22, b_Ai22.to(A.dtype.element_ty), mask=m_A22)
tl.store(p_A30, b_Ai30.to(A.dtype.element_ty), mask=m_A3)
tl.store(p_A31, b_Ai31.to(A.dtype.element_ty), mask=m_A31)
tl.store(p_A32, b_Ai32.to(A.dtype.element_ty), mask=m_A32)
tl.store(p_A33, b_Ai33.to(A.dtype.element_ty), mask=m_A33)
@dispatch('gated_delta_rule')
def chunk_gated_delta_rule_fwd_intra(
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_size: int = 64,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
r"""
GDN intra-chunk forward: fused or unfused kkt + solve_tril + recompute_w_u.
For ``chunk_size == 64``, this uses the fused kkt + solve_tril path. For
other supported chunk sizes, it computes the mathematically equivalent
representation with ``chunk_scaled_dot_kkt_fwd`` followed by ``solve_tril``.
Args:
k (torch.Tensor):
The key tensor of shape `[B, T, H, K]`.
v (torch.Tensor):
The value tensor of shape `[B, T, HV, V]`.
g (torch.Tensor):
The cumulative sum of the gate tensor of shape `[B, T, HV]`. Default: `None`.
beta (torch.Tensor):
The beta tensor of shape `[B, T, HV]`.
cu_seqlens (torch.LongTensor):
The cumulative sequence lengths. Default: `None`.
chunk_size (int):
The chunk size. Default: 64.
chunk_indices (torch.LongTensor):
Precomputed chunk indices. Default: `None`.
Returns:
w (torch.Tensor): shape `[B, T, HV, K]`
u (torch.Tensor): shape `[B, T, HV, V]`
A (torch.Tensor): shape `[B, T, HV, BT]`, the solved (I+A)^{-1} matrix
"""
if chunk_size not in (16, 32, 64):
raise ValueError(f"`chunk_size` must be 16, 32, or 64, got {chunk_size}.")
B, T, H, K, HV = *k.shape, beta.shape[2]
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
# The fused kernel keeps ten [BC, BC] fp32 accumulators live across the K loop.
# That fits NVIDIA's register file but spills on Intel GPUs, where the unfused
# two-kernel path measures 2.3-3.0x faster despite the extra HBM round-trip.
if BT == 64 and not IS_INTEL:
# Step 1: fused kkt + solve_tril
BC = 16
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
A = torch.zeros(B, T, HV, BT, device=k.device, dtype=k.dtype)
chunk_gated_delta_rule_fwd_kkt_solve_kernel[(NT, B * HV)](
k=k,
g=g,
beta=beta,
A=A,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
BT=BT,
BC=BC,
)
else:
# Step 1: mathematically equivalent unfused kkt + solve_tril
A = chunk_scaled_dot_kkt_fwd(
k=k,
g=g,
beta=beta,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
chunk_size=BT,
output_dtype=torch.float32,
)
A = solve_tril(
A=A,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
output_dtype=k.dtype,
)
# Step 2: recompute_w_u
w, u = recompute_w_u_fwd(
k=k,
v=v,
beta=beta,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
)
return w, u, A

478
ex_engine/fla_kernels/gated_delta_rule/fused_recurrent.py Normal file
View File

@@ -0,0 +1,478 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import warnings
import torch
import triton
import triton.language as tl
from fla.ops.utils.op import exp
from fla.ops.utils.softplus import softplus
from fla.utils import input_guard
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'USE_GK': lambda args: args['gk'] is not None,
'USE_GV': lambda args: args['gv'] is not None,
'USE_INITIAL_STATE': lambda args: args['h0'] is not None,
'STORE_FINAL_STATE': lambda args: args['ht'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
'USE_GATE_IN_KERNEL': lambda args: args['A_log'] is not None,
'HAS_DT_BIAS': lambda args: args['dt_bias'] is not None,
})
@triton.jit(do_not_specialize=['T'])
def fused_recurrent_gated_delta_rule_fwd_kernel(
q,
k,
v,
g,
gk,
gv,
beta,
A_log,
dt_bias,
o,
h0,
ht,
cu_seqlens,
scale,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
USE_GK: tl.constexpr,
USE_GV: tl.constexpr,
USE_QK_L2NORM_IN_KERNEL: tl.constexpr,
IS_BETA_HEADWISE: tl.constexpr,
USE_INITIAL_STATE: tl.constexpr,
STORE_FINAL_STATE: tl.constexpr,
STATE_V_FIRST: tl.constexpr,
IS_VARLEN: tl.constexpr,
USE_GATE_IN_KERNEL: tl.constexpr,
HAS_DT_BIAS: tl.constexpr,
APPLY_BETA_SIGMOID: tl.constexpr,
ALLOW_NEG_EIGVAL: tl.constexpr,
):
pid = tl.program_id(0)
NV = tl.cdiv(V, BV)
i_v, i_nh = pid % NV, (pid // NV).to(tl.int64)
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_n * T, i_n * T + T
o_k = tl.arange(0, BK)
o_v = i_v * BV + tl.arange(0, BV)
p_q = q + (bos * H + i_h) * K + o_k
p_k = k + (bos * H + i_h) * K + o_k
p_v = v + (bos * HV + i_hv) * V + o_v
if USE_G:
p_g = g + bos * HV + i_hv
if USE_GK:
p_gk = gk + (bos * HV + i_hv) * K + o_k
if USE_GV:
p_gv = gv + (bos * HV + i_hv) * V + o_v
if IS_BETA_HEADWISE:
p_beta = beta + bos * HV + i_hv
else:
p_beta = beta + (bos * HV + i_hv) * V + o_v
p_o = o + (bos * HV + i_hv) * V + o_v
mask_k = o_k < K
mask_v = o_v < V
if STATE_V_FIRST:
mask_h = mask_v[:, None] & mask_k[None, :]
else:
mask_h = mask_k[:, None] & mask_v[None, :]
if STATE_V_FIRST:
b_h = tl.zeros([BV, BK], dtype=tl.float32)
else:
b_h = tl.zeros([BK, BV], dtype=tl.float32)
if USE_INITIAL_STATE:
if STATE_V_FIRST:
p_h0 = h0 + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
else:
p_h0 = h0 + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
b_h += tl.load(p_h0, mask=mask_h, other=0).to(tl.float32)
for _ in tl.range(0, T):
b_q = tl.load(p_q, mask=mask_k, other=0).to(tl.float32)
b_k = tl.load(p_k, mask=mask_k, other=0).to(tl.float32)
b_v = tl.load(p_v, mask=mask_v, other=0).to(tl.float32)
if USE_QK_L2NORM_IN_KERNEL:
b_q = b_q / tl.sqrt(tl.sum(b_q * b_q) + 1e-6)
b_k = b_k / tl.sqrt(tl.sum(b_k * b_k) + 1e-6)
b_q = b_q * scale
if IS_BETA_HEADWISE:
b_beta = tl.load(p_beta).to(tl.float32)
else:
b_beta = tl.load(p_beta, mask=mask_v, other=0).to(tl.float32)
if APPLY_BETA_SIGMOID:
b_beta = tl.sigmoid(b_beta)
if ALLOW_NEG_EIGVAL:
b_beta = b_beta * 2
if USE_G:
b_g = tl.load(p_g).to(tl.float32)
if USE_GATE_IN_KERNEL:
b_A = tl.load(A_log + i_hv).to(tl.float32)
if HAS_DT_BIAS:
b_g = b_g + tl.load(dt_bias + i_hv).to(tl.float32)
b_g = -exp(b_A) * softplus(b_g)
b_h *= exp(b_g)
if USE_GK:
b_gk = tl.load(p_gk).to(tl.float32)
if STATE_V_FIRST:
b_h *= exp(b_gk[None, :])
else:
b_h *= exp(b_gk[:, None])
if USE_GV:
b_gv = tl.load(p_gv).to(tl.float32)
if STATE_V_FIRST:
b_h *= exp(b_gv[:, None])
else:
b_h *= exp(b_gv[None, :])
if STATE_V_FIRST:
b_v = b_beta * (b_v - tl.sum(b_h * b_k[None, :], 1))
b_h += b_v[:, None] * b_k[None, :]
b_o = tl.sum(b_h * b_q[None, :], 1)
else:
b_v = b_beta * (b_v - tl.sum(b_h * b_k[:, None], 0))
b_h += b_k[:, None] * b_v
b_o = tl.sum(b_h * b_q[:, None], 0)
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=mask_v)
p_q += H*K
p_k += H*K
p_v += HV*V
if USE_G:
p_g += HV
if USE_GK:
p_gk += HV*K
if USE_GV:
p_gv += HV*V
p_beta += HV * (1 if IS_BETA_HEADWISE else V)
p_o += HV*V
if STORE_FINAL_STATE:
if STATE_V_FIRST:
p_ht = ht + i_nh * K*V + o_v[:, None] * K + o_k[None, :]
else:
p_ht = ht + i_nh * K*V + o_k[:, None] * V + o_v[None, :]
tl.store(p_ht, b_h.to(p_ht.dtype.element_ty), mask=mask_h)
def fused_recurrent_gated_delta_rule_fwd(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
B, T, H, K, V = *k.shape, v.shape[-1]
HV = v.shape[2]
N = B if cu_seqlens is None else len(cu_seqlens) - 1
BK = triton.next_power_of_2(K)
BV = min(8, triton.next_power_of_2(V)) if gv is None else triton.next_power_of_2(V)
NV = triton.cdiv(V, BV)
o = torch.empty_like(v)
if output_final_state:
if state_v_first:
final_state = q.new_empty(N, HV, V, K, dtype=torch.float32)
else:
final_state = q.new_empty(N, HV, K, V, dtype=torch.float32)
else:
final_state = None
grid = (NV * N * HV,)
fused_recurrent_gated_delta_rule_fwd_kernel[grid](
q=q,
k=k,
v=v,
g=g,
gk=gk,
gv=gv,
beta=beta,
A_log=A_log,
dt_bias=dt_bias,
o=o,
h0=initial_state,
ht=final_state,
cu_seqlens=cu_seqlens,
scale=scale,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BK=BK,
BV=BV,
IS_BETA_HEADWISE=beta.ndim != v.ndim,
USE_QK_L2NORM_IN_KERNEL=use_qk_l2norm_in_kernel,
APPLY_BETA_SIGMOID=use_beta_sigmoid_in_kernel,
ALLOW_NEG_EIGVAL=allow_neg_eigval,
STATE_V_FIRST=state_v_first,
num_warps=1,
num_stages=3,
)
return o, final_state
class FusedRecurrentFunction(torch.autograd.Function):
@staticmethod
@input_guard
def forward(
ctx,
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
):
o, final_state = fused_recurrent_gated_delta_rule_fwd(
q=q,
k=k,
v=v,
g=g,
gk=gk,
gv=gv,
beta=beta,
A_log=A_log,
dt_bias=dt_bias,
scale=scale,
initial_state=initial_state,
output_final_state=output_final_state,
use_qk_l2norm_in_kernel=use_qk_l2norm_in_kernel,
use_beta_sigmoid_in_kernel=use_beta_sigmoid_in_kernel,
allow_neg_eigval=allow_neg_eigval,
state_v_first=state_v_first,
cu_seqlens=cu_seqlens,
)
return o, final_state
@staticmethod
@input_guard
def backward(ctx, do, dht):
raise NotImplementedError(
"Backward pass is not implemented yet and we do not have plans to implement it "
"because we haven't figured out how to compute dg without materializing the full "
"hidden states for all time steps.",
)
def fused_recurrent_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor | None = None,
gk: torch.Tensor | None = None,
gv: torch.Tensor | None = None,
beta: torch.Tensor | None = None,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
use_qk_l2norm_in_kernel: bool = False,
use_gate_in_kernel: bool = False,
A_log: torch.Tensor | None = None,
dt_bias: torch.Tensor | None = None,
use_beta_sigmoid_in_kernel: bool = False,
allow_neg_eigval: bool = False,
state_v_first: bool = False,
cu_seqlens: torch.LongTensor | None = None,
**kwargs,
) -> tuple[torch.Tensor, torch.Tensor]:
r"""
Args:
q (torch.Tensor):
queries of shape `[B, T, H, K]`.
k (torch.Tensor):
keys of shape `[B, T, H, K]`.
v (torch.Tensor):
values of shape `[B, T, HV, V]`.
GVA (Grouped Value Attention) is applied if `HV > H`, where `HV` must be divisible by `H`.
g (torch.Tensor):
g (decays) of shape `[B, T, HV]`. Default: `None`.
When `use_gate_in_kernel=False` (default), `g` must be in log space (pre-computed decay).
When `use_gate_in_kernel=True`, `g` is the raw pre-activation input; the kernel fuses
`-exp(A_log) * softplus(g + dt_bias)` internally per step.
gk (torch.Tensor):
gk (decays) of shape `[B, T, HV, K]`. Default: `None`.
gv (torch.Tensor):
gv (decays) of shape `[B, T, HV, V]`. Default: `None`.
beta (torch.Tensor):
betas of shape `[B, T, HV]`.
scale (Optional[float]):
Scale factor for the RetNet attention scores.
If not provided, it will default to `1 / sqrt(K)`. Default: `None`.
initial_state (Optional[torch.Tensor]):
Initial state of shape `[N, HV, K, V]` for `N` input sequences.
For equal-length input sequences, `N` equals the batch size `B`.
Default: `None`.
output_final_state (Optional[bool]):
Whether to output the final state of shape `[N, HV, K, V]`. Default: `False`.
use_qk_l2norm_in_kernel (Optional[bool]):
Whether to use L2 normalization in the kernel. Default: `False`.
use_gate_in_kernel (bool):
Whether to compute the log-space GDN decay internally.
When `True`, `g` is the raw input and `A_log` must be provided; the kernel fuses
gate activation into the recurrence. Default: `False`.
A_log (Optional[torch.Tensor]):
Decay parameter of shape `[HV]`. Required when `use_gate_in_kernel=True`.
dt_bias (Optional[torch.Tensor]):
Bias added to `g` before activation, of shape `[HV]`.
Only used when `use_gate_in_kernel=True`.
use_beta_sigmoid_in_kernel (Optional[bool]):
Whether to apply `torch.sigmoid(beta)` inside the kernel.
- If `True`, the passed `beta` acts as the raw beta logits.
- If `False`, `beta` is expected to already be in post-sigmoid space.
Default: `False`.
allow_neg_eigval (Optional[bool]):
Whether to allow negative eigenvalues by scaling `beta` to `[0, 2)`.
Only takes effect together with `use_beta_sigmoid_in_kernel=True`, in which case
the kernel computes `2 * sigmoid(beta)` instead of `sigmoid(beta)`. Default: `False`.
state_v_first (Optional[bool]):
Store the recurrent state in V-first ``[V, K]`` layout instead of the default ``[K, V]``. Default: ``False``.
cu_seqlens (torch.LongTensor):
Cumulative sequence lengths of shape `[N+1]` used for variable-length training,
consistent with the FlashAttention API.
Returns:
o (torch.Tensor):
Outputs of shape `[B, T, HV, V]`.
final_state (torch.Tensor):
Final state of shape `[N, HV, K, V]` if `output_final_state=True` else `None`.
Examples::
>>> import torch
>>> import torch.nn.functional as F
>>> from einops import rearrange
>>> from fla.ops.gated_delta_rule import fused_recurrent_gated_delta_rule
# inputs with equal lengths
>>> B, T, H, HV, K, V = 4, 2048, 4, 8, 512, 512
>>> q = torch.randn(B, T, H, K, device='cuda')
>>> k = F.normalize(torch.randn(B, T, H, K, device='cuda'), p=2, dim=-1)
>>> v = torch.randn(B, T, HV, V, device='cuda')
>>> g = F.logsigmoid(torch.rand(B, T, HV, device='cuda'))
>>> beta = torch.rand(B, T, HV, device='cuda').sigmoid()
>>> h0 = torch.randn(B, HV, K, V, device='cuda')
>>> o, ht = fused_gated_recurrent_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True
)
# for variable-length inputs, the batch size `B` is expected to be 1 and `cu_seqlens` is required
>>> q, k, v, g, beta = map(lambda x: rearrange(x, 'b t ... -> 1 (b t) ...'), (q, k, v, g, beta))
# for a batch with 4 sequences, `cu_seqlens` with 5 start/end positions are expected
>>> cu_seqlens = q.new_tensor([0, 2048, 4096, 6144, 8192], dtype=torch.long)
>>> o, ht = fused_gated_recurrent_delta_rule(
q, k, v, g, beta,
initial_state=h0,
output_final_state=True,
cu_seqlens=cu_seqlens
)
"""
if 'transpose_state_layout' in kwargs:
if state_v_first:
raise ValueError("Cannot pass both `state_v_first` and the deprecated `transpose_state_layout`.")
warnings.warn(
"`transpose_state_layout` is deprecated and renamed to `state_v_first`.",
DeprecationWarning,
stacklevel=2,
)
state_v_first = kwargs.pop('transpose_state_layout')
if cu_seqlens is not None:
if q.shape[0] != 1:
raise ValueError(
f"The batch size is expected to be 1 rather than {q.shape[0]} when using `cu_seqlens`."
f"Please flatten variable-length inputs before processing.",
)
if initial_state is not None and initial_state.shape[0] != len(cu_seqlens) - 1:
raise ValueError(
f"The number of initial states is expected to be equal to the number of input sequences, "
f"i.e., {len(cu_seqlens) - 1} rather than {initial_state.shape[0]}.",
)
if scale is None:
scale = k.shape[-1] ** -0.5
if beta is None:
beta = torch.ones_like(q[..., 0])
if use_gate_in_kernel:
if A_log is None:
raise ValueError("`A_log` must be provided when `use_gate_in_kernel=True`.")
if g is None:
raise ValueError("`g` (raw pre-activation) must be provided when `use_gate_in_kernel=True`.")
else:
A_log = None
dt_bias = None
if allow_neg_eigval and not use_beta_sigmoid_in_kernel:
raise ValueError("`allow_neg_eigval=True` requires `use_beta_sigmoid_in_kernel=True`.")
o, final_state = FusedRecurrentFunction.apply(
q,
k,
v,
g,
gk,
gv,
beta,
A_log,
dt_bias,
scale,
initial_state,
output_final_state,
use_qk_l2norm_in_kernel,
use_beta_sigmoid_in_kernel,
allow_neg_eigval,
state_v_first,
cu_seqlens,
)
return o, final_state
fused_recurrent_gdn = fused_recurrent_gated_delta_rule

344
ex_engine/fla_kernels/gated_delta_rule/gate.py Normal file
View File

@@ -0,0 +1,344 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import torch.nn.functional as F
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.index import prepare_chunk_indices
from fla.ops.utils.op import exp
from fla.ops.utils.softplus import softplus
from fla.utils import autocast_custom_bwd, autocast_custom_fwd, autotune_cache_kwargs, input_guard
def naive_gdn_gate(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
"""
Torch reference implementation for GDN gate computation.
Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
Args:
g (torch.Tensor):
Input tensor of shape `[..., HV]`.
A_log (torch.Tensor):
Decay parameter tensor with `HV` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[HV]`.
Returns:
Output tensor of shape `[..., HV]`.
"""
g = g.float()
if dt_bias is not None:
g = g + dt_bias.float()
return (-A_log.float().exp() * F.softplus(g)).to(output_dtype)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
'HAS_SCALE': lambda args: args['scale'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps)
for num_warps in [1, 2, 4, 8]
],
key=['H', 'BT', 'IS_VARLEN', 'REVERSE'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_chunk_cumsum_scalar_kernel(
g,
A_log,
dt_bias,
o,
scale,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
BT: tl.constexpr,
REVERSE: tl.constexpr,
HAS_BIAS: tl.constexpr,
HAS_SCALE: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + bos * H + i_h + o_t * H
p_o = o + bos * H + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_A = tl.load(A_log + i_h).to(tl.float32)
b_gate = -exp(b_A) * softplus(b_g)
b_o = tl.cumsum(b_gate, axis=0)
if REVERSE:
b_z = tl.sum(b_gate, axis=0)
b_o = -b_o + b_z[None] + b_gate
if HAS_SCALE:
b_o *= scale
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t)
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps)
for num_warps in [1, 2, 4, 8]
],
key=['H', 'BT'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_bwd_kernel(
g,
A_log,
dt_bias,
dyg,
dg,
dA,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
):
i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
b_A = tl.load(A_log + i_h).to(tl.float32)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + i_h + o_t * H
p_dg = dg + i_h + o_t * H
p_dyg = dyg + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
b_dyg = tl.load(p_dyg, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
# gate = -exp(A_log) * softplus(g + bias)
# d(gate)/d(g) = -exp(A_log) * sigmoid(g + bias) (softplus' = sigmoid)
# d(gate)/d(A_log) = -exp(A_log) * softplus(g + bias) = gate
b_neg_expA = -exp(b_A)
b_yg = b_neg_expA * softplus(b_g)
b_dg = b_neg_expA * (b_dyg * tl.sigmoid(b_g))
b_dA = tl.sum(b_dyg * b_yg, 0)
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
tl.store(dA + i_t * H + i_h, b_dA)
@input_guard
@dispatch('gated_delta_rule')
def gdn_gate_chunk_cumsum(
g: torch.Tensor,
A_log: torch.Tensor,
chunk_size: int,
scale: float = None,
dt_bias: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
output_dtype: torch.dtype | None = torch.float,
) -> torch.Tensor:
B, T, H = g.shape
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
o = torch.empty_like(g, dtype=output_dtype or g.dtype)
gdn_gate_chunk_cumsum_scalar_kernel[(NT, B * H)](
g=g,
A_log=A_log,
dt_bias=dt_bias,
o=o,
scale=scale,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
BT=BT,
REVERSE=False,
)
return o
@dispatch('gated_delta_rule')
def gdn_gate_bwd(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None,
dyg: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
H = g.shape[-1]
T = g.numel() // H
BT = 32
NT = triton.cdiv(T, BT)
dg = torch.empty_like(g, dtype=torch.float32)
dA = A_log.new_empty(NT, H, dtype=torch.float32)
gdn_gate_bwd_kernel[(NT, H)](
g=g,
A_log=A_log,
dt_bias=dt_bias,
dyg=dyg,
dg=dg,
dA=dA,
T=T,
H=H,
BT=BT,
)
dg = dg.view_as(g).type_as(g)
dA = dA.sum(0).view_as(A_log).type_as(A_log)
dbias = dg.view(-1, H).sum(0).to(dt_bias) if dt_bias is not None else None
return dg, dA, dbias
@triton.heuristics({
'HAS_BIAS': lambda args: args['dt_bias'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
for BT in [32, 64, 128]
for num_warps in [1, 2, 4, 8]
for num_stages in [2, 3]
],
key=['H'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def gdn_gate_fwd_kernel(
g,
A_log,
dt_bias,
yg,
T,
H: tl.constexpr,
BT: tl.constexpr,
HAS_BIAS: tl.constexpr,
):
i_t, i_h = tl.program_id(0).to(tl.int64), tl.program_id(1)
b_A = tl.load(A_log + i_h).to(tl.float32)
o_t = i_t * BT + tl.arange(0, BT)
m_t = o_t < T
p_g = g + i_h + o_t * H
p_yg = yg + i_h + o_t * H
b_g = tl.load(p_g, mask=m_t, other=0.0).to(tl.float32)
if HAS_BIAS:
b_g = b_g + tl.load(dt_bias + i_h).to(tl.float32)
b_yg = -exp(b_A) * softplus(b_g)
tl.store(p_yg, b_yg.to(p_yg.dtype.element_ty), mask=m_t)
@dispatch('gated_delta_rule')
def gdn_gate_fwd(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
H = g.shape[-1]
T = g.numel() // H
yg = torch.empty_like(g, dtype=output_dtype)
def grid(meta):
return (triton.cdiv(T, meta['BT']), H)
gdn_gate_fwd_kernel[grid](
g=g,
A_log=A_log,
dt_bias=dt_bias,
yg=yg,
T=T,
H=H,
)
return yg
class GDNGateFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
def forward(
ctx,
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
yg = gdn_gate_fwd(g=g, A_log=A_log, dt_bias=dt_bias, output_dtype=output_dtype)
ctx.save_for_backward(g, A_log, dt_bias)
return yg
@staticmethod
@input_guard
@autocast_custom_bwd
def backward(ctx, dyg: torch.Tensor):
g, A_log, dt_bias = ctx.saved_tensors
dg, dA, dbias = gdn_gate_bwd(g=g, A_log=A_log, dt_bias=dt_bias, dyg=dyg)
return dg, dA, dbias, None
@torch.compiler.disable
def fused_gdn_gate(
g: torch.Tensor,
A_log: torch.Tensor,
dt_bias: torch.Tensor | None = None,
output_dtype: torch.dtype = torch.float32,
) -> torch.Tensor:
r"""
Fused GDN gate computation with autograd support.
Computes: ``g = -A_log.exp() * softplus(g + dt_bias)``
Args:
g (torch.Tensor):
Input tensor of shape `[..., HV]`.
A_log (torch.Tensor):
Decay parameter tensor with `HV` elements.
dt_bias (torch.Tensor | None):
Optional bias tensor added to `g` before activation, shape `[HV]`.
output_dtype (torch.dtype):
The dtype of the output tensor. Default: `torch.float32`.
Returns:
Output tensor of shape `[..., HV]`.
"""
return GDNGateFunction.apply(g, A_log, dt_bias, output_dtype)

161
ex_engine/fla_kernels/gated_delta_rule/naive.py Normal file
View File

@@ -0,0 +1,161 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import torch.nn.functional as F
from einops import rearrange
def naive_recurrent_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
g: torch.Tensor,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
):
"""
Reference PyTorch implementation of recurrent gated delta rule.
Args:
q: [B, T, H, K]
k: [B, T, H, K]
v: [B, T, H, V]
beta: [B, T, H]
g: [B, T, H]
scale: float, optional
initial_state: [B, H, K, V], optional
output_final_state: bool
Returns:
o: [B, T, H, V]
final_state: [B, H, K, V] if output_final_state else None
"""
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
B, H, T, K, V = *k.shape, v.shape[-1]
o = torch.zeros(B, H, T, V).to(v)
h = torch.zeros(B, H, K, V).to(v)
if initial_state is not None:
h = initial_state.to(torch.float32)
if scale is None:
scale = 1 / (q.shape[-1] ** 0.5)
q = q * scale
for i in range(T):
b_q = q[:, :, i]
b_k = k[:, :, i]
b_v = v[:, :, i].clone()
h = h.clone() * g[:, :, i].exp()[..., None, None]
b_beta = beta[:, :, i]
b_v = b_v - (h.clone() * b_k[..., None]).sum(-2)
b_v = b_v * b_beta[..., None]
h = h.clone() + b_k.unsqueeze(-1) * b_v.unsqueeze(-2)
o[:, :, i] = torch.einsum('bhd,bhdm->bhm', b_q, h)
if not output_final_state:
h = None
o = o.transpose(1, 2).contiguous()
return o, h
def naive_chunk_gated_delta_rule(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
g: torch.Tensor,
beta: torch.Tensor,
chunk_size: int = 64,
scale: float = None,
initial_state: torch.Tensor = None,
output_final_state: bool = False,
):
"""
Reference PyTorch implementation of chunk gated delta rule.
Args:
q: [B, T, H, K]
k: [B, T, H, K]
v: [B, T, H, V]
g: [B, T, H]
beta: [B, T, H]
chunk_size: int
scale: float, optional
initial_state: [B, H, K, V], optional
output_final_state: bool
Returns:
o: [B, T, H, V]
final_state: [B, H, K, V] if output_final_state else None
"""
BT = chunk_size
if scale is None:
scale = 1 / (q.shape[-1] ** 0.5)
q, k, v, beta, g = map(lambda x: x.transpose(1, 2).contiguous().to(torch.float32), [q, k, v, beta, g])
T = q.shape[-2]
pad_len = (BT - (T % BT)) % BT
if pad_len > 0:
q = F.pad(q, (0, 0, 0, pad_len))
k = F.pad(k, (0, 0, 0, pad_len))
v = F.pad(v, (0, 0, 0, pad_len))
beta = F.pad(beta, (0, pad_len))
g = F.pad(g, (0, pad_len))
q, k, v, beta, g = map(lambda x: x.to(torch.float32), [q, k, v, beta, g])
decay = g
chunk_size = BT
b, h, l, d_k = q.shape
d_v = v.shape[-1]
q = q * scale
v = v * beta[..., None]
k_beta = k * beta[..., None]
assert l % chunk_size == 0
# note that diagonal is masked.
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0)
q, k, v, k_beta, decay = map(
lambda x: rearrange(x, 'b h (n c) d -> b h n c d', c=chunk_size),
[q, k, v, k_beta, decay.unsqueeze(-1)],
)
decay = decay.squeeze(-1).cumsum(-1)
decay_exp = decay.exp()[..., None]
L_mask = ((decay.unsqueeze(-1) - decay.unsqueeze(-2)).tril().exp().float()).tril()
attn = -((k_beta @ k.transpose(-1, -2)) * L_mask).masked_fill(mask, 0)
for i in range(1, chunk_size):
attn[..., i, :i] = attn[..., i, :i].clone() + (attn[..., i, :i, None].clone() * attn[..., :i, :i].clone()).sum(-2)
attn = attn + torch.eye(chunk_size, dtype=torch.float, device=q.device)
attn = attn
k_cumsum = attn @ v
k_cumdecay = attn @ (k_beta * decay_exp)
v = k_cumsum
S = k.new_zeros(b, h, d_k, d_v)
if initial_state is not None:
S = initial_state.to(torch.float32)
o = torch.zeros_like(v)
mask = torch.triu(torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1)
for i in range(0, l // chunk_size):
q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i]
attn = (q_i @ k_i.transpose(-1, -2) * L_mask[:, :, i]).masked_fill_(mask, 0)
v_prime = (k_cumdecay[:, :, i]) @ S
v_new = v_i - v_prime
o_inter = (q_i * decay[:, :, i, :, None].exp()) @ S
o[:, :, i] = o_inter + attn @ v_new
S = S * decay[:, :, i, -1, None, None].exp() + (k_i * (decay[:, :, i, -1, None] - decay[:, :, i]).exp()
[..., None]).transpose(-1, -2) @ v_new
if not output_final_state:
S = None
# unpad
o = rearrange(o, 'b h n c d -> b h (n c) d')
o = o[:, :, :T]
o = o.transpose(1, 2)
return o, S

351
ex_engine/fla_kernels/gated_delta_rule/wy_fast.py Normal file
View File

@@ -0,0 +1,351 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import torch
import triton
import triton.language as tl
from fla.ops.backends import dispatch
from fla.ops.utils import prepare_chunk_indices
from fla.ops.utils.cache import fla_cache_autotune
from fla.ops.utils.op import exp2
from fla.utils import IS_INTEL, IS_NVIDIA_BLACKWELL, autotune_cache_kwargs, check_shared_mem
# Blackwell can select unstable Triton configs for prepare_wy_repr_bwd_kernel
# during autotuning (see #913). Restrict it to the config that has been
# validated on B200 until the wider config space is re-validated.
PREPARE_WY_REPR_BWD_NUM_WARPS = [2] if IS_NVIDIA_BLACKWELL else [2, 4]
PREPARE_WY_REPR_BWD_NUM_STAGES = [4] if IS_NVIDIA_BLACKWELL else [2, 3, 4]
# Intel keeps scaling past the warp counts NVIDIA prefers: 16 warps is ~1.3x faster
# than 8 for recompute_w_u.
RECOMPUTE_W_U_NUM_WARPS = [2, 4, 8, 16] if IS_INTEL else [2, 4, 8]
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in RECOMPUTE_W_U_NUM_WARPS
for num_stages in [2, 3, 4]
],
key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def recompute_w_u_fwd_kernel(
k,
v,
beta,
w,
u,
A,
g,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
o_A = tl.arange(0, BT)
m_t = o_t < T
m_A = m_t[:, None] & (o_A[None, :] < BT)
p_b = beta + bos*HV + i_h + o_t * HV
b_b = tl.load(p_b, mask=m_t, other=0.0)
p_A = A + (bos*HV + i_h) * BT + o_t[:, None] * (HV*BT) + o_A[None, :]
b_A = tl.load(p_A, mask=m_A, other=0.0)
for i_v in range(tl.cdiv(V, BV)):
o_v = i_v * BV + tl.arange(0, BV)
m_v = m_t[:, None] & (o_v[None, :] < V)
p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_u = u + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
b_v = tl.load(p_v, mask=m_v, other=0.0)
b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
b_u = tl.dot(b_A, b_vb, allow_tf32=False)
tl.store(p_u, b_u.to(p_u.dtype.element_ty), mask=m_v)
if USE_G:
p_g = g + (bos*HV + i_h) + o_t * HV
b_g = exp2(tl.load(p_g, mask=m_t, other=0.0))
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_w = w + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
b_k = tl.load(p_k, mask=m_k, other=0.0)
b_kb = b_k * b_b[:, None]
if USE_G:
b_kb *= b_g[:, None]
b_w = tl.dot(b_A, b_kb.to(b_k.dtype))
tl.store(p_w, b_w.to(p_w.dtype.element_ty), mask=m_k)
@triton.heuristics({
'USE_G': lambda args: args['g'] is not None,
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
})
@fla_cache_autotune(
configs=[
triton.Config({}, num_warps=num_warps, num_stages=num_stages)
for num_warps in PREPARE_WY_REPR_BWD_NUM_WARPS
for num_stages in PREPARE_WY_REPR_BWD_NUM_STAGES
],
key=['H', 'HV', 'K', 'V', 'BT', 'BK', 'BV', 'IS_VARLEN'],
**autotune_cache_kwargs,
)
@triton.jit(do_not_specialize=['T'])
def prepare_wy_repr_bwd_kernel(
k,
v,
beta,
g,
A,
dw,
du,
dk,
dv,
db,
dg,
cu_seqlens,
chunk_indices,
T,
H: tl.constexpr,
HV: tl.constexpr,
K: tl.constexpr,
V: tl.constexpr,
BT: tl.constexpr,
BK: tl.constexpr,
BV: tl.constexpr,
USE_G: tl.constexpr,
IS_VARLEN: tl.constexpr,
):
i_t, i_bh = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
i_b, i_h = i_bh // HV, i_bh % HV
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int64)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(cu_seqlens + i_n + 1).to(tl.int64)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
o_t = i_t * BT + tl.arange(0, BT)
o_A = tl.arange(0, BT)
m_t = o_t < T
m_AT = (o_A[:, None] < BT) & m_t[None, :]
p_b = beta + (bos*HV + i_h) + o_t * HV
p_db = db + (bos*HV + i_h) + o_t * HV
p_A = A + (bos*HV + i_h) * BT + o_A[:, None] + o_t[None, :] * (HV*BT)
b_b = tl.load(p_b, mask=m_t, other=0.0)
b_db = tl.zeros([BT], dtype=tl.float32)
b_A = tl.load(p_A, mask=m_AT, other=0.0)
b_dA = tl.zeros([BT, BT], dtype=tl.float32)
if USE_G:
p_g = g + (bos*HV + i_h) + o_t * HV
b_g = tl.load(p_g, mask=m_t, other=0.0)
b_g_exp = exp2(b_g)
b_dg = tl.zeros([BT], dtype=tl.float32)
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
p_dw = dw + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
# [BT, BK]
b_k = tl.load(p_k, mask=m_k, other=0.0)
if USE_G:
b_kbg = b_k * (b_b * b_g_exp)[:, None]
else:
b_kbg = b_k * b_b[:, None]
b_dw = tl.load(p_dw, mask=m_k, other=0.0)
b_dA += tl.dot(b_dw, tl.trans(b_kbg).to(b_dw.dtype))
b_dkbg = tl.dot(b_A, b_dw)
if USE_G:
b_dk = b_dkbg * (b_g_exp * b_b)[:, None]
b_db += tl.sum(b_dkbg * b_k * b_g_exp[:, None], 1)
b_dg += tl.sum(b_dkbg * b_kbg, 1)
else:
b_dk = b_dkbg * b_b[:, None]
b_db += tl.sum(b_dkbg * b_k, 1)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
for i_v in range(tl.cdiv(V, BV)):
o_v = i_v * BV + tl.arange(0, BV)
m_v = m_t[:, None] & (o_v[None, :] < V)
p_v = v + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_dv = dv + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
p_du = du + (bos*HV + i_h) * V + o_t[:, None] * (HV*V) + o_v[None, :]
b_v = tl.load(p_v, mask=m_v, other=0.0)
b_vb = (b_v * b_b[:, None]).to(b_v.dtype)
b_du = tl.load(p_du, mask=m_v, other=0.0)
b_dA += tl.dot(b_du, tl.trans(b_vb))
b_dvb = tl.dot(b_A, b_du)
b_dv = b_dvb * b_b[:, None]
b_db += tl.sum(b_dvb * b_v, 1)
tl.store(p_dv, b_dv.to(p_dv.dtype.element_ty), mask=m_v)
m_A = (o_t[:, None] > o_t[None, :]) & (m_t[:, None] & m_t)
b_dA = tl.where(m_A, b_dA, 0)
b_dA = tl.dot(b_dA.to(b_A.dtype), b_A)
b_dA = tl.dot(b_A, b_dA.to(b_A.dtype))
if USE_G:
b_dA *= exp2(b_g[:, None] - b_g[None, :])
b_A = tl.zeros([BT, BT], dtype=tl.float32)
b_dA = tl.where(m_A, -b_dA, 0).to(k.dtype.element_ty)
tl.debug_barrier()
for i_k in range(tl.cdiv(K, BK)):
o_k = i_k * BK + tl.arange(0, BK)
m_k = m_t[:, None] & (o_k[None, :] < K)
p_k = k + (bos*H + i_h // (HV // H)) * K + o_t[:, None] * (H*K) + o_k[None, :]
p_dk = dk + (bos*HV + i_h) * K + o_t[:, None] * (HV*K) + o_k[None, :]
b_k = tl.load(p_k, mask=m_k, other=0.0)
b_kt = tl.trans(b_k)
b_kb = b_k * b_b[:, None]
b_A += tl.dot(b_k, b_kt)
b_dkb = tl.dot(b_dA, b_k)
b_db += tl.sum(b_dkb * b_k, 1)
b_dk = b_dkb * b_b[:, None] + tl.trans(tl.dot(tl.trans(b_kb).to(b_dA.dtype), b_dA))
b_dk += tl.load(p_dk, mask=m_k, other=0.0)
tl.store(p_dk, b_dk.to(p_dk.dtype.element_ty), mask=m_k)
tl.store(p_db, b_db.to(p_db.dtype.element_ty), mask=m_t)
b_A *= b_b[:, None]
if USE_G:
b_AdA = b_dA * b_A
p_dg = dg + (bos*HV + i_h) + o_t * HV
b_dg += tl.sum(b_AdA, axis=1) - tl.sum(b_AdA, axis=0)
tl.store(p_dg, b_dg.to(p_dg.dtype.element_ty), mask=m_t)
@dispatch('gated_delta_rule')
def recompute_w_u_fwd(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
g: torch.Tensor | None = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
BT = A.shape[-1]
BK = 64
BV = 64
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
w = k.new_empty(B, T, HV, K)
u = torch.empty_like(v)
recompute_w_u_fwd_kernel[(NT, B*HV)](
k=k,
v=v,
beta=beta,
w=w,
u=u,
A=A,
g=g,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BT=BT,
BK=BK,
BV=BV,
)
return w, u
@dispatch('gated_delta_rule')
def prepare_wy_repr_bwd(
k: torch.Tensor,
v: torch.Tensor,
beta: torch.Tensor,
A: torch.Tensor,
dw: torch.Tensor,
du: torch.Tensor,
g: torch.Tensor = None,
cu_seqlens: torch.LongTensor | None = None,
chunk_indices: torch.LongTensor | None = None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
B, T, H, K, V, HV = *k.shape, v.shape[-1], v.shape[2]
BT = A.shape[-1]
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
CONST_TILING = 64 if check_shared_mem() else 32
BK = min(max(triton.next_power_of_2(K), 16), CONST_TILING)
BV = min(max(triton.next_power_of_2(V), 16), CONST_TILING)
dk = k.new_empty(B, T, HV, K)
dv = torch.empty_like(v)
dg = torch.empty_like(g) if g is not None else None
db = torch.empty_like(beta)
prepare_wy_repr_bwd_kernel[(NT, B * HV)](
k=k,
v=v,
beta=beta,
g=g,
A=A,
dw=dw,
du=du,
dk=dk,
dv=dv,
db=db,
dg=dg,
cu_seqlens=cu_seqlens,
chunk_indices=chunk_indices,
T=T,
H=H,
HV=HV,
K=K,
V=V,
BT=BT,
BK=BK,
BV=BV,
)
if H != HV:
dk = dk.view(B, T, H, HV // H, K).sum(3)
return dk, dv, db, dg
fwd_recompute_w_u = recompute_w_u_fwd
bwd_prepare_wy_repr = prepare_wy_repr_bwd

65
ex_engine/fla_kernels/utils/__init__.py Normal file
View File

@@ -0,0 +1,65 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
from .csr import prepare_block_csr
from .cumsum import (
chunk_global_cumsum,
chunk_global_cumsum_scalar,
chunk_global_cumsum_vector,
chunk_local_cumsum,
chunk_local_cumsum_scalar,
chunk_local_cumsum_vector,
)
from .index import (
get_max_num_splits,
prepare_chunk_indices,
prepare_chunk_offsets,
prepare_cu_seqlens_from_lens,
prepare_cu_seqlens_from_mask,
prepare_lens,
prepare_lens_from_mask,
prepare_position_ids,
prepare_sequence_ids,
prepare_token_indices,
)
from .logsumexp import logsumexp_fwd
from .matmul import addmm, matmul
from .pack import pack_sequence, unpack_sequence
from .pooling import mean_pooling
from .softmax import softmax_bwd, softmax_fwd
from .softplus import softplus
from .solve_tril import solve_tril
__all__ = [
"addmm",
"chunk_global_cumsum",
"chunk_global_cumsum_scalar",
"chunk_global_cumsum_vector",
"chunk_local_cumsum",
"chunk_local_cumsum_scalar",
"chunk_local_cumsum_vector",
"get_max_num_splits",
"logsumexp_fwd",
"matmul",
"mean_pooling",
"pack_sequence",
"prepare_block_csr",
"prepare_chunk_indices",
"prepare_chunk_offsets",
"prepare_cu_seqlens_from_lens",
"prepare_cu_seqlens_from_mask",
"prepare_lens",
"prepare_lens_from_mask",
"prepare_position_ids",
"prepare_sequence_ids",
"prepare_token_indices",
"softmax_bwd",
"softmax_fwd",
"softplus",
"solve_tril",
"unpack_sequence",
]

449
ex_engine/fla_kernels/utils/cache.py Normal file
View File

@@ -0,0 +1,449 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import dataclasses
import enum
import json
import logging
import os
import re
from functools import cache, lru_cache
from pathlib import Path
from typing import Any
import torch
import triton
from packaging import version
from triton.runtime.autotuner import Autotuner
TRITON_ABOVE_3_5_1 = version.parse(triton.__version__) >= version.parse("3.5.1")
TRITON_ABOVE_3_4_0 = version.parse(triton.__version__) >= version.parse("3.4.0")
class FlaCacheMode(enum.Enum):
"""Controls how FLA loads kernel configs from its config cache (FLA_CACHE_MODE env var).
DISABLED — skip all cache lookups, always fall back to Triton autotune (default when FLA_CACHE_MODE is unset)
STRICT — exact key match only; falls back to Triton autotune if no match
FUZZY — exact key match → fuzzy key match; falls back to Triton autotune if no match
FULL — exact key match → fuzzy key match → default_config fallback
DEFAULT — use only the top-level default_config field, skip key-based lookup
ALWAYS — like DEFAULT, but re-reads config files on every kernel call;
useful for debugging: edit default_config in a JSON file and the next
kernel call picks it up without restarting the process
"""
DISABLED = "disabled"
STRICT = "strict"
FUZZY = "fuzzy"
FULL = "full"
DEFAULT = "default"
ALWAYS = "always"
def uses_default_config(self) -> bool:
"""Return True for modes that may fall back to default_config (FULL, DEFAULT, ALWAYS)."""
return self in (FlaCacheMode.FULL, FlaCacheMode.DEFAULT, FlaCacheMode.ALWAYS)
@classmethod
def from_env(cls) -> "FlaCacheMode":
mode_str = os.environ.get("FLA_CACHE_MODE", cls.DISABLED.value)
try:
return cls(mode_str)
except ValueError:
valid = [m.value for m in cls]
raise ValueError(
f"Invalid FLA_CACHE_MODE={mode_str!r}. Valid values: {valid}"
) from None
FLA_CACHE_MODE: FlaCacheMode = FlaCacheMode.from_env()
logger = logging.getLogger(__name__)
def sanitize_gpu_name(gpu_name: str) -> str:
sanitized = re.sub(r"[^0-9A-Za-z]+", "_", gpu_name)
sanitized = sanitized.strip("_")
return sanitized or "unknown_gpu"
@lru_cache(maxsize=1)
def get_gpu_info():
"""Get GPU model information.
This function detects the GPU model and returns a sanitized string identifier.
It prioritizes FLA_GPU_NAME environment variable if set, then detects from
available hardware (CUDA, ROCm, Intel GPU, or CPU).
"""
# Check if GPU name is overridden via environment variable
gpu_name = None
# Check if GPU name is overridden via environment variable
if "FLA_GPU_NAME" in os.environ:
gpu_name = os.environ["FLA_GPU_NAME"]
# Try to get device name based on availability
elif torch.cuda.is_available():
# Works for both NVIDIA and AMD GPUs (ROCm)
gpu_name = torch.cuda.get_device_name(0)
elif hasattr(torch, 'xpu') and torch.xpu.is_available():
gpu_name = torch.xpu.get_device_name(0)
if gpu_name:
return sanitize_gpu_name(gpu_name)
# Default to CPU if no GPU available
return "cpu"
def get_fla_config_dir() -> Path:
"""Get FLA's configs directory.
The directory can be overridden by setting the FLA_CONFIG_DIR environment variable.
If set, configs will be loaded directly from $FLA_CONFIG_DIR/. Otherwise FLA
falls back to the default fla/configs/{GPU}/ directory in the project.
"""
# Check if custom config dir is set via environment variable
if "FLA_CONFIG_DIR" in os.environ:
return Path(os.environ["FLA_CONFIG_DIR"])
# Default: project_dir/fla/configs/{GPU}/
project_dir = Path(__file__).parent.parent.parent
return project_dir / "configs" / get_gpu_info()
@dataclasses.dataclass(frozen=True)
class AutotuneKey:
"""Autotune key with exact/fuzzy matching, serialization, and construction helpers."""
autotune_key: tuple[Any, ...]
@staticmethod
def normalize_autotune_key(value: Any) -> Any:
if isinstance(value, (list, tuple)):
return [AutotuneKey.normalize_autotune_key(v) for v in value]
if isinstance(value, dict):
return {k: AutotuneKey.normalize_autotune_key(v) for k, v in value.items()}
return value
@staticmethod
def serialize(key: Any) -> str:
return json.dumps(AutotuneKey.normalize_autotune_key(key), separators=(",", ":"), sort_keys=True)
@staticmethod
def key_hash(key: Any) -> str:
import hashlib
return hashlib.md5(AutotuneKey.serialize(key).encode()).hexdigest()
@staticmethod
def is_numeric(value: Any) -> bool:
return isinstance(value, (int, float)) and not isinstance(value, bool)
@staticmethod
def keys_fuzzy_match(cached_key: Any, requested_key: Any) -> bool:
# Fuzzy match: numeric leaves are compatible regardless of their actual numeric values
# (e.g. a config tuned for seq_len=1024 can apply to seq_len=2048).
# Structure (type, length, dict keys) must still match exactly.
if AutotuneKey.is_numeric(cached_key) and AutotuneKey.is_numeric(requested_key):
return True
if isinstance(cached_key, (list, tuple)) and isinstance(requested_key, (list, tuple)):
return len(cached_key) == len(requested_key) and all(
AutotuneKey.keys_fuzzy_match(c, r) for c, r in zip(cached_key, requested_key)
)
if isinstance(cached_key, dict) and isinstance(requested_key, dict):
return cached_key.keys() == requested_key.keys() and all(
AutotuneKey.keys_fuzzy_match(cached_key[k], requested_key[k]) for k in cached_key
)
return cached_key == requested_key
@classmethod
def build(
cls,
arg_names: list[str],
key_names: list[str],
positional_args: tuple[Any, ...],
runtime_kwargs: dict[str, Any],
) -> "AutotuneKey":
named_args = dict(zip(arg_names, positional_args))
all_args = {**named_args, **runtime_kwargs}
tracked_args = {k: v for (k, v) in all_args.items() if k in arg_names}
tuning_key = [tracked_args[name] for name in key_names if name in tracked_args]
for arg in tracked_args.values():
if hasattr(arg, "dtype"):
tuning_key.append(str(arg.dtype))
return cls(autotune_key=tuple(tuning_key))
def exact_matches(self, entry_key: Any) -> bool:
return self.serialize(self.autotune_key) == self.serialize(entry_key)
def fuzzy_matches(self, entry_key: Any) -> bool:
self_normalized = self.normalize_autotune_key(self.autotune_key)
entry_normalized = self.normalize_autotune_key(entry_key)
return (
isinstance(self_normalized, list)
and isinstance(entry_normalized, list)
and len(self_normalized) == len(entry_normalized)
and AutotuneKey.keys_fuzzy_match(self_normalized, entry_normalized)
)
@dataclasses.dataclass(frozen=True)
class KernelConfigFile:
"""Validated in-memory representation of a {kernel_name}.json config file."""
kernel_name: str | None
triton_version: str | None
autotune_entries: dict[str, dict[str, Any]] | None
default_config: dict[str, Any] | None
@classmethod
def from_dict(cls, config_file: Path, data: Any) -> "KernelConfigFile | None":
"""Parse and validate a raw JSON dict. Returns None (with a warning) if malformed."""
def fail(msg, *args):
logger.warning(msg, *args)
raise ValueError
try:
if not isinstance(data, dict):
fail("Malformed config %s: root is %s, expected dict", config_file, type(data).__name__)
raw_entries = data.get("autotune_entries")
entries: dict[str, dict[str, Any]] | None = None
if raw_entries is not None:
if not isinstance(raw_entries, dict):
fail("Malformed config %s: 'autotune_entries' is %s, expected dict",
config_file, type(raw_entries).__name__)
for h, entry in raw_entries.items():
if not isinstance(entry, dict):
fail("Malformed config %s: autotune_entries[%r] is %s, expected dict",
config_file, h, type(entry).__name__)
if not isinstance(entry.get("config"), dict):
fail("Malformed config %s: autotune_entries[%r] missing valid 'config' field", config_file, h)
entries = raw_entries
default_config = data.get("default_config")
if default_config is not None and not isinstance(default_config, dict):
fail("Malformed config %s: 'default_config' is %s, expected dict", config_file, type(default_config).__name__)
return cls(
kernel_name=data.get("kernel_name"),
triton_version=data.get("triton_version"),
autotune_entries=entries,
default_config=default_config,
)
except ValueError:
return None
@classmethod
def from_file(cls, config_file: Path) -> "KernelConfigFile | None":
"""Read and validate a config file. Returns None if the file is missing or malformed."""
config_data = read_config_file(config_file)
if config_data is None:
return None
return cls.from_dict(config_file, config_data)
def lookup_exact(self, key: AutotuneKey) -> dict[str, Any] | None:
if self.autotune_entries is None:
return None
return self.autotune_entries.get(AutotuneKey.key_hash(key.autotune_key))
def lookup_fuzzy(self, key: AutotuneKey) -> dict[str, Any] | None:
if self.autotune_entries is None:
return None
for entry in self.autotune_entries.values():
if key.fuzzy_matches(entry.get("autotune_key")):
return entry
return None
@cache
def load_config_file(config_file: Path) -> dict[str, Any] | None:
try:
with open(config_file) as f:
return json.load(f)
except Exception as e:
logger.warning("Error reading config file %s: %s", config_file, e)
return None
def read_config_file(config_file: Path) -> dict[str, Any] | None:
"""Read a config file, bypassing the in-process cache in ALWAYS mode."""
if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return load_config_file.__wrapped__(config_file)
return load_config_file(config_file)
def load_cached_config(kernel_name: str, autotune_key: AutotuneKey | None = None) -> dict[str, Any] | None:
"""
Load cached best config for a kernel from FLA configs directory.
This function loads the cached best configuration for a given kernel name
from get_fla_config_dir()/{kernel_name}.json.
Cache files may contain multiple autotune entries keyed by Triton's
runtime tuning key plus a top-level default config.
If the config file is not found or cannot be loaded, a warning is printed
and None is returned, allowing fallback to Triton's autotune.
The lookup mode is controlled by the FLA_CACHE_MODE environment variable (see FlaCacheMode).
Args:
kernel_name: Name of the kernel (e.g., "causal_conv1d_fwd_kernel")
autotune_key: Triton autotune key for the current invocation
Returns:
Best config dictionary or None if not found or disabled
"""
if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
return None
config_dir = get_fla_config_dir()
config_file = config_dir / f"{kernel_name}.json"
if not config_file.exists():
return None
config_data = read_config_file(config_file)
if config_data is None:
return None
config = KernelConfigFile.from_dict(config_file, config_data)
if config is None:
return None
if FLA_CACHE_MODE is FlaCacheMode.DEFAULT or FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return config.default_config
# STRICT mode: exact match only, no fuzzy fallback
if FLA_CACHE_MODE is FlaCacheMode.STRICT:
if autotune_key is not None:
entry = config.lookup_exact(autotune_key)
if entry is not None:
return entry["config"]
return None
# FULL and FUZZY modes: try exact key match first, then fuzzy match
if autotune_key is not None:
entry = config.lookup_exact(autotune_key) or config.lookup_fuzzy(autotune_key)
if entry is not None:
return entry["config"]
if FLA_CACHE_MODE is FlaCacheMode.FUZZY:
return None
# FULL mode: fall back to default_config, then legacy raw config (no autotune_entries)
if config.default_config is not None:
return config.default_config
if config.autotune_entries is not None:
return None
return config_data
class CachedAutotuner(Autotuner):
"""
A modified autotuner that loads best config from FLA's config directory.
This class extends Triton's Autotuner but overrides the run method to
try loading cached configuration first before falling back to autotune.
"""
def __init__(self, fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs):
super().__init__(fn, arg_names, configs, key, reset_to_zero, restore_value, **kwargs)
self.kernel_name = fn.fn.__name__ if hasattr(fn, 'fn') else fn.__name__
# None-safe pre/post hooks: Triton's defaults crash when a restore_value / reset_to_zero arg
# is None (idiomatic for optional pointers gated by a tl.constexpr flag).
# Fixed upstream in triton-lang/triton#10295 — remove this override once FLA's minimum Triton version has it.
if not self.user_defined_pre_hook and (self.reset_to_zero or self.restore_value):
def _pre_hook(kw, reset_only=False):
for n in self.reset_to_zero:
if kw[n] is not None:
kw[n].zero_()
if not reset_only:
self.restore_copies = {n: kw[n].clone() for n in self.restore_value if kw[n] is not None}
self.pre_hook = _pre_hook
if not self.user_defined_post_hook and self.restore_value:
def _post_hook(kw, exception):
for n, copy in self.restore_copies.items():
kw[n].copy_(copy)
self.restore_copies = {}
self.post_hook = _post_hook
def should_check_fla_cache(self, key: AutotuneKey) -> bool:
if FLA_CACHE_MODE is FlaCacheMode.DISABLED:
return False
if FLA_CACHE_MODE is FlaCacheMode.ALWAYS:
return True
return key.autotune_key not in self.cache
def run(self, *args, **kwargs):
key = AutotuneKey.build(self.arg_names, self.keys, args, kwargs)
if self.should_check_fla_cache(key):
self.maybe_load_cached_config(key)
return super().run(*args, **kwargs)
def maybe_load_cached_config(self, key: AutotuneKey):
best_config = load_cached_config(self.kernel_name, key)
if best_config is not None:
kw = best_config["kwargs"]
num_warps = best_config["num_warps"]
num_stages = best_config["num_stages"]
extra = {
"num_ctas": best_config["num_ctas"],
"maxnreg": best_config.get("maxnreg"),
"pre_hook": None,
"ir_override": best_config.get("ir_override"),
} if TRITON_ABOVE_3_5_1 else {}
cfg = triton.Config(kw, num_warps=num_warps, num_stages=num_stages, **extra)
self.cache[key.autotune_key] = cfg
else:
logger.debug(
"No cached config found for kernel %s and key %s; falling back to Triton autotune",
self.kernel_name,
list(key.autotune_key),
)
def fla_cache_autotune(configs, key=None, prune_configs_by=None, reset_to_zero=None, restore_value=None,
pre_hook=None, post_hook=None, warmup=None, rep=None, use_cuda_graph=False,
do_bench=None, cache_results=False):
"""
Decorator for auto-tuning a :code:`triton.jit`'d function with FLA config support.
Extends Triton's autotune to load best configurations from FLA's config directory
(default: fla/configs/{GPU}/, or FLA_CONFIG_DIR/ when overridden), keyed by kernel
name from {kernel_name}.json. Lookup behaviour is controlled by FLA_CACHE_MODE.
Falls back to normal Triton autotuning when no cached config is found.
"""
# key can be None when we want to use cache only (no fallback autotune)
if key is None:
key = []
def decorator(fn):
kwargs = {}
if TRITON_ABOVE_3_4_0:
kwargs = {"cache_results": cache_results}
return CachedAutotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value,
pre_hook=pre_hook, post_hook=post_hook,
prune_configs_by=prune_configs_by, warmup=warmup, rep=rep,
use_cuda_graph=use_cuda_graph, do_bench=do_bench,
**kwargs,
)
return decorator
def configure_fla_cache_autotune():
triton.autotune = fla_cache_autotune
logger.info(
"configure_fla_cache_autotune() is enabling FLA fla_cache_autotune; "
"triton.autotune will be replaced with fla_cache_autotune."
)
def restore_autotune_backend():
from triton.runtime.autotuner import autotune as original_autotune
triton.autotune = original_autotune
logger.info(
"restore_autotune_backend() is restoring Triton's original autotune; "
"triton.autotune will be replaced with triton.runtime.autotuner.autotune."
)

101
ex_engine/fla_kernels/utils/op.py Normal file
View File

@@ -0,0 +1,101 @@
# Copyright (c) 2023-2026, Songlin Yang, Yu Zhang, Zhiyuan Li
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
# For a list of all contributors, visit:
# https://github.com/fla-org/flash-linear-attention/graphs/contributors
import os
import triton
import triton.language as tl
import triton.language.extra.libdevice as tldevice
from fla.utils import IS_GATHER_SUPPORTED, IS_NVIDIA_BLACKWELL
if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
@triton.jit
def exp(x): return tldevice.fast_expf(x.to(tl.float32))
@triton.jit
def exp2(x): return tldevice.exp2(x.to(tl.float32))
@triton.jit
def log(x): return tldevice.fast_logf(x.to(tl.float32))
@triton.jit
def log2(x): return tldevice.fast_log2f(x.to(tl.float32))
@triton.jit
def tanh(x): return tldevice.fast_tanhf(x.to(tl.float32))
else:
@triton.jit
def exp(x): return tl.exp(x.to(tl.float32))
@triton.jit
def exp2(x): return tl.math.exp2(x.to(tl.float32))
@triton.jit
def log(x): return tl.log(x.to(tl.float32))
@triton.jit
def log2(x): return tl.log2(x.to(tl.float32))
@triton.jit
def tanh(x): return tldevice.tanh(x.to(tl.float32))
if IS_NVIDIA_BLACKWELL:
"""
Compute tl.dot with Blackwell workaround.
On SM100 datacenter and SM120 consumer Blackwell GPUs, wraps the result in
inline assembly to prevent the TritonGPUHoistTMEMAlloc pass from incorrectly
fusing add and dot operations.
See: https://github.com/fla-org/flash-linear-attention/issues/638
TODO: Remove this workaround once the Triton compiler bug is fixed.
Track upstream issue at: https://github.com/triton-lang/triton/issues/8695
"""
@triton.jit
def safe_dot(a, b, allow_tf32: tl.constexpr = None):
return tl.inline_asm_elementwise(
asm="mov.f32 $0, $1;",
constraints="=r,r",
args=[tl.dot(a, b, allow_tf32=allow_tf32)],
dtype=tl.float32,
is_pure=True,
pack=1,
)
else:
@triton.jit
def safe_dot(a, b, allow_tf32: tl.constexpr = None):
return tl.dot(a, b, allow_tf32=allow_tf32)
if not IS_GATHER_SUPPORTED:
@triton.jit
def gather(src, index, axis, _builder=None):
"""
Gather operation that works when tl.gather is not supported.
This is a fallback implementation that returns None.
Just to make triton compiler happy.
"""
return None
else:
gather = tl.gather
if hasattr(triton.language, '_experimental_make_tensor_descriptor'):
# For Triton 3.3.x
make_tensor_descriptor = triton.language._experimental_make_tensor_descriptor
elif hasattr(triton.language, 'make_tensor_descriptor'):
# For Triton 3.4.x and later
make_tensor_descriptor = triton.language.make_tensor_descriptor
else:
"""
Fallback implementation when TMA is not supported.
Returns None to indicate TMA descriptors are unavailable.
Just make triton compiler happy.
"""
@triton.jit
def make_tensor_descriptor(
base,
shape,
strides,
block_shape,
_builder=None,
):
return None

188
ex_engine/xllm_kernels/cuda/activation.cu Normal file
View File

@@ -0,0 +1,188 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <c10/cuda/CUDAGuard.h>
#include <torch/cuda.h>
#include <cstdint>
#include "cuda_ops_api.h"
#include "device_utils.cuh"
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/activation_kernels.cu
namespace {
using ::xllm::kernel::cuda::xllm_ldg;
template <typename scalar_t,
scalar_t (*ACT_FN)(const scalar_t&),
bool act_first>
__device__ __forceinline__ scalar_t compute(const scalar_t& x,
const scalar_t& y) {
return act_first ? ACT_FN(x) * y : x * ACT_FN(y);
}
// Check if pointer is 16-byte aligned for int4 vectorized access
__device__ __forceinline__ bool is_16byte_aligned(const void* ptr) {
return (reinterpret_cast<uintptr_t>(ptr) & 15) == 0;
}
// Activation and gating kernel template with 128-bit vectorized access
// optimization.
template <typename scalar_t,
scalar_t (*ACT_FN)(const scalar_t&),
bool act_first>
__global__ void XLLM_KERNEL_ATTR(1024)
act_and_mul_kernel(scalar_t* __restrict__ out, // [..., d]
const scalar_t* __restrict__ input, // [..., 2, d]
const int d) {
constexpr int kVecSize = 16 / sizeof(scalar_t);
const int64_t token_idx = blockIdx.x;
const scalar_t* x_ptr = input + token_idx * 2 * d;
const scalar_t* y_ptr = x_ptr + d;
scalar_t* out_ptr = out + token_idx * d;
// Check alignment for 128-bit vectorized access.
// All three pointers must be 16-byte aligned for safe int4 operations.
const bool aligned = is_16byte_aligned(x_ptr) && is_16byte_aligned(y_ptr) &&
is_16byte_aligned(out_ptr);
if (aligned && d >= kVecSize) {
// Fast path: 128-bit vectorized loop
const int4* x_vec = reinterpret_cast<const int4*>(x_ptr);
const int4* y_vec = reinterpret_cast<const int4*>(y_ptr);
int4* out_vec = reinterpret_cast<int4*>(out_ptr);
const int num_vecs = d / kVecSize;
const int vec_end = num_vecs * kVecSize;
for (int i = threadIdx.x; i < num_vecs; i += blockDim.x) {
int4 x = xllm_ldg(&x_vec[i]), y = xllm_ldg(&y_vec[i]), r;
auto* xp = reinterpret_cast<scalar_t*>(&x);
auto* yp = reinterpret_cast<scalar_t*>(&y);
auto* rp = reinterpret_cast<scalar_t*>(&r);
#pragma unroll
for (int j = 0; j < kVecSize; j++) {
rp[j] = compute<scalar_t, ACT_FN, act_first>(xp[j], yp[j]);
}
out_vec[i] = r;
}
// Scalar cleanup for remaining elements
for (int i = vec_end + threadIdx.x; i < d; i += blockDim.x) {
out_ptr[i] = compute<scalar_t, ACT_FN, act_first>(xllm_ldg(&x_ptr[i]),
xllm_ldg(&y_ptr[i]));
}
} else {
// Scalar fallback for unaligned data or small d
for (int64_t idx = threadIdx.x; idx < d; idx += blockDim.x) {
const scalar_t x = xllm_ldg(&x_ptr[idx]);
const scalar_t y = xllm_ldg(&y_ptr[idx]);
out_ptr[idx] = compute<scalar_t, ACT_FN, act_first>(x, y);
}
}
}
template <typename T>
__device__ __forceinline__ T silu_kernel(const T& x) {
// x * sigmoid(x)
const float f = static_cast<float>(x);
return static_cast<T>(f / (1.0f + expf(-f)));
}
template <typename T>
__device__ __forceinline__ T gelu_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'none' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L36-L38
const float f = static_cast<float>(x);
constexpr float kAlpha = M_SQRT1_2;
return static_cast<T>(f * 0.5f * (1.0f + ::erf(f * kAlpha)));
}
template <typename T>
__device__ __forceinline__ T gelu_tanh_kernel(const T& x) {
// Equivalent to PyTorch GELU with 'tanh' approximation.
// Refer to:
// https://github.com/pytorch/pytorch/blob/8ac9b20d4b090c213799e81acf48a55ea8d437d6/aten/src/ATen/native/cuda/ActivationGeluKernel.cu#L25-L30
const float f = static_cast<float>(x);
constexpr float kBeta = M_SQRT2 * M_2_SQRTPI * 0.5f;
constexpr float kKappa = 0.044715;
float x_cube = f * f * f;
float inner = kBeta * (f + kKappa * x_cube);
return static_cast<T>(0.5f * f * (1.0f + ::tanhf(inner)));
}
#define LAUNCH_ACTIVATION_GATE_KERNEL(KERNEL, ACT_FIRST) \
int d = input.size(-1) / 2; \
int64_t num_tokens = input.numel() / input.size(-1); \
dim3 grid(num_tokens); \
dim3 block(std::min(d, 1024)); \
if (num_tokens == 0) { \
return; \
} \
const at::cuda::OptionalCUDAGuard device_guard(device_of(input)); \
const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); \
DISPATCH_FLOATING_TYPES(input.scalar_type(), "act_and_mul_kernel", [&] { \
act_and_mul_kernel<scalar_t, KERNEL<scalar_t>, ACT_FIRST> \
<<<grid, block, 0, stream>>>( \
out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), d); \
});
void silu_and_mul(torch::Tensor out, // [..., d]
torch::Tensor input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(silu_kernel, true);
}
void gelu_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(gelu_kernel, true);
}
void gelu_tanh_and_mul(torch::Tensor& out, // [..., d]
torch::Tensor& input) // [..., 2 * d]
{
LAUNCH_ACTIVATION_GATE_KERNEL(gelu_tanh_kernel, true);
}
} // namespace
namespace xllm::kernel::cuda {
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode) {
if (act_mode != "silu" && act_mode != "gelu" && act_mode != "gelu_tanh" &&
act_mode != "gelu_pytorch_tanh") {
LOG(FATAL) << "Unsupported act mode: " << act_mode
<< ", only support silu, gelu, gelu_tanh, gelu_pytorch_tanh";
}
// flashinfer act_and_mul ops
// std::string uri = act_mode + "_and_mul";
// FunctionFactory::get_instance().act_and_mul(uri).call(
// out, input, support_pdl());
if (act_mode == "silu") {
silu_and_mul(out, input);
} else if (act_mode == "gelu") {
gelu_and_mul(out, input);
} else if (act_mode == "gelu_tanh" || act_mode == "gelu_pytorch_tanh") {
// gelu_tanh or gelu_pytorch_tanh (mathematically equivalent)
gelu_tanh_and_mul(out, input);
}
}
} // namespace xllm::kernel::cuda

209
ex_engine/xllm_kernels/cuda/block_copy.cu Normal file
View File

@@ -0,0 +1,209 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAException.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAStream.h>
#include <cuda_runtime.h>
#include <torch/extension.h>
#include <cstdint>
#include <type_traits>
#include "cuda_ops_api.h"
#include "utils.h"
namespace xllm::kernel::cuda {
namespace {
template <typename scalar_t>
struct VecType;
template <>
struct VecType<c10::Half> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<c10::BFloat16> {
using type = uint4;
static constexpr int32_t vec_width = 8;
};
template <>
struct VecType<float> {
using type = float4;
static constexpr int32_t vec_width = 4;
};
DEVICE_INLINE int32_t find_group_idx(const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int32_t dst_idx) {
int32_t left = 0;
int32_t right = num_groups - 1;
while (left < right) {
const int32_t mid = left + ((right - left) >> 1);
const bool move_left = dst_idx < cum_sum[mid];
right = move_left ? mid : right;
left = move_left ? left : mid + 1;
}
return left;
}
template <typename scalar_t, bool kVectorized>
__global__ void block_copy_kernel(const int64_t* __restrict__ key_cache_ptrs,
const int64_t* __restrict__ value_cache_ptrs,
const int32_t* __restrict__ src_block_indices,
const int32_t* __restrict__ dst_block_indices,
const int32_t* __restrict__ cum_sum,
const int32_t num_groups,
const int64_t numel_per_block) {
const int64_t layer_idx = static_cast<int64_t>(blockIdx.x);
const int32_t dst_linear_idx = static_cast<int32_t>(blockIdx.y);
const int64_t tile_idx = static_cast<int64_t>(blockIdx.z);
scalar_t* __restrict__ key_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(key_cache_ptrs[layer_idx]));
scalar_t* __restrict__ value_cache = reinterpret_cast<scalar_t*>(
static_cast<uintptr_t>(value_cache_ptrs[layer_idx]));
const int32_t group_idx = find_group_idx(cum_sum, num_groups, dst_linear_idx);
const int32_t src_block = src_block_indices[group_idx];
const int32_t dst_block = dst_block_indices[dst_linear_idx];
const int64_t src_offset = static_cast<int64_t>(src_block) * numel_per_block;
const int64_t dst_offset = static_cast<int64_t>(dst_block) * numel_per_block;
if constexpr (kVectorized) {
using VecTypeT = typename VecType<scalar_t>::type;
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
const int64_t num_vecs_per_block = numel_per_block / kVecWidth;
const int64_t vec_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (vec_idx >= num_vecs_per_block) {
return;
}
const int64_t elem_offset = vec_idx * kVecWidth;
const auto* key_src_vec =
reinterpret_cast<const VecTypeT*>(key_cache + src_offset + elem_offset);
const auto* value_src_vec = reinterpret_cast<const VecTypeT*>(
value_cache + src_offset + elem_offset);
auto* key_dst_vec =
reinterpret_cast<VecTypeT*>(key_cache + dst_offset + elem_offset);
auto* value_dst_vec =
reinterpret_cast<VecTypeT*>(value_cache + dst_offset + elem_offset);
*key_dst_vec = *key_src_vec;
*value_dst_vec = *value_src_vec;
} else {
const int64_t elem_idx = tile_idx * static_cast<int64_t>(blockDim.x) +
static_cast<int64_t>(threadIdx.x);
if (elem_idx >= numel_per_block) {
return;
}
key_cache[dst_offset + elem_idx] = key_cache[src_offset + elem_idx];
value_cache[dst_offset + elem_idx] = value_cache[src_offset + elem_idx];
}
}
} // namespace
void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices,
torch::Tensor dst_block_indices,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype) {
if (src_block_indices.numel() == 0) {
return;
}
CHECK(key_cache_ptrs.is_cuda());
CHECK(value_cache_ptrs.is_cuda());
CHECK(src_block_indices.is_cuda());
CHECK(dst_block_indices.is_cuda());
CHECK(cum_sum.is_cuda());
CHECK_EQ(key_cache_ptrs.scalar_type(), torch::kInt64);
CHECK_EQ(value_cache_ptrs.scalar_type(), torch::kInt64);
CHECK_EQ(src_block_indices.scalar_type(), torch::kInt32);
CHECK_EQ(dst_block_indices.scalar_type(), torch::kInt32);
CHECK_EQ(cum_sum.scalar_type(), torch::kInt32);
CHECK_EQ(key_cache_ptrs.dim(), 1);
CHECK_EQ(value_cache_ptrs.dim(), 1);
CHECK_EQ(src_block_indices.dim(), 1);
CHECK_EQ(dst_block_indices.dim(), 1);
CHECK_EQ(cum_sum.dim(), 1);
CHECK(key_cache_ptrs.is_contiguous());
CHECK(value_cache_ptrs.is_contiguous());
CHECK(src_block_indices.is_contiguous());
CHECK(dst_block_indices.is_contiguous());
CHECK(cum_sum.is_contiguous());
CHECK_EQ(key_cache_ptrs.size(0), value_cache_ptrs.size(0));
CHECK_EQ(src_block_indices.size(0), cum_sum.size(0));
CHECK_GT(numel_per_block, 0);
const at::cuda::OptionalCUDAGuard device_guard(key_cache_ptrs.device());
constexpr int32_t kThreadsPerBlock = 256;
const int32_t num_layers = static_cast<int32_t>(key_cache_ptrs.size(0));
const int32_t num_groups = static_cast<int32_t>(src_block_indices.size(0));
const int32_t num_dst_blocks =
static_cast<int32_t>(dst_block_indices.size(0));
const cudaStream_t stream =
c10::cuda::getCurrentCUDAStream(key_cache_ptrs.get_device());
DISPATCH_FLOATING_TYPES(cache_dtype, "block_copy_kernel", [&] {
constexpr bool kHasVecType = std::is_same_v<scalar_t, float> ||
std::is_same_v<scalar_t, c10::Half> ||
std::is_same_v<scalar_t, c10::BFloat16>;
if constexpr (kHasVecType) {
constexpr int32_t kVecWidth = VecType<scalar_t>::vec_width;
if (numel_per_block % kVecWidth == 0) {
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block / kVecWidth, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, true>
<<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
return;
}
}
const int64_t tiles_per_block =
ceil_div<int64_t>(numel_per_block, kThreadsPerBlock);
const dim3 grid(num_layers, num_dst_blocks, tiles_per_block);
block_copy_kernel<scalar_t, false><<<grid, kThreadsPerBlock, 0, stream>>>(
key_cache_ptrs.data_ptr<int64_t>(),
value_cache_ptrs.data_ptr<int64_t>(),
src_block_indices.data_ptr<int32_t>(),
dst_block_indices.data_ptr<int32_t>(),
cum_sum.data_ptr<int32_t>(),
num_groups,
numel_per_block);
C10_CUDA_KERNEL_LAUNCH_CHECK();
});
}
} // namespace xllm::kernel::cuda

306
ex_engine/xllm_kernels/cuda/headers/cuda_ops_api.h Normal file
View File

@@ -0,0 +1,306 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <glog/logging.h>
#include <optional>
#include <tuple>
#include <vector>
#include "utils.h"
namespace xllm::kernel::cuda {
// TODO: add head_size parameter
void rotary_embedding(torch::Tensor& positions,
torch::Tensor& query,
std::optional<torch::Tensor> key,
torch::Tensor& cos_sin_cache,
// int64_t head_size,
bool is_neox);
// act_mode only support silu, gelu, gelu_tanh
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode);
void reshape_paged_cache(
torch::Tensor slot_ids, // [n_tokens]
torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor values, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim]
torch::Tensor value_cache);
void block_copy(torch::Tensor key_cache_ptrs,
torch::Tensor value_cache_ptrs,
torch::Tensor src_block_indices,
torch::Tensor dst_block_indices,
torch::Tensor cum_sum,
int64_t numel_per_block,
torch::ScalarType cache_dtype);
#if !defined(USE_DCU)
void batch_prefill(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& mask = std::nullopt);
// Wrapper function for batch_prefill that conditionally uses AttentionRunner
// for piecewise CUDA Graph capture
void batch_prefill_with_optional_piecewise_capture(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse);
void batch_prefill_non_causal(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor key,
torch::Tensor value,
torch::Tensor q_cu_seq_lens,
torch::Tensor kv_cu_seq_lens,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& mask = std::nullopt);
void batch_chunked_prefill(
const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor paged_kv_indptr,
torch::Tensor paged_kv_indices,
torch::Tensor paged_kv_last_page_len,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
std::optional<torch::Tensor> qo_indptr = std::nullopt,
bool causal = true);
void batch_decode(const std::string& uri,
ffi::Array<int64_t> plan_info,
torch::Tensor float_workspace_buffer,
torch::Tensor int_workspace_buffer,
torch::Tensor page_locked_int_workspace_buffer,
torch::Tensor query,
torch::Tensor k_cache,
torch::Tensor v_cache,
torch::Tensor paged_kv_indptr,
torch::Tensor paged_kv_indices,
torch::Tensor paged_kv_last_page_len,
int64_t window_left,
double sm_scale,
torch::Tensor output,
std::optional<torch::Tensor>& output_lse,
bool use_tensor_core,
std::optional<torch::Tensor> qo_indptr = std::nullopt);
#endif // !defined(USE_DCU)
void rms_norm(torch::Tensor output,
torch::Tensor input,
torch::Tensor weight,
double eps);
void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
torch::Tensor& residual, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
double epsilon);
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias);
void cutlass_scaled_mm(torch::Tensor& c,
torch::Tensor const& a,
torch::Tensor const& b,
torch::Tensor const& a_scales,
torch::Tensor const& b_scales,
std::optional<torch::Tensor> const& bias);
// Static scaled FP8 quantization
// Quantizes input tensor to FP8 using a pre-computed scale factor
void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
torch::Tensor const& input, // [..., d]
torch::Tensor const& scale); // [1]
// FP8 scaled quantize: quantizes input tensor to FP8 e4m3 format
// Returns: (quantized_output, scale)
std::tuple<torch::Tensor, torch::Tensor> fp8_scaled_quantize(
const torch::Tensor& input,
const std::optional<torch::Tensor>& output = std::nullopt,
const std::optional<torch::Tensor>& scale = std::nullopt);
// ============================================================================
// Fused RMSNorm + Static FP8 Quantization
// ============================================================================
// These functions combine RMSNorm and FP8 quantization to reduce memory
// bandwidth by avoiding the intermediate write-back to global memory.
// Fused RMSNorm + Static FP8 Quantization (without residual)
// Combines RMSNorm normalization and FP8 quantization in a single kernel.
// This is optimal for the first layer where no residual connection exists.
void rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size], FP8 output
torch::Tensor& input, // [..., hidden_size], input tensor
torch::Tensor& weight, // [hidden_size], RMSNorm weight
torch::Tensor& scale, // [1], FP8 quantization scale
double epsilon); // RMSNorm epsilon
// Fused Add + RMSNorm + Static FP8 Quantization (with residual)
// Combines residual addition, RMSNorm, and FP8 quantization in a single kernel.
// The residual tensor is updated in-place with the sum of input and residual.
void fused_add_rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size], FP8 output
torch::Tensor& input, // [..., hidden_size], input tensor
torch::Tensor& residual, // [..., hidden_size], residual (updated in-place)
torch::Tensor& weight, // [hidden_size], RMSNorm weight
torch::Tensor& scale, // [1], FP8 quantization scale
double epsilon); // RMSNorm epsilon
// FP8 scaled matmul for W8A8 quantization using CUTLASS kernels
// Performs: c = (a @ b.T) with scales applied
torch::Tensor fp8_scaled_matmul(
const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& a_scale,
const torch::Tensor& b_scale,
torch::ScalarType output_dtype,
const std::optional<torch::Tensor>& bias = std::nullopt,
const std::optional<torch::Tensor>& output = std::nullopt);
std::pair<torch::Tensor, torch::Tensor> compute_topk_for_beam_search(
torch::Tensor combined_probs,
uint32_t batch_size,
uint32_t beam_size,
uint32_t top_k,
torch::Device device);
std::pair<torch::Tensor, torch::Tensor> compute_topk_general(
torch::Tensor input,
uint32_t batch_size,
uint32_t input_length,
uint32_t k,
torch::Device device);
torch::Tensor air_log_softmax_last_dim(const torch::Tensor& input,
const torch::Tensor& temperatures);
void fused_qk_norm_rope(
torch::Tensor& qkv, // Combined QKV tensor [num_tokens,
// (num_heads_q+num_heads_k+num_heads_v)*head_dim]
int64_t num_heads_q, // Number of query heads
int64_t num_heads_k, // Number of key heads
int64_t num_heads_v, // Number of value heads
int64_t head_dim, // Dimension per head
double eps, // Epsilon for RMS normalization
const torch::Tensor& q_weight, // RMSNorm weights for query [head_dim]
const torch::Tensor& k_weight, // RMSNorm weights for key [head_dim]
const torch::Tensor&
cos_sin_cache, // Cos/sin cache [max_position, rotary_dim]
bool interleaved, // Whether RoPE is applied in interleaved style
const torch::Tensor& position_ids // Position IDs for RoPE [num_tokens]
);
std::tuple<torch::Tensor, torch::Tensor> moe_fused_topk(
torch::Tensor& gating_output,
int64_t topk,
bool renormalize,
const std::optional<torch::Tensor>& correction_bias,
const std::string& scoring_func);
torch::Tensor random_sample(const torch::Tensor& probs);
torch::Tensor cutlass_fused_moe(
const torch::Tensor& input, // [num_tokens, hidden]
const torch::Tensor& token_selected_experts, // [num_tokens, top_k]
const torch::Tensor& token_final_scales, // [num_tokens, top_k]
const torch::Tensor&
fc1_expert_weights, // [num_experts, inter_dim, hidden]
const torch::Tensor&
fc2_expert_weights, // [num_experts, hidden, inter_dim]
torch::ScalarType output_dtype,
const std::vector<torch::Tensor>& quant_scales,
int32_t tp_size,
int32_t tp_rank,
int32_t ep_size,
int32_t ep_rank,
int32_t cluster_size,
int32_t cluster_rank,
const std::optional<torch::Tensor>& fc1_expert_biases = std::nullopt,
const std::optional<torch::Tensor>& fc2_expert_biases = std::nullopt,
const std::optional<torch::Tensor>& input_sf = std::nullopt,
const std::optional<torch::Tensor>& swiglu_alpha = std::nullopt,
const std::optional<torch::Tensor>& swiglu_beta = std::nullopt,
const std::optional<torch::Tensor>& swiglu_limit = std::nullopt,
const std::optional<torch::Tensor>& output = std::nullopt,
bool enable_alltoall = false,
bool use_deepseek_fp8_block_scale = false,
bool use_w4_group_scaling = false,
bool use_mxfp8_act_scaling = false,
bool min_latency_mode = false,
bool use_packed_weights = false,
int32_t tune_max_num_tokens = 8192,
ActivationType activation_type = ActivationType::SWIGLU);
// ---- moe_compute_index (moe_compute_index.cu) ----
// Fused routing index: bincount + argsort replacement.
// Returns {src_dst, dst_src, expert_sizes}.
std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> moe_compute_index(
const torch::Tensor& expert_id,
int64_t num_experts);
// ---- moe_combine_result (moe_combine.cu) ----
// Fused combine: reorder + weighted sum in one pass.
torch::Tensor moe_combine_result(const torch::Tensor& gemm2,
const torch::Tensor& reduce_weight,
int64_t N,
int32_t topk);
} // namespace xllm::kernel::cuda

116
ex_engine/xllm_kernels/cuda/headers/device_utils.cuh Normal file
View File

@@ -0,0 +1,116 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#if defined(USE_DCU)
#include <hip/amd_detail/amd_hip_bf16.h>
#include <hipcub/hipcub.hpp>
namespace cub = hipcub;
#else
#include <cub/cub.cuh>
#if CUB_VERSION >= 200800
#include <cuda/functional>
#endif
#endif
namespace xllm::kernel::cuda {
#if !defined(USE_DCU)
using BFloat16Type = __nv_bfloat16;
#define WARP_SIZE 32
#define XLLM_KERNEL_ATTR(MAX_THREADS)
#else
using BFloat16Type = hip_bfloat16;
#define WARP_SIZE 64
#define XLLM_KERNEL_ATTR(MAX_THREADS) __launch_bounds__(MAX_THREADS, 1)
#endif
#define MAX(a, b) ((a) > (b) ? (a) : (b))
#define MIN(a, b) ((a) < (b) ? (a) : (b))
// Aligned array type
template <typename T,
// Number of elements in the array
int N,
// Alignment requirement in bytes
int Alignment = sizeof(T) * N>
class alignas(Alignment) AlignedArray {
T data[N];
};
#define XLLM_SHFL_XOR_SYNC(mask, var, lane_mask) \
__shfl_xor_sync((mask), (var), (lane_mask))
#define XLLM_SHFL_XOR_SYNC_WIDTH(mask, var, lane_mask, width) \
__shfl_xor_sync((mask), (var), (lane_mask), (width))
template <typename T>
__device__ __forceinline__ T xllm_ldg(const T* ptr) {
#if defined(USE_DCU)
return *ptr;
#else
return __ldg(ptr);
#endif
}
// Define reduction operators based on CUB version.
#if defined(USE_DCU)
using MaxReduceOp = hipcub::Max;
using MinReduceOp = hipcub::Min;
#elif CUB_VERSION >= 200800
using MaxReduceOp = ::cuda::maximum<>;
using MinReduceOp = ::cuda::minimum<>;
#else
using MaxReduceOp = cub::Max;
using MinReduceOp = cub::Min;
#endif
template <typename T>
__device__ float convert_to_float(T x) {
if constexpr (std::is_same_v<T, __half>) {
return __half2float(x);
#if defined(USE_DCU)
} else if constexpr (std::is_same_v<T, hip_bfloat16>) {
return __bfloat162float(reinterpret_cast<const __hip_bfloat16&>(x));
#else
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
return __bfloat162float(x);
#endif
} else if constexpr (std::is_same_v<T, float>) {
return x;
} else {
return static_cast<float>(x);
}
}
// Constructs some constants needed to partition the work across threads at
// compile time.
template <typename T, int EXPERTS, int BYTES_PER_LDG>
struct TopkConstants {
static constexpr int ELTS_PER_LDG = BYTES_PER_LDG / sizeof(T);
static_assert(EXPERTS / (ELTS_PER_LDG * WARP_SIZE) == 0 ||
EXPERTS % (ELTS_PER_LDG * WARP_SIZE) == 0,
"");
static constexpr int VECs_PER_THREAD =
MAX(1, EXPERTS / (ELTS_PER_LDG * WARP_SIZE));
static constexpr int VPT = VECs_PER_THREAD * ELTS_PER_LDG;
static constexpr int THREADS_PER_ROW = EXPERTS / VPT;
static constexpr int ROWS_PER_WARP = WARP_SIZE / THREADS_PER_ROW;
};
} // namespace xllm::kernel::cuda

239
ex_engine/xllm_kernels/cuda/headers/fp8_quant_utils.cuh Normal file
View File

@@ -0,0 +1,239 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* https://github.com/jd-opensource/xllm/blob/main/LICENSE
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
* ===========================================================================*/
#pragma once
// clang-format off
#include <c10/util/Float8_e4m3fn.h>
#include <cmath>
#include <torch/types.h>
// clang-format on
namespace xllm {
namespace kernel {
namespace cuda {
// FP8 type max value definitions
template <typename T,
typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
std::is_same_v<T, int8_t>>>
struct quant_type_max {
static constexpr T val() { return std::numeric_limits<T>::max(); }
};
template <typename T>
__host__ __device__ static constexpr T quant_type_max_v =
quant_type_max<T>::val();
// Minimum scaling factor for quantization types
template <typename T,
typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
std::is_same_v<T, int8_t>>>
struct min_scaling_factor {
__device__ __host__ static inline float val() {
return 1.0f / (quant_type_max_v<T> * 512.0f);
}
};
template <>
struct min_scaling_factor<int8_t> {
__device__ __host__ static inline float val() {
return std::numeric_limits<float>::epsilon();
}
};
// Vectorization containers
template <typename scalar_t, size_t vec_size>
struct __align__(vec_size * sizeof(scalar_t)) vec_n_t {
scalar_t val[vec_size];
};
template <typename quant_type_t, size_t vec_size>
struct __align__(vec_size * sizeof(quant_type_t)) q8_n_t {
static_assert(std::is_same_v<quant_type_t, int8_t> ||
std::is_same_v<quant_type_t, c10::Float8_e4m3fn>);
quant_type_t val[vec_size];
};
// Atomic max for float
__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
float old;
old = (value >= 0)
? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
: __uint_as_float(
atomicMin((unsigned int*)addr, __float_as_uint(value)));
return old;
}
// FP8 conversion functions
namespace fp8 {
#ifdef ENABLE_FP8
#include <cuda_fp8.h>
// float -> c10::Float8_e4m3fn conversion
template <typename Tout, typename Tin>
__inline__ __device__ Tout
vec_conversion(const Tin& x,
const __nv_fp8_interpretation_t fp8_type = __NV_E4M3) {
return x;
}
template <>
__inline__ __device__ c10::Float8_e4m3fn
vec_conversion<c10::Float8_e4m3fn, float>(
const float& a,
const __nv_fp8_interpretation_t fp8_type) {
#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
return static_cast<c10::Float8_e4m3fn>(a);
#else
return c10::Float8_e4m3fn(__nv_cvt_float_to_fp8(a, __NV_SATFINITE, fp8_type),
c10::Float8_e4m3fn::from_bits());
#endif
}
#endif // ENABLE_FP8
} // namespace fp8
// Scaled FP8 conversion with saturation
template <bool is_scale_inverted, typename fp8_type>
__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
float const scale) {
float x = 0.0f;
if constexpr (is_scale_inverted) {
x = val * scale;
} else {
x = val / scale;
}
float r =
fmaxf(-quant_type_max_v<fp8_type>, fminf(x, quant_type_max_v<fp8_type>));
#ifdef ENABLE_FP8
// Use hardware cvt instruction for fp8 on nvidia
return fp8::vec_conversion<fp8_type, float>(r);
#else
return static_cast<fp8_type>(r);
#endif
}
// Vectorization utilities
template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
struct DefaultVecOp {
ScaOp scalar_op;
__device__ __forceinline__ void operator()(
vec_n_t<OutT, VEC_SIZE>& dst,
const vec_n_t<InT, VEC_SIZE>& src) const {
#pragma unroll
for (int i = 0; i < VEC_SIZE; ++i) {
scalar_op(dst.val[i], src.val[i]);
}
}
};
template <int VEC_SIZE,
typename InT,
typename OutT,
typename VecOp,
typename ScaOp>
__device__ inline void vectorize_with_alignment(
const InT* in,
OutT* out,
int len,
int tid,
int stride,
VecOp&& vec_op, // vec_n_t<InT,16> -> vec_n_t<OutT,16>
ScaOp&& scalar_op) { // InT -> OutT
static_assert(VEC_SIZE > 0 && (VEC_SIZE & (VEC_SIZE - 1)) == 0,
"VEC_SIZE must be a positive power-of-two");
constexpr int WIDTH = VEC_SIZE * sizeof(InT);
uintptr_t addr = reinterpret_cast<uintptr_t>(in);
// Fast path when the whole region is already aligned
bool can_vec = ((addr & (WIDTH - 1)) == 0) && ((len & (VEC_SIZE - 1)) == 0);
if (can_vec) {
int num_vec = len / VEC_SIZE;
using vin_t = vec_n_t<InT, VEC_SIZE>;
using vout_t = vec_n_t<OutT, VEC_SIZE>;
auto* v_in = reinterpret_cast<const vin_t*>(in);
auto* v_out = reinterpret_cast<vout_t*>(out);
for (int i = tid; i < num_vec; i += stride) {
vout_t tmp;
vin_t src = v_in[i];
vec_op(tmp, src);
v_out[i] = tmp;
}
return;
}
int misalignment_offset = addr & (WIDTH - 1);
int alignment_bytes = WIDTH - misalignment_offset;
int prefix_elems = alignment_bytes & (WIDTH - 1);
prefix_elems /= sizeof(InT);
prefix_elems = min(prefix_elems, len);
// Prefix handling
for (int i = tid; i < prefix_elems; i += stride) {
scalar_op(out[i], in[i]);
}
in += prefix_elems;
out += prefix_elems;
len -= prefix_elems;
int num_vec = len / VEC_SIZE;
using vin_t = vec_n_t<InT, VEC_SIZE>;
using vout_t = vec_n_t<OutT, VEC_SIZE>;
auto* v_in = reinterpret_cast<const vin_t*>(in);
auto* v_out = reinterpret_cast<vout_t*>(out);
// Vectorized main part
for (int i = tid; i < num_vec; i += stride) {
vout_t tmp;
vin_t src = v_in[i];
vec_op(tmp, src);
v_out[i] = tmp;
}
// Tail handling
int tail_start = num_vec * VEC_SIZE;
for (int i = tid + tail_start; i < len; i += stride) {
scalar_op(out[i], in[i]);
}
}
template <int VEC_SIZE, typename InT, typename OutT, typename ScaOp>
__device__ __forceinline__ void vectorize_with_alignment(const InT* in,
OutT* out,
int len,
int tid,
int stride,
ScaOp&& scalar_op) {
using Vec = DefaultVecOp<VEC_SIZE, InT, OutT, std::decay_t<ScaOp>>;
vectorize_with_alignment<VEC_SIZE>(in,
out,
len,
tid,
stride,
Vec{scalar_op},
std::forward<ScaOp>(scalar_op));
}
} // namespace cuda
} // namespace kernel
} // namespace xllm

231
ex_engine/xllm_kernels/cuda/headers/type_convert.cuh Normal file
View File

@@ -0,0 +1,231 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <torch/all.h>
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/type_convert.cuh
/* Converter helpers for the conversion from torch types to HIP/CUDA types,
and the associated type conversions within HIP/CUDA. These helpers need
to be implemented for now because the relevant type conversion
operators/constructors are not consistently implemented by HIP/CUDA, so
a generic conversion via type casts cannot be implemented.
Each helper should have the member static constexpr bool `exists`:
If false, the optimized kernel is not used for the corresponding torch type.
If true, the helper should be fully defined as shown in the examples below.
*/
namespace xllm::kernel::cuda {
template <typename torch_type>
class _typeConvert {
public:
static constexpr bool exists = false;
};
template <>
class _typeConvert<float> {
public:
static constexpr bool exists = true;
using hip_type = float;
using packed_hip_type = float2;
using packed_hip_type4 = float4; // For 128-bit vectorization
__device__ static __forceinline__ float convert(hip_type x) { return x; }
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return x;
}
__device__ static __forceinline__ float4 convert(packed_hip_type4 x) {
return x;
}
};
#if defined(USE_DCU) || (defined(CUDA_VERSION) && (CUDA_VERSION >= 12000)) || \
defined(USE_MACA)
// CUDA < 12.0 runs into issues with packed type conversion
template <>
class _typeConvert<c10::Half> {
public:
static constexpr bool exists = true;
using hip_type = __half;
using packed_hip_type = __half2;
__device__ static __forceinline__ float convert(hip_type x) {
return __half2float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __half22float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2half_rn(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22half2_rn(x);
}
};
#endif // defined(USE_DCU) || CUDA_VERSION >= 12000
#if defined(USE_DCU)
template <>
class _typeConvert<c10::BFloat16> {
public:
static constexpr bool exists = true;
using hip_type = __hip_bfloat16;
using packed_hip_type = __hip_bfloat162;
__device__ static __forceinline__ float convert(hip_type x) {
return __bfloat162float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __bfloat1622float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2bfloat16(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22bfloat162_rn(x);
}
};
#elif defined(CUDA_VERSION) && (CUDA_VERSION >= 12000) && \
defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) || \
defined(USE_MACA)
// CUDA_ARCH < 800 does not have BF16 support.
template <>
class _typeConvert<c10::BFloat16> {
public:
static constexpr bool exists = true;
using hip_type = __nv_bfloat16;
using packed_hip_type = __nv_bfloat162;
__device__ static __forceinline__ float convert(hip_type x) {
return __bfloat162float(x);
}
__device__ static __forceinline__ float2 convert(packed_hip_type x) {
return __bfloat1622float2(x);
}
__device__ static __forceinline__ hip_type convert(float x) {
return __float2bfloat16(x);
}
__device__ static __forceinline__ packed_hip_type convert(float2 x) {
return __float22bfloat162_rn(x);
}
};
#endif
/* Vector helper to generate vectorized and packed FP16/BF16 ops
for appropriate specializations of fused_add_rms_norm_kernel.
Only functions that are necessary in that kernel are implemented.
Alignment to 16 bytes is required to use 128-bit global memory ops.
*/
template <typename scalar_t, int width>
class alignas(16) _f16Vec {
public:
/* Not theoretically necessary that width is a power of 2 but should
almost always be the case for optimization purposes */
static_assert(width > 0 && (width & (width - 1)) == 0,
"Width is not a positive power of 2!");
using Converter = _typeConvert<scalar_t>;
using T1 = typename Converter::hip_type;
using T2 = typename Converter::packed_hip_type;
T1 data[width];
__device__ _f16Vec& operator+=(const _f16Vec<scalar_t, width>& other) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
if constexpr (std::is_same_v<T2, float2>) {
data[i] += other.data[i];
data[i + 1] += other.data[i + 1];
} else {
T2 temp{data[i], data[i + 1]};
temp += T2{other.data[i], other.data[i + 1]};
data[i] = temp.x;
data[i + 1] = temp.y;
}
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) data[i] += other.data[i];
}
return *this;
}
__device__ _f16Vec& operator*=(const _f16Vec<scalar_t, width>& other) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
if constexpr (std::is_same_v<T2, float2>) {
data[i] *= other.data[i];
data[i + 1] *= other.data[i + 1];
} else {
T2 temp{data[i], data[i + 1]};
temp *= T2{other.data[i], other.data[i + 1]};
data[i] = temp.x;
data[i + 1] = temp.y;
}
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) data[i] *= other.data[i];
}
return *this;
}
__device__ _f16Vec& operator*=(const float scale) {
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
float2 temp_f = Converter::convert(T2{data[i], data[i + 1]});
temp_f.x *= scale;
temp_f.y *= scale;
T2 temp = Converter::convert(temp_f);
data[i] = temp.x;
data[i + 1] = temp.y;
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) {
float temp = Converter::convert(data[i]) * scale;
data[i] = Converter::convert(temp);
}
}
return *this;
}
__device__ float sum_squares() const {
float result = 0.0f;
if constexpr (width % 2 == 0) {
#pragma unroll
for (int i = 0; i < width; i += 2) {
float2 z = Converter::convert(T2{data[i], data[i + 1]});
result += z.x * z.x + z.y * z.y;
}
} else {
#pragma unroll
for (int i = 0; i < width; ++i) {
float x = Converter::convert(data[i]);
result += x * x;
}
}
return result;
}
};
} // namespace xllm::kernel::cuda

163
ex_engine/xllm_kernels/cuda/headers/utils.h Normal file
View File

@@ -0,0 +1,163 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#if defined(USE_DCU)
#include <c10/hip/HIPGuard.h>
#else
#include <c10/cuda/CUDAGuard.h>
#endif
#include <glog/logging.h>
#include <torch/torch.h>
#if !defined(USE_DCU)
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/extra/c_env_api.h>
#include <tvm/ffi/extra/module.h>
#include <tvm/ffi/optional.h>
#endif
#include <string>
#include <tuple>
#include <type_traits>
#include <unordered_map>
#if defined(__CUDACC__) || defined(_NVHPC_CUDA) || defined(__HIPCC__)
#define HOST_DEVICE_INLINE __host__ __device__ __forceinline__
#define DEVICE_INLINE __device__ __forceinline__
#define HOST_INLINE __host__ __forceinline__
#else
#define HOST_DEVICE_INLINE inline
#define DEVICE_INLINE inline
#define HOST_INLINE inline
#endif
#if !defined(USE_DCU)
namespace ffi = tvm::ffi;
#endif
namespace xllm::kernel::cuda {
template <typename T>
HOST_DEVICE_INLINE constexpr std::enable_if_t<std::is_integral_v<T>, T>
ceil_div(T a, T b) {
return (a + b - 1) / b;
}
enum class ActivationType : int8_t {
GELU = 0,
RELU = 1,
SILU = 2,
SWIGLU = 3,
GEGLU = 4,
SWIGLU_BIAS = 5,
RELU2 = 6,
IDENTITY = 7,
INVALID_TYPE = 8
};
// torch tensor is only on cpu
torch::Tensor get_cache_buffer(const int32_t seq_len,
const torch::Device& device);
// NOLINTBEGIN(cppcoreguidelines-macro-usage)
#define DISPATCH_CASE_FLOATING_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
#define DISPATCH_CASE_HALF_TYPES(...) \
AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
#define DISPATCH_HALF_TYPES(TYPE, NAME, ...) \
AT_DISPATCH_SWITCH(TYPE, NAME, DISPATCH_CASE_HALF_TYPES(__VA_ARGS__))
// NOLINTEND(cppcoreguidelines-macro-usage)
bool should_use_tensor_core(torch::ScalarType kv_cache_dtype,
int64_t num_attention_heads,
int64_t num_kv_heads);
bool support_pdl();
std::string path_to_uri_so_lib(const std::string& uri);
std::string determine_attention_backend(int64_t pos_encoding_mode,
bool use_fp16_qk_reduction,
bool use_custom_mask);
std::string get_batch_prefill_uri(const std::string& backend,
torch::ScalarType dtype_q,
torch::ScalarType dtype_kv,
torch::ScalarType dtype_o,
torch::ScalarType dtype_idx,
int64_t head_dim_qk,
int64_t head_dim_vo,
int64_t pos_encoding_mode,
bool use_sliding_window,
bool use_logits_soft_cap,
bool use_fp16_qk_reduction);
std::string get_batch_decode_uri(torch::ScalarType dtype_q,
torch::ScalarType dtype_kv,
torch::ScalarType dtype_o,
torch::ScalarType dtype_idx,
int64_t head_dim_qk,
int64_t head_dim_vo,
int64_t pos_encoding_mode,
bool use_sliding_window,
bool use_logits_soft_cap);
std::tuple<torch::Tensor, double> split_scale_param(const torch::Tensor& scale);
#if !defined(USE_DCU)
DLDataType to_dl_data_type(torch::ScalarType scalar_type);
// below are tvm-ffi related functions
ffi::Tensor to_ffi_tensor(const torch::Tensor& torch_tensor);
ffi::Optional<ffi::Tensor> to_ffi_optional_tensor(
const std::optional<torch::Tensor>& optional);
ffi::Array<ffi::Tensor> to_ffi_array_tensors(
const std::vector<torch::Tensor>& torch_tensors);
ffi::Optional<ffi::Array<ffi::Tensor>> to_ffi_optional_array_tensors(
const std::optional<std::vector<torch::Tensor>>& optional);
ffi::Module get_module(const std::string& uri);
ffi::Function get_function(const std::string& uri,
const std::string& func_name);
inline void bind_tvmffi_stream_to_current_torch_stream(
const torch::Device& device) {
const auto cur = c10::cuda::getCurrentCUDAStream(device.index());
// DLPack device type for CUDA is 2 (kDLCUDA).
void* original_stream = nullptr;
const int rc = TVMFFIEnvSetStream(
/*device_type=*/2,
/*device_id=*/device.index(),
reinterpret_cast<void*>(cur.stream()),
&original_stream);
if (rc != 0) {
LOG(WARNING) << "[tvmffi.stream] failed to set stream, rc=" << rc
<< " dev=" << device.index();
}
}
#endif // !defined(USE_DCU)
} // namespace xllm::kernel::cuda

600
ex_engine/xllm_kernels/cuda/norm.cu Normal file
View File

@@ -0,0 +1,600 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <c10/cuda/CUDAGuard.h>
#include <torch/cuda.h>
#include <cstdint>
#include <cub/cub.cuh>
#include "cuda_ops_api.h"
#include "device_utils.cuh"
#include "fp8_quant_utils.cuh"
#include "type_convert.cuh"
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/layernorm_kernels.cu
#if CUB_VERSION >= 200800
#include <cuda/std/functional>
using CubAddOp = ::cuda::std::plus<>;
using CubMaxOp = ::cuda::maximum<>;
#else // if CUB_VERSION < 200800
using CubAddOp = cub::Sum;
using CubMaxOp = cub::Max;
#endif // CUB_VERSION
namespace {
using namespace xllm::kernel::cuda;
template <typename scalar_t>
__global__ void XLLM_KERNEL_ATTR(1024)
rms_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size]
const scalar_t* __restrict__ input, // [..., hidden_size]
const int64_t input_stride,
const scalar_t* __restrict__ weight, // [hidden_size]
const float epsilon,
const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
const float x = static_cast<float>(input[blockIdx.x * input_stride + idx]);
variance += x * x;
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = static_cast<float>(input[blockIdx.x * input_stride + idx]);
out[blockIdx.x * hidden_size + idx] =
(static_cast<scalar_t>(x * s_variance)) * weight[idx];
}
}
/* Function specialization in the case of FP16/BF16 tensors.
Additional optimizations we can make in this case are
packed and vectorized operations, which help with the
memory latency bottleneck. */
template <typename scalar_t, int width>
__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel(
scalar_t* __restrict__ input, // [..., hidden_size]
const int64_t input_stride,
scalar_t* __restrict__ residual, // [..., hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
const float epsilon,
const int num_tokens,
const int hidden_size) {
// Sanity checks on our vector struct and type-punned pointer arithmetic
static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
const int vec_hidden_size = hidden_size / width;
const int64_t vec_input_stride = input_stride / width;
__shared__ float s_variance;
float variance = 0.0f;
/* These and the argument pointers are all declared `restrict` as they are
not aliased in practice. Argument pointers should not be dereferenced
in this kernel as that would be undefined behavior */
auto* __restrict__ input_v =
reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
auto* __restrict__ residual_v =
reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
auto* __restrict__ weight_v =
reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
_f16Vec<scalar_t, width> temp = input_v[strided_id];
temp += residual_v[id];
variance += temp.sum_squares();
residual_v[id] = temp;
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
input_v[strided_id] = temp;
}
}
/* Generic fused_add_rms_norm_kernel
The width field is not used here but necessary for other specializations.
*/
template <typename scalar_t, int width>
__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
XLLM_KERNEL_ATTR(1024) fused_add_rms_norm_kernel(
scalar_t* __restrict__ input, // [..., hidden_size]
const int64_t input_stride,
scalar_t* __restrict__ residual, // [..., hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
const float epsilon,
const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
scalar_t z = input[blockIdx.x * input_stride + idx];
z += residual[blockIdx.x * hidden_size + idx];
float x = static_cast<float>(z);
variance += x * x;
residual[blockIdx.x * hidden_size + idx] = z;
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = static_cast<float>(residual[blockIdx.x * hidden_size + idx]);
input[blockIdx.x * input_stride + idx] =
(static_cast<scalar_t>(x * s_variance)) * weight[idx];
}
}
#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
fused_add_rms_norm_kernel<scalar_t, width> \
<<<grid, block, 0, stream>>>(input.data_ptr<scalar_t>(), \
input_stride, \
residual.data_ptr<scalar_t>(), \
weight.data_ptr<scalar_t>(), \
epsilon, \
num_tokens, \
hidden_size); \
});
// ============================================================================
// Fused RMSNorm + Static FP8 Quantization Kernels
// ============================================================================
// These kernels combine RMSNorm and FP8 quantization to reduce memory
// bandwidth by avoiding the intermediate write-back to global memory.
// Dispatch macro for FP8 types
#define DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
[&] { \
const auto& the_type = TYPE; \
switch (the_type) { \
case at::ScalarType::Float8_e4m3fn: { \
using fp8_t = c10::Float8_e4m3fn; \
return __VA_ARGS__(); \
} \
default: \
AT_ERROR(#NAME, \
" not implemented for FP8 type '", \
toString(the_type), \
"'"); \
} \
}()
/**
* Fused RMSNorm + Static FP8 Quantization kernel (without residual)
* Combines RMSNorm and FP8 quantization in a single kernel to reduce
* memory bandwidth by avoiding intermediate write-back.
*
* @tparam scalar_t Input data type (float, half, bfloat16)
* @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn)
* @param out Output FP8 tensor [num_tokens, hidden_size]
* @param input Input tensor [num_tokens, hidden_size]
* @param input_stride Stride of input tensor in the token dimension
* @param weight RMSNorm weight tensor [hidden_size]
* @param scale FP8 quantization scale (scalar)
* @param epsilon RMSNorm epsilon
* @param num_tokens Number of tokens
* @param hidden_size Hidden dimension size
*/
template <typename scalar_t, typename fp8_type>
__global__ void rms_norm_static_fp8_quant_kernel(
fp8_type* __restrict__ out, // [num_tokens, hidden_size]
const scalar_t* __restrict__ input, // [num_tokens, hidden_size]
const int64_t input_stride,
const scalar_t* __restrict__ weight, // [hidden_size]
const float* __restrict__ scale, // [1]
const float epsilon,
const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
const scalar_t* input_row = input + blockIdx.x * input_stride;
// Step 1: Compute variance for RMSNorm
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
const float x = static_cast<float>(input_row[idx]);
variance += x * x;
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
// Step 2: Precompute scale inverse to avoid division
const float scale_inv = 1.0f / (*scale);
// Step 3: Fused RMSNorm + FP8 quantization
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = static_cast<float>(input_row[idx]);
float out_norm = (static_cast<scalar_t>(x * s_variance)) *
static_cast<float>(weight[idx]);
out[blockIdx.x * hidden_size + idx] =
xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(out_norm,
scale_inv);
}
}
/**
* Fused Add + RMSNorm + Static FP8 Quantization kernel (with residual)
* Optimized version with packed + vectorized operations for FP16/BF16.
*
* @tparam scalar_t Input data type (float, half, bfloat16)
* @tparam width Vector width for optimization (0, 8)
* @tparam fp8_type Output FP8 type (c10::Float8_e4m3fn)
*/
template <typename scalar_t, int width, typename fp8_type>
__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
fused_add_rms_norm_static_fp8_quant_kernel(
fp8_type* __restrict__ out, // [num_tokens, hidden_size]
scalar_t* __restrict__ input, // [num_tokens, hidden_size]
const int64_t input_stride,
scalar_t* __restrict__ residual, // [num_tokens, hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
const float* __restrict__ scale, // [1]
const float epsilon,
const int num_tokens,
const int hidden_size) {
static_assert(std::is_pod_v<_f16Vec<scalar_t, width>>);
static_assert(sizeof(_f16Vec<scalar_t, width>) == sizeof(scalar_t) * width);
const int vec_hidden_size = hidden_size / width;
const int64_t vec_input_stride = input_stride / width;
__shared__ float s_variance;
float variance = 0.0f;
auto* __restrict__ input_v =
reinterpret_cast<_f16Vec<scalar_t, width>*>(input);
auto* __restrict__ residual_v =
reinterpret_cast<_f16Vec<scalar_t, width>*>(residual);
auto* __restrict__ weight_v =
reinterpret_cast<const _f16Vec<scalar_t, width>*>(weight);
// Step 1: Fused add and compute variance
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
int64_t strided_id = blockIdx.x * vec_input_stride + idx;
_f16Vec<scalar_t, width> temp = input_v[strided_id];
temp += residual_v[id];
variance += temp.sum_squares();
residual_v[id] = temp; // Store updated residual
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
// Step 2: Precompute scale inverse
const float scale_inv = 1.0f / (*scale);
// Step 3: Fused RMSNorm + FP8 quantization
for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
int id = blockIdx.x * vec_hidden_size + idx;
_f16Vec<scalar_t, width> temp = residual_v[id];
temp *= s_variance;
temp *= weight_v[idx];
// Convert each element to FP8
#pragma unroll
for (int i = 0; i < width; ++i) {
float val = _typeConvert<scalar_t>::convert(temp.data[i]);
out[id * width + i] =
xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(val,
scale_inv);
}
}
}
/**
* Generic fused add + RMSNorm + FP8 quant kernel (fallback for unaligned data)
*/
template <typename scalar_t, int width, typename fp8_type>
__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
fused_add_rms_norm_static_fp8_quant_kernel(
fp8_type* __restrict__ out, // [num_tokens, hidden_size]
scalar_t* __restrict__ input, // [num_tokens, hidden_size]
const int64_t input_stride,
scalar_t* __restrict__ residual, // [num_tokens, hidden_size]
const scalar_t* __restrict__ weight, // [hidden_size]
const float* __restrict__ scale, // [1]
const float epsilon,
const int num_tokens,
const int hidden_size) {
__shared__ float s_variance;
float variance = 0.0f;
// Step 1: Fused add and compute variance
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
scalar_t z = input[blockIdx.x * input_stride + idx];
z += residual[blockIdx.x * hidden_size + idx];
float x = static_cast<float>(z);
variance += x * x;
residual[blockIdx.x * hidden_size + idx] = z; // Store updated residual
}
using BlockReduce = cub::BlockReduce<float, 1024>;
__shared__ typename BlockReduce::TempStorage reduceStore;
variance = BlockReduce(reduceStore).Reduce(variance, CubAddOp{}, blockDim.x);
if (threadIdx.x == 0) {
s_variance = rsqrtf(variance / hidden_size + epsilon);
}
__syncthreads();
// Step 2: Precompute scale inverse
const float scale_inv = 1.0f / (*scale);
// Step 3: Fused RMSNorm + FP8 quantization
for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
float x = static_cast<float>(residual[blockIdx.x * hidden_size + idx]);
float out_norm = (static_cast<scalar_t>(x * s_variance)) *
static_cast<float>(weight[idx]);
out[blockIdx.x * hidden_size + idx] =
xllm::kernel::cuda::scaled_fp8_conversion<true, fp8_type>(out_norm,
scale_inv);
}
}
#define LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(width) \
DISPATCH_FLOATING_TYPES( \
input.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \
DISPATCH_FP8_TYPES( \
out.scalar_type(), "fused_add_rms_norm_static_fp8_quant", [&] { \
fused_add_rms_norm_static_fp8_quant_kernel<scalar_t, \
width, \
fp8_t> \
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(), \
input.data_ptr<scalar_t>(), \
input_stride, \
residual.data_ptr<scalar_t>(), \
weight.data_ptr<scalar_t>(), \
scale.data_ptr<float>(), \
epsilon, \
num_tokens, \
hidden_size); \
}); \
});
} // namespace
namespace xllm::kernel::cuda {
// flashinfer rmsnorm ops
// void rmsnorm(torch::Tensor output,
// torch::Tensor input,
// torch::Tensor weight,
// double eps) {
// FunctionFactory::get_instance().rmsnorm_func("norm").call(
// output, input, weight, eps, support_pdl());
// }
void rms_norm(torch::Tensor output, // [..., hidden_size]
torch::Tensor input, // [..., hidden_size]
torch::Tensor weight, // [hidden_size]
double eps) {
CHECK(output.is_contiguous());
CHECK(weight.is_contiguous());
// The kernel addresses tokens as `blockIdx.x * input_stride + idx`, which
// can only represent contiguous inputs or simple 2D strided rows. Flux q/k
// tensors reach this path as high-dimensional transposed views, so make that
// layout explicit before flattening tokens for the kernel.
if (input.dim() > 2 && !input.is_contiguous()) {
input = input.contiguous();
}
CHECK(input.stride(-1) == 1);
int hidden_size = input.size(-1);
int num_tokens = input.numel() / hidden_size;
int64_t input_stride = input.stride(-2);
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, 1024));
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
rms_norm_kernel<scalar_t>
<<<grid, block, 0, stream>>>(output.data_ptr<scalar_t>(),
input.data_ptr<scalar_t>(),
input_stride,
weight.data_ptr<scalar_t>(),
eps,
num_tokens,
hidden_size);
});
}
void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
torch::Tensor& residual, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
double epsilon) {
CHECK(weight.scalar_type() == input.scalar_type());
CHECK(input.scalar_type() == residual.scalar_type());
CHECK(residual.is_contiguous());
CHECK(weight.is_contiguous());
int hidden_size = input.size(-1);
int64_t input_stride = input.stride(-2);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
/* This kernel is memory-latency bound in many scenarios.
When num_tokens is large, a smaller block size allows
for increased block occupancy on CUs and better latency
hiding on global mem ops. */
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 block(std::min(hidden_size, max_block_size));
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
/*If the tensor types are FP16/BF16, try to use the optimized kernel
with packed + vectorized ops.
Max optimization is achieved with a width-8 vector of FP16/BF16s
since we can load at most 128 bits at once in a global memory op.
However, this requires each tensor's data to be aligned to 16
bytes.
*/
auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
constexpr int kVectorWidth = 8;
constexpr int kReqAlignmentBytes =
kVectorWidth * 2; // kVectorWidth * sizeof(bfloat16 or float16) (float32
// falls back to non-vectorized version anyway)
bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 &&
res_ptr % kReqAlignmentBytes == 0 &&
wt_ptr % kReqAlignmentBytes == 0;
bool offsets_are_multiple_of_vector_width =
hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0;
if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {
LAUNCH_FUSED_ADD_RMS_NORM(8);
} else {
LAUNCH_FUSED_ADD_RMS_NORM(0);
}
}
// ============================================================================
// Fused RMSNorm + Static FP8 Quantization Host Functions
// ============================================================================
void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size], FP8
torch::Tensor& input, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
torch::Tensor& scale, // [1]
double epsilon) {
CHECK(out.is_contiguous());
CHECK(input.stride(-1) == 1);
CHECK(weight.is_contiguous());
CHECK(scale.is_contiguous());
int hidden_size = input.size(-1);
int64_t input_stride = input.stride(-2);
int num_tokens = input.numel() / hidden_size;
// For large num_tokens, use smaller blocks to increase SM concurrency
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 grid(num_tokens);
dim3 block(std::min(hidden_size, max_block_size));
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_FLOATING_TYPES(
input.scalar_type(), "rms_norm_static_fp8_quant", [&] {
DISPATCH_FP8_TYPES(out.scalar_type(), "rms_norm_static_fp8_quant", [&] {
rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t>
<<<grid, block, 0, stream>>>(out.data_ptr<fp8_t>(),
input.data_ptr<scalar_t>(),
input_stride,
weight.data_ptr<scalar_t>(),
scale.data_ptr<float>(),
epsilon,
num_tokens,
hidden_size);
});
});
}
void fused_add_rms_norm_static_fp8_quant(
torch::Tensor& out, // [..., hidden_size], FP8
torch::Tensor& input, // [..., hidden_size]
torch::Tensor& residual, // [..., hidden_size]
torch::Tensor& weight, // [hidden_size]
torch::Tensor& scale, // [1]
double epsilon) {
CHECK(out.is_contiguous());
CHECK(residual.is_contiguous());
CHECK(weight.is_contiguous());
CHECK(scale.is_contiguous());
CHECK(residual.scalar_type() == input.scalar_type());
CHECK(weight.scalar_type() == input.scalar_type());
int hidden_size = input.size(-1);
int64_t input_stride = input.stride(-2);
int num_tokens = input.numel() / hidden_size;
dim3 grid(num_tokens);
const int max_block_size = (num_tokens < 256) ? 1024 : 256;
dim3 block(std::min(hidden_size, max_block_size));
const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
// Check alignment for vectorized kernel
auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
auto res_ptr = reinterpret_cast<std::uintptr_t>(residual.data_ptr());
auto wt_ptr = reinterpret_cast<std::uintptr_t>(weight.data_ptr());
constexpr int kVectorWidth = 8;
constexpr int kReqAlignmentBytes = kVectorWidth * 2;
bool ptrs_are_aligned = inp_ptr % kReqAlignmentBytes == 0 &&
res_ptr % kReqAlignmentBytes == 0 &&
wt_ptr % kReqAlignmentBytes == 0;
bool offsets_are_multiple_of_vector_width =
hidden_size % kVectorWidth == 0 && input_stride % kVectorWidth == 0;
if (ptrs_are_aligned && offsets_are_multiple_of_vector_width) {
LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(8);
} else {
LAUNCH_FUSED_ADD_RMS_NORM_STATIC_FP8_QUANT(0);
}
}
} // namespace xllm::kernel::cuda

101
ex_engine/xllm_kernels/cuda/reshape_paged_cache.cu Normal file
View File

@@ -0,0 +1,101 @@
/* Copyright 2025-2026 The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <c10/cuda/CUDAStream.h>
#include "cuda_ops_api.h"
#include "device_utils.cuh"
namespace xllm::kernel::cuda {
template <typename T>
__global__ void XLLM_KERNEL_ATTR(1024) reshape_paged_cache_kernel(
const int* __restrict__ slot_ids, // [n_tokens]
const T* __restrict__ keys, // [n_tokens, n_heads, head_dim]
const T* __restrict__ values, // [n_tokens, n_heads, head_dim]
T* __restrict__ key_cache,
T* __restrict__ value_cache,
int64_t k_stride,
int64_t v_stride,
int64_t n_kv_heads,
int64_t head_dim,
int64_t block_size) {
// block/token index
const int64_t bid = blockIdx.x;
// which slot to write to
const int64_t slot_id = slot_ids[bid];
if (slot_id < 0) {
return;
}
// block index
const int64_t block_idx = slot_id / block_size;
// offset within block
const int64_t block_offset = slot_id % block_size;
// base index for the block in cache
const int64_t block_base_idx = block_idx * block_size * n_kv_heads * head_dim;
// copy value one by one for the token
for (int64_t i = threadIdx.x; i < n_kv_heads * head_dim; i += blockDim.x) {
const int64_t k_src_idx = bid * k_stride + i;
const int64_t v_src_idx = bid * v_stride + i;
// cache: [n_blocks, block_size, n_heads, head_dim]
const int64_t head_base_idx =
block_base_idx + block_offset * n_kv_heads * head_dim;
// which head to write to
const int head_idx = i / head_dim;
// which dim within head to write to
const int head_offset = i % head_dim;
const int64_t dst_idx = head_base_idx + head_idx * head_dim + head_offset;
key_cache[dst_idx] = keys[k_src_idx];
value_cache[dst_idx] = values[v_src_idx];
}
}
void reshape_paged_cache(
torch::Tensor slot_ids, // [n_tokens]
torch::Tensor keys, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor values, // [n_tokens, n_kv_heads, head_dim]
torch::Tensor key_cache, // [n_blocks, block_size, n_heads, head_dim]
torch::Tensor value_cache) {
// keys and values should be continuous at n_kv_heads and head_dim dims
CHECK(keys.stride(-1) == 1 && keys.stride(-2) == keys.size(-1));
CHECK(values.stride(-1) == 1 && values.stride(-2) == values.size(-1));
const int64_t n_tokens = keys.size(-3);
const int64_t n_kv_heads = keys.size(-2);
const int64_t head_dim = keys.size(-1);
const int64_t block_size = key_cache.size(-3);
// it is possible that keys and values have different strides
const int64_t k_stride = keys.stride(-3);
const int64_t v_stride = values.stride(-3);
const int64_t n = n_kv_heads * head_dim;
dim3 grid(n_tokens);
dim3 block(std::min<int>(n, 1024));
DISPATCH_FLOATING_TYPES(
keys.scalar_type(), "reshape_paged_cache_kernel", [&] {
reshape_paged_cache_kernel<scalar_t>
<<<grid, block, 0, c10::cuda::getCurrentCUDAStream()>>>(
slot_ids.data_ptr<int>(),
keys.data_ptr<scalar_t>(),
values.data_ptr<scalar_t>(),
key_cache.data_ptr<scalar_t>(),
value_cache.data_ptr<scalar_t>(),
k_stride,
v_stride,
n_kv_heads,
head_dim,
block_size);
});
}
} // namespace xllm::kernel::cuda

258
ex_engine/xllm_kernels/cuda/rope.cu Normal file
View File

@@ -0,0 +1,258 @@
/* Copyright 2025 The vLLM Authors and The xLLM Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <ATen/cuda/CUDAContext.h>
#include <c10/cuda/CUDAGuard.h>
#include <torch/all.h>
#include "cuda_ops_api.h"
#include "device_utils.cuh"
// ref to:
// https://github.com/vllm-project/vllm/blob/main/csrc/pos_encoding_kernels.cu
namespace {
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_token_rotary_embedding(
scalar_t* __restrict__ arr,
const scalar_t* __restrict__ cos_ptr,
const scalar_t* __restrict__ sin_ptr,
int rot_offset,
int embed_dim) {
int x_index, y_index;
scalar_t cos, sin;
if (IS_NEOX) {
// GPT-NeoX style rotary embedding.
x_index = rot_offset;
y_index = embed_dim + rot_offset;
cos = *(cos_ptr + x_index);
sin = *(sin_ptr + x_index);
} else {
// GPT-J style rotary embedding.
x_index = 2 * rot_offset;
y_index = 2 * rot_offset + 1;
cos = *(cos_ptr + x_index / 2);
sin = *(sin_ptr + x_index / 2);
}
const scalar_t x = arr[x_index];
const scalar_t y = arr[y_index];
arr[x_index] = x * cos - y * sin;
arr[y_index] = y * cos + x * sin;
}
template <typename scalar_t, bool IS_NEOX>
inline __device__ void apply_rotary_embedding(
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* cache_ptr,
const int head_size,
const int num_heads,
const int num_kv_heads,
const int rot_dim,
const int token_idx,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride) {
const int embed_dim = rot_dim / 2;
const scalar_t* cos_ptr = cache_ptr;
const scalar_t* sin_ptr = cache_ptr + embed_dim;
const int nq = num_heads * embed_dim;
for (int i = threadIdx.x; i < nq; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * query_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
query + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
if (key != nullptr) {
const int nk = num_kv_heads * embed_dim;
for (int i = threadIdx.x; i < nk; i += blockDim.x) {
const int head_idx = i / embed_dim;
const int64_t token_head =
token_idx * key_stride + head_idx * head_stride;
const int rot_offset = i % embed_dim;
apply_token_rotary_embedding<scalar_t, IS_NEOX>(
key + token_head, cos_ptr, sin_ptr, rot_offset, embed_dim);
}
}
}
template <typename scalar_t, bool IS_NEOX>
__global__ void XLLM_KERNEL_ATTR(512) rotary_embedding_kernel(
const int64_t* __restrict__ positions, // [batch_size, seq_len] or
// [num_tokens]
scalar_t* __restrict__ query, // [batch_size, seq_len, num_heads,
// head_size] or [num_tokens, num_heads,
// head_size]
scalar_t* __restrict__ key, // nullptr or
// [batch_size, seq_len, num_kv_heads,
// head_size] or [num_tokens, num_kv_heads,
// head_size]
const scalar_t* __restrict__ cos_sin_cache, // [max_position, 2,
// rot_dim // 2]
const int rot_dim,
const int64_t query_stride,
const int64_t key_stride,
const int64_t head_stride,
const int num_heads,
const int num_kv_heads,
const int head_size) {
// Each thread block is responsible for one token.
const int token_idx = blockIdx.x;
int64_t pos = positions[token_idx];
const scalar_t* cache_ptr = cos_sin_cache + pos * rot_dim;
apply_rotary_embedding<scalar_t, IS_NEOX>(query,
key,
cache_ptr,
head_size,
num_heads,
num_kv_heads,
rot_dim,
token_idx,
query_stride,
key_stride,
head_stride);
}
} // namespace
namespace xllm::kernel::cuda {
// flashinfer rope ops
// void apply_rope_pos_ids_cos_sin_cache(torch::Tensor q,
// torch::Tensor k,
// torch::Tensor cos_sin_cache,
// torch::Tensor pos_ids,
// bool interleave) {
// const int64_t head_dim = cos_sin_cache.size(-1) / 2;
// q = q.view({q.size(0), -1, head_dim});
// k = k.view({k.size(0), -1, head_dim});
// FunctionFactory::get_instance().rope_func("rope").call(
// q, k, q, k, cos_sin_cache, pos_ids, interleave);
// }
void rotary_embedding(
torch::Tensor& positions, // [batch_size, seq_len] or [num_tokens]
torch::Tensor& query, // [batch_size, seq_len, num_heads * head_size] or
// [num_tokens, num_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
std::optional<torch::Tensor> key,
// null or
// [batch_size, seq_len, num_kv_heads * head_size] or
// [num_tokens, num_kv_heads * head_size] or
// [batch_size, seq_len, num_heads, head_size] or
// [num_tokens, num_heads, head_size]
// int64_t head_size,
torch::Tensor& cos_sin_cache, // [max_position, rot_dim]
bool is_neox) {
// num_tokens = batch_size * seq_len
const int positions_ndim = positions.dim();
const int query_ndim = query.dim();
// For partial rotary models, e.g. MiniMax-M2 with head_dim=128 and
// rotary_dim=64, the cache width is the rotary dimension rather than the
// physical per-head stride. When query is already shaped as
// [*, num_heads, head_size], infer the real head_size from query itself.
int64_t head_size = (query_ndim == positions_ndim + 2)
? query.size(-1)
: cos_sin_cache.size(-1);
int64_t num_tokens = positions.numel();
// Make sure num_tokens dim is consistent across positions, query, and key
CHECK(positions_ndim == 1 || positions_ndim == 2)
<< "positions must have shape [num_tokens] or [batch_size, seq_len]";
if (positions_ndim == 1) {
CHECK(query.size(0) == positions.size(0) &&
(!key.has_value() || key->size(0) == positions.size(0)))
<< "query, key and positions must have the same number of tokens";
}
if (positions_ndim == 2) {
CHECK(query.size(0) == positions.size(0) &&
(!key.has_value() || key->size(0) == positions.size(0)) &&
query.size(1) == positions.size(1) &&
(!key.has_value() || key->size(1) == positions.size(1)))
<< "query, key and positions must have the same batch_size and seq_len";
}
// Make sure head_size is valid for query and key
// hidden_size = num_heads * head_size
int query_hidden_size = query.numel() / num_tokens;
int key_hidden_size = key.has_value() ? key->numel() / num_tokens : 0;
CHECK(query_hidden_size % head_size == 0);
CHECK(key_hidden_size % head_size == 0);
// Make sure query and key have consistent number of heads
int num_heads = query_hidden_size / head_size;
int num_kv_heads = key.has_value() ? key_hidden_size / head_size : num_heads;
CHECK(num_heads % num_kv_heads == 0);
int rot_dim = cos_sin_cache.size(1);
int seq_dim_idx = positions_ndim - 1;
int64_t query_stride = query.stride(seq_dim_idx);
int64_t key_stride = key.has_value() ? key->stride(seq_dim_idx) : 0;
// Determine head stride: for [*, heads, head_size] use stride of last dim;
// for flat [*, heads*head_size], heads blocks are contiguous of size
// head_size
int64_t head_stride =
(query_ndim == positions_ndim + 2) ? query.stride(-2) : head_size;
dim3 grid(num_tokens);
dim3 block(std::min<int64_t>(num_heads * rot_dim / 2, 512));
const at::cuda::OptionalCUDAGuard device_guard(device_of(query));
const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
DISPATCH_FLOATING_TYPES(
query.scalar_type(), "apply_rope_pos_ids_cos_sin_cache", [&] {
if (is_neox) {
rotary_embedding_kernel<scalar_t, true><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
} else {
rotary_embedding_kernel<scalar_t, false><<<grid, block, 0, stream>>>(
positions.data_ptr<int64_t>(),
query.data_ptr<scalar_t>(),
key.has_value() ? key->data_ptr<scalar_t>() : nullptr,
cos_sin_cache.data_ptr<scalar_t>(),
rot_dim,
query_stride,
key_stride,
head_stride,
num_heads,
num_kv_heads,
head_size);
}
});
}
} // namespace xllm::kernel::cuda

32
ex_engine/xllm_kernels/ilu/activation.cpp Normal file
View File

@@ -0,0 +1,32 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode) {
if (act_mode == "silu") {
infer::silu_and_mul(input, out);
} else {
LOG(FATAL) << "Unsupported act mode: " << act_mode
<< ", only support silu, gelu, gelu_tanh";
}
}
} // namespace xllm::kernel::ilu

163
ex_engine/xllm_kernels/ilu/attention.cpp Normal file
View File

@@ -0,0 +1,163 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "ixinfer.h"
#include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void reshape_paged_cache(torch::Tensor& key,
std::optional<torch::Tensor>& value,
torch::Tensor& key_cache,
std::optional<torch::Tensor>& value_cache,
torch::Tensor& slot_mapping) {
auto value_ = value.value_or(torch::Tensor());
auto value_cache_ = value_cache.value_or(torch::Tensor());
int64_t key_token_stride = key.stride(0);
int64_t value_token_stride = 0;
if (value_.defined()) {
value_token_stride = value_.stride(0);
}
slot_mapping = slot_mapping.to(at::kLong);
infer::xllm_reshape_and_cache(key,
value_,
key_cache,
value_cache_,
slot_mapping,
key_token_stride,
value_token_stride);
}
void batch_prefill(torch::Tensor& query,
const torch::Tensor& key,
const std::optional<torch::Tensor>& value,
torch::Tensor& output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_cu_seq_lens,
const std::optional<torch::Tensor>& kv_cu_seq_lens,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& attn_bias,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_quant_scale,
const std::optional<torch::Tensor>& v_quant_scale,
const torch::Tensor& block_tables,
int64_t max_query_len,
int64_t max_seq_len,
float scale,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const std::string& compute_dtype,
bool return_lse) {
double softcap = 0.0;
bool sqrt_alibi = false;
auto q_cu_seq_lens_ = q_cu_seq_lens.value_or(torch::Tensor());
auto kv_cu_seq_lens_ = kv_cu_seq_lens.value_or(torch::Tensor());
auto q_quant_scale_ = q_quant_scale.value_or(torch::Tensor());
auto k_quant_scale_ = k_quant_scale.value_or(torch::Tensor());
auto v_quant_scale_ = v_quant_scale.value_or(torch::Tensor());
auto block_tables_ = block_tables;
auto key_ = key;
auto value_ = value.value();
infer::ixinfer_flash_attn_unpad_with_block_tables(query,
key_,
value_,
output,
block_tables_,
q_cu_seq_lens_,
kv_cu_seq_lens_,
max_query_len,
max_seq_len,
is_causal,
window_size_left,
window_size_right,
static_cast<double>(scale),
softcap,
sqrt_alibi,
alibi_slope,
c10::nullopt,
output_lse);
}
void batch_decode(torch::Tensor& query,
const torch::Tensor& k_cache,
torch::Tensor& output,
const torch::Tensor& block_table,
const torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& v_cache,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_cache_quant_scale,
const std::optional<torch::Tensor>& v_cache_quant_scale,
const std::optional<torch::Tensor>& out_quant_scale,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& mask,
const std::string& compute_dtype,
int64_t max_seq_len,
int64_t window_size_left,
int64_t window_size_right,
float scale,
bool return_lse,
bool is_causal,
int64_t kv_cache_quant_bit_size) {
if (query.dim() == 4) {
query =
query
.view({query.size(0) * query.size(1), query.size(2), query.size(3)})
.contiguous();
}
if (output.dim() == 4) {
output = output
.view({output.size(0) * output.size(1),
output.size(2),
output.size(3)})
.contiguous();
;
}
auto v_cache_ = v_cache.value_or(torch::Tensor());
int64_t num_kv_heads = k_cache.size(1);
int64_t page_block_size = k_cache.size(2);
double softcap = 0.0;
bool enable_cuda_graph = false;
bool use_sqrt_alibi = false;
auto block_table_ = block_table;
auto k_cache_ = k_cache;
auto seq_lens_ = seq_lens;
infer::xllm_paged_attention(output,
query,
k_cache_,
v_cache_,
num_kv_heads,
scale,
block_table_,
seq_lens_,
page_block_size,
max_seq_len,
alibi_slope,
is_causal,
(int32_t)window_size_left,
(int32_t)window_size_right,
softcap,
enable_cuda_graph,
use_sqrt_alibi,
c10::nullopt);
}
} // namespace xllm::kernel::ilu

99
ex_engine/xllm_kernels/ilu/fused_moe.cpp Normal file
View File

@@ -0,0 +1,99 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <glog/logging.h>
#include "ilu_ops_api.h"
namespace xllm::kernel::ilu {
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias) {
torch::Tensor input_ = input.to(torch::kFloat32);
auto reduce_weight =
torch::empty({input.size(0), topk},
torch::dtype(torch::kFloat).device(input.device()));
auto topk_indices =
torch::empty({input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
auto token_expert_indices =
torch::empty({input.size(0), topk},
torch::dtype(torch::kInt32).device(input.device()));
infer::topk_softmax(
reduce_weight, topk_indices, token_expert_indices, input_, false);
auto tt = reduce_weight.sum(-1);
if (normalize) {
reduce_weight = reduce_weight / reduce_weight.sum(-1).unsqueeze(-1);
}
return std::make_tuple(reduce_weight, topk_indices);
}
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
int64_t expert_num) {
auto src_dst = expert_id.new_empty({expert_id.numel()});
auto dst_src = torch::empty_like(src_dst);
auto expert_sizes_gpu = expert_id.new_empty({expert_num});
auto expert_sizes_gpu_cumsum = expert_id.new_zeros({expert_id.numel() + 1});
infer::moe_compute_token_index_api(expert_id,
src_dst,
dst_src,
expert_sizes_gpu,
/*expert_mask=*/std::nullopt,
/*expert_sizes_cpu*/ std::nullopt,
/*expert_sizes_gpu*/ std::nullopt,
0,
expert_num,
expert_num);
expert_sizes_gpu_cumsum = expert_sizes_gpu.cumsum(-1);
return {src_dst, dst_src, expert_sizes_gpu, expert_sizes_gpu_cumsum};
}
torch::Tensor moe_expand_input(const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk) {
int64_t dst_tokens = input.size(0) * topk;
auto output = input.new_empty({dst_tokens, input.size(1)});
infer::moe_expand_input(
output, input, combine_idx, gather_index, dst_tokens, topk);
return output;
}
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight) {
input = input.view({-1, weight.size(1), input.size(1)});
auto output = input.new_empty({input.size(0), input.size(2)});
infer::moe_output_reduce_sum(output,
input,
weight,
/*mask=*/std::nullopt,
/*extra_residual*/ std::nullopt,
/*scaling_factor=*/1.0);
return output;
}
} // namespace xllm::kernel::ilu

39
ex_engine/xllm_kernels/ilu/group_gemm.cpp Normal file
View File

@@ -0,0 +1,39 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
namespace xllm::kernel::ilu {
torch::Tensor group_gemm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output) {
infer::moe_w16a16_group_gemm(
output,
input,
weight,
tokens_per_experts,
dst_to_src,
/*bias=*/std::nullopt,
/*format=*/"TN",
/*persistent=*/0,
/*output_n=*/tokens_per_experts.sum().item<int64_t>());
return output;
}
} // namespace xllm::kernel::ilu

153
ex_engine/xllm_kernels/ilu/ilu_ops_api.h Normal file
View File

@@ -0,0 +1,153 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
#include <ATen/DynamicLibrary.h>
#include <ATen/core/dispatch/Dispatcher.h>
#include <cuda_runtime.h>
#include <glog/logging.h>
#include <torch/all.h>
#include <optional>
#include "ATen/Tensor.h"
#include "ATen/cuda/CUDAEvent.h"
#include "c10/core/Device.h"
#include "c10/core/DeviceGuard.h"
#include "c10/core/GradMode.h"
#include "c10/core/InferenceMode.h"
#include "c10/core/MemoryFormat.h"
#include "c10/core/ScalarType.h"
#include "c10/core/TensorOptions.h"
#include "c10/cuda/CUDAFunctions.h"
#include "c10/cuda/CUDAGuard.h"
#include "c10/cuda/CUDAStream.h"
#include "ixformer.h"
#include "kernels/kernels.h"
// #include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& cos_sin_cache,
torch::Tensor& positions,
bool interleave);
// act_mode only support silu, gelu, gelu_tanh
void act_and_mul(torch::Tensor out,
torch::Tensor input,
const std::string& act_mode);
void reshape_paged_cache(
torch::Tensor& key, // (num_tokens, num_heads, head_size)
std::optional<torch::Tensor>& value, // (num_tokens, num_heads, head_size)
torch::Tensor& key_cache, // (num_blocks, num_heads, block_size, head_size)
std::optional<torch::Tensor>&
value_cache, // (num_blocks, num_heads, block_size, head_size)
torch::Tensor& slot_mapping); //(num_tokens)
void batch_prefill(torch::Tensor& query,
const torch::Tensor& key,
const std::optional<torch::Tensor>& value,
torch::Tensor& output,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_cu_seq_lens,
const std::optional<torch::Tensor>& kv_cu_seq_lens,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& attn_bias,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_quant_scale,
const std::optional<torch::Tensor>& v_quant_scale,
const torch::Tensor& block_tables,
int64_t max_query_len,
int64_t max_seq_len,
float scale,
bool is_causal,
int64_t window_size_left,
int64_t window_size_right,
const std::string& compute_dtype,
bool return_lse);
void batch_decode(torch::Tensor& query,
const torch::Tensor& k_cache,
torch::Tensor& output,
const torch::Tensor& block_table,
const torch::Tensor& seq_lens,
const std::optional<torch::Tensor>& v_cache,
std::optional<torch::Tensor>& output_lse,
const std::optional<torch::Tensor>& q_quant_scale,
const std::optional<torch::Tensor>& k_cache_quant_scale,
const std::optional<torch::Tensor>& v_cache_quant_scale,
const std::optional<torch::Tensor>& out_quant_scale,
const std::optional<torch::Tensor>& alibi_slope,
const std::optional<torch::Tensor>& mask,
const std::string& compute_dtype,
int64_t max_seq_len,
int64_t window_size_left,
int64_t window_size_right,
float scale,
bool return_lse,
bool is_causal,
int64_t kv_cache_quant_bit_size);
void residual_layer_norm(torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps);
void rms_norm(torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps);
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias);
std::tuple<torch::Tensor, torch::Tensor> moe_active_topk(
const torch::Tensor& input,
int64_t topk,
int64_t num_expert_group,
int64_t topk_group,
bool normalize,
const std::optional<torch::Tensor>& mask,
const std::string& normed_by,
const std::string& scoring_func,
double route_scale,
const std::optional<torch::Tensor>& e_score_correction_bias);
std::vector<torch::Tensor> moe_gen_idx(torch::Tensor& expert_id,
int64_t expert_num);
torch::Tensor moe_expand_input(const torch::Tensor& input,
const torch::Tensor& gather_index,
const torch::Tensor& combine_idx,
int64_t topk);
torch::Tensor group_gemm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& tokens_per_experts,
const std::optional<torch::Tensor>& dst_to_src,
torch::Tensor& output);
torch::Tensor moe_combine_result(torch::Tensor& input, torch::Tensor& weight);
} // namespace xllm::kernel::ilu

147
ex_engine/xllm_kernels/ilu/ixformer.h Normal file
View File

@@ -0,0 +1,147 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include <torch/all.h>
#include "ATen/Tensor.h"
#include "utils.h"
namespace ixformer::infer {
torch::Tensor ixinfer_flash_attn_unpad_with_block_tables(
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& out,
torch::Tensor& block_tables,
torch::Tensor& cu_seq_q,
torch::Tensor& cu_seq_k,
int64_t max_seq_q,
int64_t max_seq_k,
bool is_causal,
int64_t window_left,
int64_t window_right,
double scale,
double softcap,
bool sqrt_alibi,
const std::optional<torch::Tensor>& alibi_slopes,
const std::optional<torch::Tensor>& sinks,
std::optional<torch::Tensor>& lse);
void silu_and_mul(torch::Tensor& input, torch::Tensor& output);
torch::Tensor xllm_paged_attention(
torch::Tensor& out,
torch::Tensor& query,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
int64_t num_kv_heads,
double scale,
torch::Tensor& block_tables,
torch::Tensor& context_lens,
int64_t block_size,
int64_t max_context_len,
const std::optional<torch::Tensor>& alibi_slopes,
bool causal,
int32_t window_left,
int32_t window_right,
double softcap,
bool enable_cuda_graph,
bool use_sqrt_alibi,
const std::optional<torch::Tensor>& sinks);
torch::Tensor ixformer_linear(torch::Tensor& input,
torch::Tensor& weight,
int64_t act_type,
const std::optional<torch::Tensor>& bias,
const std::optional<torch::Tensor>& out,
const std::optional<bool> persistent);
torch::Tensor ixformer_linear_ex(torch::Tensor& input,
torch::Tensor& weight,
const c10::optional<torch::Tensor>& bias,
const c10::optional<torch::Tensor>& out);
void xllm_reshape_and_cache(torch::Tensor& key,
torch::Tensor& value,
torch::Tensor& key_cache,
torch::Tensor& value_cache,
torch::Tensor& slot_mapping,
int64_t key_token_stride,
int64_t value_token_stride);
void xllm_rotary_embedding(torch::Tensor& positions,
torch::Tensor& query,
torch::Tensor& key,
int64_t head_size,
torch::Tensor& cos_sin_cache,
bool is_neox);
void residual_rms_norm(torch::Tensor& input,
torch::Tensor& residual,
torch::Tensor& weight,
torch::Tensor& output,
torch::Tensor& residual_output,
const std::optional<torch::Tensor>& fused_bias,
double alpha,
double eps,
bool is_post);
void rms_norm(torch::Tensor& input,
torch::Tensor& weight,
torch::Tensor& output,
const std::optional<torch::Tensor>& fused_bias,
double eps);
void topk_softmax(torch::Tensor& topk_weights,
torch::Tensor& topk_indices,
torch::Tensor& token_expert_indices,
torch::Tensor& gating_output,
bool renormalize);
void moe_compute_token_index_api(
torch::Tensor& topk_ids,
torch::Tensor& src_dst,
torch::Tensor& dst_src,
torch::Tensor& expert_sizes_gpu,
const c10::optional<torch::Tensor>& expert_mask,
const c10::optional<torch::Tensor>& expert_sizes_cpu,
const c10::optional<torch::Tensor>& expand_tokens_gpu,
int64_t start_expert_id,
int64_t end_expert_id,
int64_t num_experts);
void moe_expand_input(torch::Tensor outputs,
torch::Tensor inputs,
torch::Tensor dst_to_src,
const c10::optional<torch::Tensor>& src_to_dst,
int64_t dst_tokens,
int64_t expand_factor);
void moe_w16a16_group_gemm(torch::Tensor output,
torch::Tensor inputs,
torch::Tensor weights,
torch::Tensor tokens_per_experts,
const c10::optional<torch::Tensor>& dst_to_src,
const c10::optional<torch::Tensor>& bias,
std::string format,
int64_t persistent,
int64_t output_n);
void moe_output_reduce_sum(torch::Tensor outputs,
torch::Tensor inputs,
const c10::optional<torch::Tensor>& mul_weight,
const c10::optional<torch::Tensor>& mask,
const c10::optional<torch::Tensor>& extra_residual,
double scaling_factor);
} // namespace ixformer::infer

73
ex_engine/xllm_kernels/ilu/matmul.cpp Normal file
View File

@@ -0,0 +1,73 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "util/env_var.h"
namespace xllm::kernel::ilu {
bool gemv_conditions(const torch::Tensor& input,
const torch::Tensor& weight,
const torch::Tensor& bias,
int64_t gemv_max_batch) {
// gemv input:[m,k] weight:[n,k]
// 1. m <= gemv_max_batch
// 2. k % 32 == 0 && n % 2 == 0
// 3. bias is None
torch::Tensor input_view = input.view({-1, input.size(-1)});
torch::Tensor weight_view = weight.view({-1, weight.size(-1)});
int64_t m = input_view.size(0);
int64_t k = input_view.size(1);
int64_t n = weight_view.size(0);
if (bias.defined() == false && m <= gemv_max_batch && k % 32 == 0 &&
n % 2 == 0) {
return true;
}
return false;
}
torch::Tensor matmul(torch::Tensor a,
torch::Tensor b,
std::optional<torch::Tensor> bias) {
int64_t act_type = -1;
bool persistent = false;
std::vector<int64_t> output_shape = a.sizes().vec();
if (!output_shape.empty()) {
output_shape[output_shape.size() - 1] = b.size(0);
}
torch::Tensor output = a.new_empty(output_shape);
bool use_gemv = true;
const int64_t gemv_max_batch = 1;
const bool disable_infer_gemm_ex =
xllm::util::get_bool_env("DISABLE_INFER_GEMM_EX", false);
use_gemv =
use_gemv &&
gemv_conditions(a, b, bias.value_or(at::Tensor()), gemv_max_batch) &&
!disable_infer_gemm_ex && (act_type == -1);
if (use_gemv) {
output = infer::ixformer_linear_ex(a, b, bias, output);
} else {
output = infer::ixformer_linear(a, b, act_type, bias, output, persistent);
}
return output;
}
} // namespace xllm::kernel::ilu

51
ex_engine/xllm_kernels/ilu/norm.cpp Normal file
View File

@@ -0,0 +1,51 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "utils.h"
using namespace ixformer;
namespace xllm::kernel::ilu {
void residual_layer_norm(torch::Tensor& input,
torch::Tensor& output,
std::optional<torch::Tensor>& residual,
torch::Tensor& weight,
std::optional<torch::Tensor>& bias,
std::optional<torch::Tensor>& residual_out,
double eps) {
auto residual_ = residual.value_or(torch::zeros_like(input));
torch::Tensor residual_out_ = residual_out.value_or(torch::zeros_like(input));
infer::residual_rms_norm(input,
residual_,
weight,
output,
residual_out_,
bias,
/*alpha=*/1.0,
eps,
false);
}
void rms_norm(torch::Tensor& output,
torch::Tensor& input,
torch::Tensor& weight,
double eps) {
std::optional<torch::Tensor> fused_bias = std::nullopt;
infer::rms_norm(input, weight, output, fused_bias, eps);
}
} // namespace xllm::kernel::ilu

31
ex_engine/xllm_kernels/ilu/rope.cpp Normal file
View File

@@ -0,0 +1,31 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#include "ilu_ops_api.h"
#include "utils.h"
namespace xllm::kernel::ilu {
void apply_rope_pos_ids_cos_sin_cache(torch::Tensor& query,
torch::Tensor& key,
torch::Tensor& cos_sin_cache,
torch::Tensor& positions,
bool interleave) {
const int64_t head_size = cos_sin_cache.size(-1);
infer::xllm_rotary_embedding(
positions, query, key, head_size, cos_sin_cache, !interleave);
}
} // namespace xllm::kernel::ilu

63
ex_engine/xllm_kernels/ilu/utils.h Normal file
View File

@@ -0,0 +1,63 @@
/* Copyright 2025-2026 The xLLM Authors.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
https://github.com/jd-opensource/xllm/blob/main/LICENSE
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
#pragma once
namespace xllm::kernel::ilu {
#undef check_tensor_contiguous
#define check_tensor_contiguous(x, type) \
TORCH_CHECK(x.scalar_type() == type); \
TORCH_CHECK(x.is_cuda()); \
TORCH_CHECK(x.is_contiguous());
#undef check_tensor_half_bf_float
#define check_tensor_half_bf_float(x) \
TORCH_CHECK(x.scalar_type() == at::ScalarType::Half || \
x.scalar_type() == at::ScalarType::Float || \
x.scalar_type() == at::ScalarType::BFloat16); \
TORCH_CHECK(x.is_cuda());
// from torchCheckMsgImpl
inline const char* ixformer_check_msg_impl(const char* msg) { return msg; }
// // If there is just 1 user-provided C-string argument, use it.
#define IXFORMER_CHECK_MSG(cond, type, ...) \
(ixformer_check_msg_impl( \
"Expected " #cond \
" to be true, but got false. " \
"(Could this error message be improved? If so, " \
"please report an enhancement request to ixformer.)", \
##__VA_ARGS__))
#define IXFORMER_CHECK(cond, ...) \
{ \
if (!(cond)) { \
std::cerr << __FILE__ << " (" << __LINE__ << ")" \
<< "-" << __FUNCTION__ << " : " \
<< IXFORMER_CHECK_MSG(cond, "", ##__VA_ARGS__) << std::endl; \
throw std::runtime_error("IXFORMER_CHECK ERROR"); \
} \
}
#undef CUINFER_CHECK
#define CUINFER_CHECK(func) \
do { \
cuinferStatus_t status = (func); \
if (status != CUINFER_STATUS_SUCCESS) { \
std::cerr << "Error in file " << __FILE__ << " on line " << __LINE__ \
<< ": " << cuinferGetErrorString(status) << std::endl; \
throw std::runtime_error("CUINFER_CHECK ERROR"); \
} \
} while (0)
} // namespace xllm::kernel::ilu