Initial K3 snapshot: 0.5B KDA/MLA/MoE train path
Standalone tree split from LLMRL/projects/kda. Includes Triton dt_bias backward fix, train_k3 --preset 0.5b, SFT, Docker runtime, and tests.
This commit is contained in:
dela
@@ -0,0 +1,468 @@
|
||||
# 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 kda._fla.ops.backends import dispatch
|
||||
from kda._fla.ops.utils.cache import fla_cache_autotune
|
||||
from kda._fla.ops.utils.index import prepare_chunk_indices
|
||||
from kda._fla.utils import autotune_cache_kwargs, check_shared_mem, input_guard
|
||||
|
||||
BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
|
||||
|
||||
|
||||
@triton.heuristics({
|
||||
'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=['B', 'H', 'BT', 'IS_VARLEN', 'REVERSE'],
|
||||
**autotune_cache_kwargs,
|
||||
)
|
||||
@triton.jit(do_not_specialize=['T'])
|
||||
def chunk_local_cumsum_scalar_kernel(
|
||||
s,
|
||||
o,
|
||||
scale,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
REVERSE: 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_s = s + bos*H + i_h + o_t * H
|
||||
p_o = o + bos*H + i_h + o_t * H
|
||||
# [BT]
|
||||
b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32)
|
||||
if REVERSE:
|
||||
b_o = tl.cumsum(b_s, axis=0, reverse=True)
|
||||
else:
|
||||
b_o = tl.cumsum(b_s, axis=0)
|
||||
if HAS_SCALE:
|
||||
b_o *= scale
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t)
|
||||
|
||||
|
||||
@triton.heuristics({
|
||||
'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({'BS': BS}, num_warps=num_warps)
|
||||
for BS in BS_LIST
|
||||
for num_warps in [2, 4, 8]
|
||||
],
|
||||
key=['B', 'H', 'S', 'BT', 'IS_VARLEN', 'REVERSE'],
|
||||
**autotune_cache_kwargs,
|
||||
)
|
||||
@triton.jit(do_not_specialize=['T'])
|
||||
def chunk_local_cumsum_vector_kernel(
|
||||
s,
|
||||
o,
|
||||
scale,
|
||||
cu_seqlens,
|
||||
chunk_indices,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
S: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
BS: tl.constexpr,
|
||||
REVERSE: tl.constexpr,
|
||||
HAS_SCALE: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
):
|
||||
i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1).to(tl.int64), tl.program_id(2).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)
|
||||
o_s = i_s * BS + tl.arange(0, BS)
|
||||
m_s = (o_t[:, None] < T) & (o_s[None, :] < S)
|
||||
p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :]
|
||||
p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :]
|
||||
# [BT, BS]
|
||||
b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32)
|
||||
if REVERSE:
|
||||
b_o = tl.cumsum(b_s, axis=0, reverse=True)
|
||||
else:
|
||||
b_o = tl.cumsum(b_s, axis=0)
|
||||
if HAS_SCALE:
|
||||
b_o *= scale
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_s)
|
||||
|
||||
|
||||
@triton.heuristics({
|
||||
'HAS_SCALE': lambda args: args['scale'] is not None,
|
||||
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
|
||||
})
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
|
||||
for BT in [32, 64, 128, 256]
|
||||
for num_warps in [2, 4, 8]
|
||||
for num_stages in [1, 2, 3, 4]
|
||||
],
|
||||
key=['B', 'H', 'IS_VARLEN', 'REVERSE'],
|
||||
**autotune_cache_kwargs,
|
||||
)
|
||||
@triton.jit(do_not_specialize=['T'])
|
||||
def chunk_global_cumsum_scalar_kernel(
|
||||
s,
|
||||
o,
|
||||
scale,
|
||||
cu_seqlens,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
REVERSE: tl.constexpr,
|
||||
HAS_SCALE: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
):
|
||||
i_nh = tl.program_id(0).to(tl.int64)
|
||||
i_n, i_h = i_nh // H, i_nh % 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)
|
||||
else:
|
||||
bos, eos = i_n * T, i_n * T + T
|
||||
T = eos - bos
|
||||
|
||||
b_z = tl.zeros([], dtype=tl.float32)
|
||||
NT = tl.cdiv(T, BT)
|
||||
for i_c in range(NT):
|
||||
i_t = NT - 1 - i_c if REVERSE else i_c
|
||||
o_t = i_t * BT + tl.arange(0, BT)
|
||||
m_t = o_t < T
|
||||
p_s = s + bos*H + i_h + o_t * H
|
||||
p_o = o + bos*H + i_h + o_t * H
|
||||
b_s = tl.load(p_s, mask=m_t, other=0.0).to(tl.float32)
|
||||
if REVERSE:
|
||||
b_o = tl.cumsum(b_s, axis=0, reverse=True)
|
||||
else:
|
||||
b_o = tl.cumsum(b_s, axis=0)
|
||||
b_ss = tl.sum(b_s, 0)
|
||||
b_o += b_z
|
||||
if i_c >= 0:
|
||||
b_z += b_ss
|
||||
if HAS_SCALE:
|
||||
b_o *= scale
|
||||
tl.store(p_o, b_o.to(p_o.dtype.element_ty), mask=m_t)
|
||||
|
||||
|
||||
@triton.heuristics({
|
||||
'HAS_SCALE': lambda args: args['scale'] is not None,
|
||||
'IS_VARLEN': lambda args: args['cu_seqlens'] is not None,
|
||||
})
|
||||
@triton.autotune(
|
||||
configs=[
|
||||
triton.Config({'BT': BT}, num_warps=num_warps, num_stages=num_stages)
|
||||
for BT in [16, 32, 64, 128]
|
||||
for num_warps in [2, 4, 8]
|
||||
for num_stages in [1, 2, 3, 4]
|
||||
],
|
||||
key=['B', 'H', 'S', 'IS_VARLEN', 'REVERSE'],
|
||||
**autotune_cache_kwargs,
|
||||
)
|
||||
@triton.jit(do_not_specialize=['T'])
|
||||
def chunk_global_cumsum_vector_kernel(
|
||||
s,
|
||||
o,
|
||||
scale,
|
||||
cu_seqlens,
|
||||
T,
|
||||
B: tl.constexpr,
|
||||
H: tl.constexpr,
|
||||
S: tl.constexpr,
|
||||
BT: tl.constexpr,
|
||||
BS: tl.constexpr,
|
||||
REVERSE: tl.constexpr,
|
||||
HAS_SCALE: tl.constexpr,
|
||||
IS_VARLEN: tl.constexpr,
|
||||
):
|
||||
i_s, i_nh = tl.program_id(0), tl.program_id(1).to(tl.int64)
|
||||
i_n, i_h = i_nh // H, i_nh % 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)
|
||||
else:
|
||||
bos, eos = i_n * T, i_n * T + T
|
||||
T = eos - bos
|
||||
|
||||
b_z = tl.zeros([BS], dtype=tl.float32)
|
||||
NT = tl.cdiv(T, BT)
|
||||
for i_c in range(NT):
|
||||
i_t = NT - 1 - i_c if REVERSE else i_c
|
||||
o_t = i_t * BT + tl.arange(0, BT)
|
||||
o_s = i_s * BS + tl.arange(0, BS)
|
||||
m_s = (o_t[:, None] < T) & (o_s[None, :] < S)
|
||||
p_s = s + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :]
|
||||
p_o = o + (bos * H + i_h) * S + o_t[:, None] * (H*S) + o_s[None, :]
|
||||
# [BT, BS]
|
||||
b_s = tl.load(p_s, mask=m_s, other=0.0).to(tl.float32)
|
||||
if REVERSE:
|
||||
b_c = b_z[None, :] + tl.cumsum(b_s, axis=0, reverse=True)
|
||||
else:
|
||||
b_c = b_z[None, :] + tl.cumsum(b_s, axis=0)
|
||||
if HAS_SCALE:
|
||||
b_c *= scale
|
||||
tl.store(p_o, b_c.to(p_o.dtype.element_ty), mask=m_s)
|
||||
b_z += tl.sum(b_s, 0)
|
||||
|
||||
|
||||
def chunk_local_cumsum_scalar(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
scale: float = None,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
chunk_indices: torch.LongTensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
B, T, H = g.shape
|
||||
assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
|
||||
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)
|
||||
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
|
||||
grid = (NT, B * H)
|
||||
chunk_local_cumsum_scalar_kernel[grid](
|
||||
s=g_org,
|
||||
o=g,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
BT=BT,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return g
|
||||
|
||||
|
||||
def chunk_local_cumsum_vector(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
scale: float = None,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
chunk_indices: torch.LongTensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
B, T, H, S = 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)
|
||||
assert chunk_size == 2**(chunk_size.bit_length()-1), "chunk_size must be a power of 2"
|
||||
|
||||
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
|
||||
def grid(meta): return (triton.cdiv(meta['S'], meta['BS']), NT, B * H)
|
||||
# keep cummulative normalizer in fp32
|
||||
# this kernel is equivalent to
|
||||
# g = g.view(B, H, NT, BT, -1).cumsum(-2).view(B, H, T, -1)
|
||||
chunk_local_cumsum_vector_kernel[grid](
|
||||
s=g_org,
|
||||
o=g,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
chunk_indices=chunk_indices,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
S=S,
|
||||
BT=BT,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return g
|
||||
|
||||
|
||||
@input_guard
|
||||
def chunk_global_cumsum_scalar(
|
||||
s: torch.Tensor,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
scale: float = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
B, T, H = s.shape
|
||||
N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
|
||||
|
||||
z = torch.empty_like(s, dtype=output_dtype or s.dtype)
|
||||
grid = (N * H,)
|
||||
chunk_global_cumsum_scalar_kernel[grid](
|
||||
s=s,
|
||||
o=z,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return z
|
||||
|
||||
|
||||
@input_guard
|
||||
def chunk_global_cumsum_vector(
|
||||
s: torch.Tensor,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
scale: float = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
B, T, H, S = s.shape
|
||||
N = len(cu_seqlens) - 1 if cu_seqlens is not None else B
|
||||
BS = min(32, triton.next_power_of_2(S))
|
||||
|
||||
z = torch.empty_like(s, dtype=output_dtype or s.dtype)
|
||||
grid = (triton.cdiv(S, BS), N * H)
|
||||
chunk_global_cumsum_vector_kernel[grid](
|
||||
s=s,
|
||||
o=z,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
T=T,
|
||||
B=B,
|
||||
H=H,
|
||||
S=S,
|
||||
BS=BS,
|
||||
REVERSE=reverse,
|
||||
)
|
||||
return z
|
||||
|
||||
|
||||
@input_guard
|
||||
@dispatch('utils')
|
||||
def chunk_global_cumsum(
|
||||
s: torch.Tensor,
|
||||
reverse: bool = False,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
scale: float = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
if cu_seqlens is not None:
|
||||
assert s.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"
|
||||
if len(s.shape) == 3:
|
||||
return chunk_global_cumsum_scalar(
|
||||
s=s,
|
||||
reverse=reverse,
|
||||
cu_seqlens=cu_seqlens,
|
||||
scale=scale,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
elif len(s.shape) == 4:
|
||||
return chunk_global_cumsum_vector(
|
||||
s=s,
|
||||
reverse=reverse,
|
||||
cu_seqlens=cu_seqlens,
|
||||
scale=scale,
|
||||
output_dtype=output_dtype,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported input shape {s.shape}, "
|
||||
f"which should be [B, T, H] or [B, T, H, D]",
|
||||
)
|
||||
|
||||
|
||||
@input_guard
|
||||
@dispatch('utils')
|
||||
def chunk_local_cumsum(
|
||||
g: torch.Tensor,
|
||||
chunk_size: int,
|
||||
reverse: bool = False,
|
||||
scale: float = None,
|
||||
cu_seqlens: torch.Tensor | None = None,
|
||||
output_dtype: torch.dtype | None = torch.float,
|
||||
chunk_indices: torch.LongTensor | None = None,
|
||||
**kwargs,
|
||||
) -> torch.Tensor:
|
||||
if 'head_first' in kwargs:
|
||||
raise DeprecationWarning(
|
||||
"head_first has been removed. Inputs must be in `[B, T, H, ...]` format.",
|
||||
)
|
||||
if cu_seqlens is not None:
|
||||
assert g.shape[0] == 1, "Only batch size 1 is supported when cu_seqlens are provided"
|
||||
if len(g.shape) == 3:
|
||||
return chunk_local_cumsum_scalar(
|
||||
g=g,
|
||||
chunk_size=chunk_size,
|
||||
reverse=reverse,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
output_dtype=output_dtype,
|
||||
chunk_indices=chunk_indices,
|
||||
)
|
||||
elif len(g.shape) == 4:
|
||||
return chunk_local_cumsum_vector(
|
||||
g=g,
|
||||
chunk_size=chunk_size,
|
||||
reverse=reverse,
|
||||
scale=scale,
|
||||
cu_seqlens=cu_seqlens,
|
||||
output_dtype=output_dtype,
|
||||
chunk_indices=chunk_indices,
|
||||
)
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported input shape {g.shape}, "
|
||||
f"which should be (B, T, H) or (B, T, H, D)",
|
||||
)
|
||||
Reference in New Issue
Block a user