Add files using upload-large-folder tool

create_yaml.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python3
+import os
+import yaml
+import sys
+
+def main():
+    # 从环境变量获取值，设置默认值
+    master_addr = os.getenv("MASTER_ADDR", "127.0.0.1")
+    master_port = int(os.getenv("MASTER_PORT", 29500))
+    
+    # 获取节点和GPU数量，确保转换为整数
+    try:
+        num_nodes = int(os.getenv("SENSECORE_PYTORCH_NNODES", 1))
+    except (ValueError, TypeError):
+        num_nodes = 1
+    
+    try:
+        gpus_per_node = int(os.getenv("SENSECORE_ACCELERATE_DEVICE_COUNT", 1))
+    except (ValueError, TypeError):
+        gpus_per_node = 1
+    
+    try:
+        node_rank = int(os.getenv("SENSECORE_PYTORCH_NODE_RANK", 0))
+    except (ValueError, TypeError):
+        node_rank = 0
+    
+    # 计算总进程数
+    num_processes = num_nodes * gpus_per_node
+    
+    # 配置字典
+    config = {
+        "compute_environment": "LOCAL_MACHINE",
+        "distributed_type": "DEEPSPEED",
+        "deepspeed_config": {
+            "deepspeed_config_file": "configs/ds_config.json",
+            "zero3_init_flag": True,
+            "deepspeed_multinode_launcher": "standard",
+            "deepspeed_hostfile": '/mnt/jfzn/msj/flash-linear-attention/legacy/training/hostfile.txt',
+        },
+        "machine_rank": node_rank,
+        "main_process_ip": master_addr,
+        "main_process_port": master_port,
+        "main_training_function": "main",
+        "num_machines": num_nodes,
+        "num_processes": num_processes,
+        "same_network": True,
+        "use_cpu": False,
+        "rdzv_backend": "c10d",
+        "tpu_env": [],
+        "tpu_use_cluster": False,
+        "tpu_use_sudo": False,
+    }
+    
+    # 打印配置信息
+    print("Generated Configuration:")
+    print(f"  Master: {master_addr}:{master_port}")
+    print(f"  Number of nodes: {num_nodes}")
+    print(f"  GPUs per node: {gpus_per_node}")
+    print(f"  Total processes: {num_processes}")
+    print(f"  Node rank: {node_rank}")
+    
+    # 确保配置目录存在
+    os.makedirs("configs", exist_ok=True)
+    
+    # 写入YAML文件
+    output_file = "/mnt/jfzn/msj/flash-linear-attention/legacy/training/configs/deepspeed_sencore.yaml"
+    with open(output_file, "w") as f:
+        yaml.dump(config, f, default_flow_style=False)
+    
+    
+
+
+    print(f"\nConfiguration saved to: {output_file}")
+    
+    # 同时输出文件内容供检查
+    print("\nFile content:")
+    with open(output_file, "r") as f:
+        print(f.read())
+    
+    # 读取原始 SSH 配置
+
+    input_file = '/mnt/jfzn/msj/flash-linear-attention/legacy/training/ssh_config/config'   # 你的文件名
+    output_file = '/mnt/jfzn/msj/flash-linear-attention/legacy/training/hostfile.txt'    # 输出 hostfile
+
+    hostnames = []
+
+    with open(input_file, "r") as f:
+        for line in f:
+            line = line.strip()
+            if line.startswith("Hostname"):
+                # 提取 Hostname 后面的内容
+                hostname = line.split(None, 1)[1]
+                hostnames.append(hostname)
+
+    # 写入到 hostfile，每行一个 hostname
+    with open(output_file, "w") as f:
+        for host in hostnames:
+            f.write(host+ " slots=8\n")
+
+    print(f"提取了 {len(hostnames)} 个 hostname，已写入 {output_file}")
+
+
+    # output_path='/mnt/jfzn/msj/flash-linear-attention/legacy/training/hostfile.txt'
+    # with open(output_path, "w") as f:
+    #     addr = master_addr
+    #     f.write(f"{addr} slots={8}\n")
+    #     q = addr.split('-')
+    #     q[2] = 'worker'
+    #     addr = '-'.join(q)
+    #     f.write(f"{addr} slots={8}\n")
+    # print(f"hostfile 已生成: {output_path}")
+    return 0
+
+if __name__ == "__main__":
+    sys.exit(main())

fla3/ops/utils/logsumexp.py

@@ -0,0 +1,80 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2024, Songlin Yang, Yu Zhang
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp, log
+
+
+@triton.heuristics({
+    'HAS_SCALE': lambda args: args['scale'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [1, 2, 4, 8, 16, 32]
+    ],
+    key=['D']
+)
+@triton.jit
+def logsumexp_fwd_kernel(
+    x,
+    z,
+    scale,
+    D: tl.constexpr,
+    B: tl.constexpr,
+    HAS_SCALE: tl.constexpr
+):
+    i_n, i_d = tl.program_id(0).to(tl.int64), tl.program_id(1).to(tl.int64)
+    o_d = i_d * B + tl.arange(0, B)
+    m_d = o_d < D
+
+    b_x = tl.load(x + i_n * D + o_d, mask=m_d, other=-float('inf'))
+    if HAS_SCALE:
+        b_x = b_x * scale
+    b_m = tl.max(b_x, 0)
+    b_z = log(tl.sum(exp(b_x - b_m), 0)) + b_m
+    tl.store(z + i_n * tl.cdiv(D, B) + i_d, b_z)
+
+
+def logsumexp_fwd(
+    x,
+    scale: Optional[float] = None,
+    dtype: Optional[torch.dtype] = None
+):
+    r"""
+    Compute the logsumexp of the input tensor over the last dimension.
+
+    Args:
+        x (Tensor):
+            The input tensor of any shape.
+        scale (Optional[float]):
+            The scale applied to the input tensor. Default: `None`.
+        dtype (Optional[torch.dtype]):
+            The data type of the output tensor. Default: `None`.
+    Returns:
+        Tensor: The logsumexp of the input tensor.
+    """
+
+    shape = x.shape
+    x = x.view(-1, shape[-1])
+    N, D = x.shape
+    B = min(triton.next_power_of_2(D), 64 * 1024)
+    ND = triton.cdiv(D, B)
+
+    z = x.new_empty(N, ND, dtype=torch.float)
+    logsumexp_fwd_kernel[(N, ND)](
+        x=x,
+        z=z,
+        scale=scale,
+        D=D,
+        B=B
+    )
+    z = z.logsumexp(-1).view(*shape[:-1])
+    if dtype is not None and dtype != torch.float:
+        z = z.to(dtype)
+    return z

fla3/ops/utils/matmul.py

@@ -0,0 +1,245 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# code adapted from
+# https://triton-lang.org/main/getting-started/tutorials/03-matrix-multiplication.html
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.op import exp
+from ...utils import input_guard
+
+
+# `triton.jit`'ed functions can be auto-tuned by using the `triton.autotune` decorator, which consumes:
+#   - A list of `triton.Config` objects that define different configurations of
+#       meta-parameters (e.g., `BM`) and compilation options (e.g., `num_warps`) to try
+#   - An auto-tuning *key* whose change in values will trigger evaluation of all the
+#       provided configs
+@triton.heuristics({
+    'HAS_ALPHA': lambda args: args['alpha'] is not None,
+    'HAS_BETA': lambda args: args['beta'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BM': 128, 'BK': 64, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 128, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4),
+        triton.Config({'BM': 64, 'BK': 32, 'BN': 32, 'G': 4}, num_stages=5, num_warps=2),
+        triton.Config({'BM': 32, 'BK': 32, 'BN': 64, 'G': 4}, num_stages=5, num_warps=2),
+        # Good config for fp8 inputs.
+        # triton.Config({'BM': 128, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=3, num_warps=8),
+        # triton.Config({'BM': 256, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=3, num_warps=8),
+        # triton.Config({'BM': 256, 'BK': 128, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 64, 'BK': 128, 'BN': 256, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 128, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 64, 'BN': 64, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 64, 'BK': 64, 'BN': 128, 'G': 4}, num_stages=4, num_warps=4),
+        # triton.Config({'BM': 128, 'BK': 64, 'BN': 32, 'G': 4}, num_stages=4, num_warps=4)
+    ],
+    key=['M', 'N', 'K']
+)
+@triton.jit
+def matmul_kernel(
+    # Pointers to matrices
+    a,
+    b,
+    c,
+    input,
+    alpha,
+    beta,
+    # Matrix dimensions
+    M,
+    N,
+    K,
+    # The stride variables represent how much to increase the ptr by when moving by 1
+    # element in a particular dimension. E.g. `s_am` is how much to increase `a`
+    # by to get the element one row down (A has M rows).
+    stride_ab, stride_am, stride_ak,  # a: batch, M, K
+    stride_bk, stride_bn,             # b: K, N
+    stride_cb, stride_cm, stride_cn,  # c: batch, M, N
+    # Meta-parameters
+    BM: tl.constexpr,
+    BK: tl.constexpr,
+    BN: tl.constexpr,
+    G: tl.constexpr,
+    ACTIVATION: tl.constexpr,
+    HAS_INPUT: tl.constexpr,
+    HAS_ALPHA: tl.constexpr,
+    HAS_BETA: tl.constexpr,
+    ALLOW_TF32: tl.constexpr,
+    X_DIM: tl.constexpr = 1,
+):
+    """Kernel for computing the matmul C = A x B.
+    A has shape (M, K), B has shape (K, N) and C has shape (M, N)
+    """
+    # -----------------------------------------------------------
+    # Map program ids `pid` to the block of C it should compute.
+    # This is done in a grouped ordering to promote L2 data reuse.
+    # See above `L2 Cache Optimizations` section for details.
+    i_b, i_m, i_n = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+
+    NM, NN = tl.num_programs(1), tl.num_programs(2)
+    i_m, i_n = tl.swizzle2d(i_m, i_n, NM, NN, G)
+
+    # ----------------------------------------------------------
+    # Create pointers for the first blocks of A and B.
+    # We will advance this pointer as we move in the K direction
+    # and accumulate
+    # `p_a` is a block of [BM, BK] pointers
+    # `p_b` is a block of [BK, BN] pointers
+    # See above `Pointer Arithmetic` section for details
+    a_batch_ptr = a + i_b * stride_ab
+    o_am = (i_m * BM + tl.arange(0, BM)) % M
+    o_bn = (i_n * BN + tl.arange(0, BN)) % N
+    o_k = tl.arange(0, BK)
+
+    p_a = a_batch_ptr + (o_am[:, None] * stride_am + o_k[None, :] * stride_ak)
+    p_b = b + (o_k[:, None] * stride_bk + o_bn[None, :] * stride_bn)
+
+    b_acc = tl.zeros((BM, BN), dtype=tl.float32)
+    for k in range(0, tl.cdiv(K, BK)):
+        # Load the next block of A and B, generate a mask by checking the K dimension.
+        # If it is out of bounds, set it to 0.
+        b_a = tl.load(p_a, mask=o_k[None, :] < K - k * BK, other=0.0)
+        b_b = tl.load(p_b, mask=o_k[:, None] < K - k * BK, other=0.0)
+        # We accumulate along the K dimension.
+        b_acc = tl.dot(b_a, b_b, acc=b_acc, allow_tf32=ALLOW_TF32)
+        # Advance the ptrs to the next K block.
+        p_a += BK * stride_ak
+        p_b += BK * stride_bk
+
+    o_cm = i_m * BM + tl.arange(0, BM)
+    o_cn = i_n * BN + tl.arange(0, BN)
+    mask = (o_cm[:, None] < M) & (o_cn[None, :] < N)
+
+    b_c = b_acc
+    # You can fuse arbitrary activation functions here
+    # while the b_acc is still in FP32!
+    if ACTIVATION == "leaky_relu":
+        b_c = leaky_relu(b_c)
+    elif ACTIVATION == "relu":
+        b_c = relu(b_c)
+    elif ACTIVATION == "sigmoid":
+        b_c = sigmoid(b_c)
+    elif ACTIVATION == "tanh":
+        b_c = tanh(b_c)
+
+    if HAS_ALPHA:
+        b_c *= tl.load(alpha)
+
+    if HAS_INPUT:
+        p_i = input + (stride_cm * o_cm[:, None] if X_DIM == 2 else 0) + stride_cn * o_cn[None, :]
+        mask_p = (o_cn[None, :] < N) if X_DIM == 1 else mask
+        b_i = tl.load(p_i, mask=mask_p, other=0.0).to(tl.float32)
+        if HAS_BETA:
+            b_i *= tl.load(beta)
+        b_c += b_i
+
+    # -----------------------------------------------------------
+    # Write back the block of the output matrix C with masks.
+    c_batch_ptr = c + i_b * stride_cb
+    p_c = c_batch_ptr + stride_cm * o_cm[:, None] + stride_cn * o_cn[None, :]
+    tl.store(p_c, b_c.to(c.dtype.element_ty), mask=mask)
+
+
+# We can fuse `leaky_relu` by providing it as an `ACTIVATION` meta-parameter in `matmul_kernel`.
+@triton.jit
+def leaky_relu(x):
+    return tl.where(x >= 0, x, 0.01 * x)
+
+
+@triton.jit
+def sigmoid(x):
+    # σ(x) = 1 / (1 + exp(-x))
+    return 1.0 / (1.0 + exp(-x))
+
+
+@triton.jit
+def tanh(x):
+    # tanh(x) = (exp(x) - exp(-x)) / (exp(x) + exp(-x))
+    # 2 * sigmoid(2x) - 1
+    return (exp(x) - exp(-x)) / (exp(x) + exp(-x))
+
+
+@triton.jit
+def relu(x):
+    # ReLU(x) = max(0, x)
+    return tl.maximum(x, 0.0)
+
+
+@input_guard
+def matmul(a, b, activation=''):
+    assert a.dim() in [2, 3], "a must be 2D or 3D"
+    assert b.dim() == 2, "b must be 2D"
+    assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}"
+
+    if a.dim() == 2:
+        a_dim = 2
+        a = a.unsqueeze(0).contiguous()  # (1, M, K)
+    else:
+        a_dim = 3
+    allow_tf32 = False if a.dtype == torch.float32 else True
+
+    B, M, K = a.shape[0], a.shape[1], a.shape[2]
+    K_b, N = b.shape
+    assert K == K_b, f"Incompatible K dimension: A {K} vs B {K_b}"
+    c = a.new_empty(B, M, N)
+
+    def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN']))
+    matmul_kernel[grid](
+        a, b, c, None, None, None,
+        M, N, K,
+        a.stride(0), a.stride(1), a.stride(2),  # stride_ab, stride_am, stride_ak
+        b.stride(0), b.stride(1),               # stride_bk, stride_bn (b.dim() == 2)
+        c.stride(0), c.stride(1), c.stride(2),  # stride_cb, stride_cm, stride_cn
+        ACTIVATION=activation,
+        ALLOW_TF32=allow_tf32,
+        HAS_INPUT=False,
+    )
+    return c.squeeze(0) if a_dim == 2 else c
+
+
+@input_guard
+def addmm(
+    x: torch.Tensor,
+    a: torch.Tensor,
+    b: torch.Tensor,
+    alpha: Optional[float] = None,
+    beta: Optional[float] = None,
+) -> torch.Tensor:
+    assert a.dim() in [2, 3], "a must be 2D or 3D"
+    assert b.dim() == 2, "b must be 2D"
+    assert a.shape[-1] == b.shape[0], f"Incompatible dimensions: A {a.shape}, B {b.shape}"
+
+    if a.dim() == 2:
+        a_dim = 2
+        a = a.unsqueeze(0).contiguous()  # (1, M, K)
+    else:
+        a_dim = 3
+    allow_tf32 = False if a.dtype == torch.float32 else True
+
+    B, M, K = a.shape[0], a.shape[1], a.shape[2]
+    K_b, N = b.shape
+    assert K == K_b, f"Incompatible K dimension: A {K} vs B {K_b}"
+    c = a.new_empty(B, M, N)
+
+    def grid(meta): return (B, triton.cdiv(M, meta['BM']), triton.cdiv(N, meta['BN']))
+    matmul_kernel[grid](
+        a, b, c, x, alpha, beta,
+        M, N, K,
+        a.stride(0), a.stride(1), a.stride(2),  # stride_ab, stride_am, stride_ak
+        b.stride(0), b.stride(1),               # stride_bk, stride_bn (b.dim() == 2)
+        c.stride(0), c.stride(1), c.stride(2),  # stride_cb, stride_cm, stride_cn
+        ACTIVATION=None,
+        ALLOW_TF32=allow_tf32,
+        HAS_INPUT=True,
+        X_DIM=x.dim(),
+    )
+    return c.squeeze(0) if a_dim == 2 else c

fla3/ops/utils/op.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2024, Songlin Yang, Yu Zhang
+
+import os
+
+import triton
+import triton.language as tl
+import triton.language.extra.libdevice as tldevice
+
+from ...utils import is_gather_supported
+
+if os.environ.get('FLA_USE_FAST_OPS', '0') == '1':
+    div = tldevice.fast_dividef
+    exp = tldevice.fast_expf
+    log = tldevice.fast_logf
+    log2 = tldevice.fast_log2f
+else:
+    @triton.jit
+    def div_normal(x, y):
+        return x / y
+    div = div_normal
+    exp = tl.exp
+    log = tl.log
+    log2 = tl.log2
+
+
+@triton.jit
+def safe_exp(x):
+    return exp(tl.where(x <= 0, x, float('-inf')))
+
+
+if not is_gather_supported:
+    @triton.jit
+    def gather(src, index, axis, _builder=None):
+        # This is a fallback implementation when tl.gather is not supported
+        # In order to pass triton compiler, there is no actual gather operation
+        return src
+else:
+    gather = tl.gather

fla3/ops/utils/pack.py

@@ -0,0 +1,208 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+# Code adapted from https://github.com/mayank31398/cute-kernels
+
+from typing import Optional
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_lens
+from ...utils import input_guard
+
+
+@triton.autotune(
+    configs=[
+        triton.Config({}, num_warps=num_warps)
+        for num_warps in [4, 8, 16, 32]
+    ],
+    key=['D', 'PADDING_SIDE', 'PACK']
+)
+@triton.jit
+def packunpack_sequence_kernel(
+    x,
+    y,
+    cu_seqlens,
+    S,
+    D,
+    BD: tl.constexpr,
+    PADDING_SIDE: tl.constexpr,
+    PACK: tl.constexpr,
+):
+    i_d, i_s, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    bos, eos = tl.load(cu_seqlens + i_b), tl.load(cu_seqlens + i_b + 1)
+
+    T = eos - bos
+    if PADDING_SIDE == 'left':
+        NP = S - T
+        if i_s < NP:
+            return
+        i_t = bos + (i_s - NP)
+    else:
+        if i_s >= T:
+            return
+        i_t = bos + i_s
+
+    o_d = i_d * BD + tl.arange(0, BD)
+    mask = o_d < D
+
+    if PACK:
+        b_x = tl.load(x + (i_b * S + i_s) * D + o_d, mask=mask)
+        tl.store(y + i_t * D + o_d, b_x, mask=mask)
+    else:
+        b_x = tl.load(x + i_t * D + o_d, mask=mask)
+        tl.store(y + (i_b * S + i_s) * D + o_d, b_x, mask=mask)
+
+
+def pack_sequence_fwdbwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str,
+) -> torch.Tensor:
+    B, S = x.shape[:2]
+    D = x.numel() // (B * S)
+    BD = min(triton.next_power_of_2(D), 4096)
+    ND = triton.cdiv(D, BD)
+
+    y = torch.empty(cu_seqlens[-1].item(), *x.shape[2:], device=x.device, dtype=x.dtype)
+    packunpack_sequence_kernel[ND, S, B](
+        x=x,
+        y=y,
+        cu_seqlens=cu_seqlens,
+        S=S,
+        D=D,
+        BD=BD,
+        PADDING_SIDE=padding_side,
+        PACK=True,
+    )
+    return y
+
+
+def unpack_sequence_fwdbwd(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str,
+    desired_shape: torch.Size,
+) -> torch.Tensor:
+    if desired_shape is None:
+        desired_shape = (len(cu_seqlens) - 1, prepare_lens(cu_seqlens).max().item(), *x.shape[1:])
+    y = torch.zeros(desired_shape, device=x.device, dtype=x.dtype)
+    B, S = y.shape[:2]
+    D = y.numel() // (B * S)
+    BD = min(triton.next_power_of_2(D), 4096)
+    ND = triton.cdiv(D, BD)
+
+    packunpack_sequence_kernel[ND, S, B](
+        x=x,
+        y=y,
+        cu_seqlens=cu_seqlens,
+        S=S,
+        D=D,
+        BD=BD,
+        PADDING_SIDE=padding_side,
+        PACK=False,
+    )
+    return y
+
+
+class PackSequenceFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        padding_side: str,
+    ) -> torch.Tensor:
+        assert padding_side in ['left', 'right']
+        assert x.ndim >= 2
+
+        ctx.cu_seqlens = cu_seqlens
+        ctx.padding_side = padding_side
+        ctx.desired_shape = x.shape
+
+        y = pack_sequence_fwdbwd(
+            x=x,
+            cu_seqlens=cu_seqlens,
+            padding_side=padding_side,
+        )
+        return y
+
+    @staticmethod
+    @input_guard
+    def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]:
+        dx = unpack_sequence_fwdbwd(
+            x=dy,
+            cu_seqlens=ctx.cu_seqlens,
+            padding_side=ctx.padding_side,
+            desired_shape=ctx.desired_shape,
+        )
+        return dx, *[None] * 10
+
+
+class UnpackSequenceFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        cu_seqlens: torch.Tensor,
+        padding_side: str,
+        desired_shape: Optional[torch.Size] = None,
+    ) -> torch.Tensor:
+        assert padding_side in ['left', 'right']
+        assert x.ndim >= 2
+        if desired_shape is not None:
+            assert desired_shape[0] == cu_seqlens.shape[0] - 1
+            assert desired_shape[2:] == x.shape[1:]
+
+        ctx.cu_seqlens = cu_seqlens
+        ctx.padding_side = padding_side
+
+        y = unpack_sequence_fwdbwd(
+            x=x,
+            cu_seqlens=cu_seqlens,
+            padding_side=padding_side,
+            desired_shape=desired_shape,
+        )
+        return y
+
+    @staticmethod
+    @input_guard
+    def backward(ctx, dy: torch.Tensor) -> tuple[torch.Tensor | None]:
+        dx = pack_sequence_fwdbwd(
+            x=dy,
+            cu_seqlens=ctx.cu_seqlens,
+            padding_side=ctx.padding_side,
+        )
+        return dx, None, None, None
+
+
+def pack_sequence(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str = 'left'
+) -> torch.Tensor:
+    return PackSequenceFunction.apply(
+        x,
+        cu_seqlens,
+        padding_side,
+    )
+
+
+def unpack_sequence(
+    x: torch.Tensor,
+    cu_seqlens: torch.Tensor,
+    padding_side: str = 'left',
+    desired_shape: Optional[torch.Size] = None,
+) -> torch.Tensor:
+    return UnpackSequenceFunction.apply(
+        x,
+        cu_seqlens,
+        padding_side,
+        desired_shape,
+    )

fla3/ops/utils/pooling.py

@@ -0,0 +1,207 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) 2023-2025, Songlin Yang, Yu Zhang
+
+from typing import Optional, Tuple
+
+import torch
+import triton
+import triton.language as tl
+
+from ...ops.utils.index import prepare_chunk_indices
+from ...utils import autocast_custom_bwd, autocast_custom_fwd, input_guard
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BD': BD}, num_warps=num_warps)
+        for BD in [16, 32, 64, 128]
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['BT']
+)
+@triton.jit(do_not_specialize=['T'])
+def mean_pooling_fwd_kernel(
+    x,
+    o,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_tg = i_t
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    p_x = tl.make_block_ptr(x + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
+    p_o = tl.make_block_ptr(o + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,))
+    # [BT, BD]
+    b_x = tl.load(p_x, boundary_check=(0, 1)).to(tl.float32)
+    # [BD]
+    b_o = tl.sum(b_x, axis=0) / min(BT, T - i_t * BT)
+    tl.store(p_o, b_o.to(p_o.dtype.element_ty), boundary_check=(0,))
+
+
+@triton.heuristics({
+    'IS_VARLEN': lambda args: args['cu_seqlens'] is not None
+})
+@triton.autotune(
+    configs=[
+        triton.Config({'BD': BD}, num_warps=num_warps)
+        for BD in [16, 32, 64, 128]
+        for num_warps in [1, 2, 4, 8]
+    ],
+    key=['BT']
+)
+@triton.jit(do_not_specialize=['T'])
+def mean_pooling_bwd_kernel(
+    do,
+    dx,
+    cu_seqlens,
+    chunk_indices,
+    T,
+    H: tl.constexpr,
+    D: tl.constexpr,
+    BT: tl.constexpr,
+    BD: tl.constexpr,
+    IS_VARLEN: tl.constexpr
+):
+    i_d, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, i_h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_tg = i_t
+        i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+        T = eos - bos
+        NT = tl.cdiv(T, BT)
+    else:
+        NT = tl.cdiv(T, BT)
+        i_tg = i_b * NT + i_t
+        bos, eos = i_b * T, i_b * T + T
+
+    p_dx = tl.make_block_ptr(dx + (bos * H + i_h) * D, (T, D), (H*D, 1), (i_t * BT, i_d * BD), (BT, BD), (1, 0))
+    p_do = tl.make_block_ptr(do + (i_tg * H + i_h) * D, (D,), (1,), (i_d * BD,), (BD,), (0,))
+    # [BD]
+    b_do = tl.load(p_do, boundary_check=(0,)).to(tl.float32)
+    # [BT, BD]
+    b_dx = b_do / tl.full((BT,), min(BT, T - i_t * BT), dtype=tl.float32)[:, None]
+    tl.store(p_dx, b_dx.to(p_dx.dtype.element_ty), boundary_check=(0, 1))
+
+
+def mean_pooling_fwd(
+    x: torch.Tensor,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None
+) -> torch.Tensor:
+    B, T, H, D = x.shape
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    o = x.new_empty(B, NT, H, D)
+    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H)
+    mean_pooling_fwd_kernel[grid](
+        x,
+        o,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        H=H,
+        D=D,
+        BT=BT,
+    )
+    return o
+
+
+def mean_pooling_bwd(
+    do: torch.Tensor,
+    batch_size: int,
+    seq_len: int,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None
+) -> torch.Tensor:
+    B, T, H, D = batch_size, seq_len, *do.shape[-2:]
+    BT = chunk_size
+    chunk_indices = prepare_chunk_indices(cu_seqlens, chunk_size) if cu_seqlens is not None else None
+    NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
+
+    dx = do.new_empty(B, T, H, D)
+    def grid(meta): return (triton.cdiv(D, meta['BD']), NT, B * H)
+    mean_pooling_bwd_kernel[grid](
+        do,
+        dx,
+        cu_seqlens,
+        chunk_indices,
+        T=T,
+        H=H,
+        D=D,
+        BT=BT,
+    )
+    return dx
+
+
+class MeanPoolingFunction(torch.autograd.Function):
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_fwd
+    def forward(
+        ctx,
+        x: torch.Tensor,
+        chunk_size: int,
+        cu_seqlens: Optional[torch.LongTensor] = None
+    ) -> torch.Tensor:
+        o = mean_pooling_fwd(x, chunk_size, cu_seqlens)
+        ctx.batch_size = x.shape[0]
+        ctx.seq_len = x.shape[1]
+        ctx.chunk_size = chunk_size
+        ctx.cu_seqlens = cu_seqlens
+        return o
+
+    @staticmethod
+    @input_guard
+    @autocast_custom_bwd
+    def backward(
+        ctx, do
+    ) -> Tuple[torch.Tensor, None, None]:
+        batch_size = ctx.batch_size
+        seq_len = ctx.seq_len
+        chunk_size = ctx.chunk_size
+        cu_seqlens = ctx.cu_seqlens
+        dx = mean_pooling_bwd(do, batch_size, seq_len, chunk_size, cu_seqlens)
+        return dx, None, None
+
+
+def mean_pooling(
+    x: torch.Tensor,
+    chunk_size: int,
+    cu_seqlens: Optional[torch.LongTensor] = None,
+    head_first: bool = False
+) -> torch.Tensor:
+    if head_first:
+        x = x.transpose(1, 2)
+    if cu_seqlens is not None:
+        if x.shape[0] != 1:
+            raise ValueError(
+                f"The batch size is expected to be 1 rather than {x.shape[0]} when using `cu_seqlens`."
+                f"Please ..tten variable-length inputs before processing."
+            )
+    o = MeanPoolingFunction.apply(x, chunk_size, cu_seqlens)
+    if head_first:
+        o = o.transpose(1, 2)
+    return o

flame/data.py

@@ -0,0 +1,246 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from copy import deepcopy
+from dataclasses import dataclass
+from typing import Any, Dict, Iterable, List, Union
+
+import numpy as np
+import torch
+from datasets import Dataset, IterableDataset
+from flame.logging import get_logger
+from transformers import PreTrainedTokenizer
+
+logger = get_logger(__name__)
+
+
+class HuggingfaceDataset(IterableDataset):
+
+    def __init__(
+        self,
+        dataset: Dataset,
+        tokenizer: PreTrainedTokenizer,
+        context_len: int = 2048,
+        rank: int = 0,
+        world_size: int = 1,
+        buffer_size: int = 1024
+    ) -> HuggingfaceDataset:
+
+        self.dataset = dataset
+        self.tokenizer = tokenizer
+
+        self.data = dataset.shard(world_size, rank)
+        self.context_len = context_len
+        self.rank = rank
+        self.world_size = world_size
+        self.buffer_size = buffer_size
+
+        if tokenizer.vocab_size < torch.iinfo(torch.int16).max:
+            self.dtype = torch.int16
+        elif tokenizer.vocab_size < torch.iinfo(torch.int32).max:
+            self.dtype = torch.int32
+        else:
+            self.dtype = torch.int64
+        self.states = None
+        self.buffer = torch.tensor([], dtype=self.dtype)
+        self.tokens = []
+        self.rand_id = 0
+        self.token_id = 0
+        self.rng_state = None
+        self._epoch = 0
+
+    def __iter__(self):
+        g = torch.Generator()
+        g.manual_seed(self._epoch + self.rank)
+        if self.rng_state is not None:
+            g.set_state(self.rng_state)
+
+        rand_it = self.randint(0, self.buffer_size, g=g)
+        if self.states is not None:
+            self.data.load_state_dict(self.states)
+
+        # max number of tokens allowed in the chunk buffer
+        n_tokens = self.buffer_size * self.context_len
+
+        while True:
+            for sample in self.tokenize(self.data):
+                # keep appending the samples to the token buffer
+                self.tokens += sample
+                # if the token buffer is full, start sampling
+                # NOTE: we first convert the token ids to a tensor of shape [n_chunks, context_len] for efficiency
+                if len(self.buffer) == 0 and len(self.tokens) >= n_tokens:
+                    self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=self.dtype).view(self.buffer_size, -1)
+                    self.tokens = self.tokens[n_tokens:]
+                if len(self.buffer) == self.buffer_size:
+                    yield from self.sample(rand_it)
+
+            n_chunks = len(self.tokens) // self.context_len
+            # handle the left tokens in the buffer
+            if n_chunks > 0:
+                n_tokens = n_chunks * self.context_len
+                indices = torch.randperm(n_chunks, generator=g).tolist()
+                self.buffer = torch.tensor(self.tokens[:n_tokens], dtype=torch.long).view(n_chunks, -1)
+                self.tokens = self.tokens[n_tokens:]
+                for i in indices:
+                    yield {'input_ids': self.buffer[i]}
+
+    def tokenize(self, data, batch_size: int = 64):
+        texts, states = [], []
+        for sample in data:
+            texts.append(sample['text'])
+            states.append(self.data.state_dict())
+            if len(texts) == batch_size:
+                for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                    self.states = s
+                    yield tokenized
+                texts, states = [], []
+        if len(texts) > 0:
+            for s, tokenized in zip(states, self.tokenizer(texts, return_attention_mask=False)['input_ids']):
+                self.states = s
+                yield tokenized
+
+    def sample(self, indices):
+        n_tokens = (len(self.tokens) // self.context_len) * self.context_len
+        while self.token_id < n_tokens:
+            i = next(indices)
+            start, end = self.token_id, self.token_id + self.context_len
+            self.token_id += self.context_len
+            yield {'input_ids': self.buffer[i].to(torch.long)}
+            self.buffer[i] = torch.tensor(self.tokens[start:end], dtype=self.dtype)
+        self.token_id = 0
+        self.tokens = self.tokens[n_tokens:]
+
+    def randint(
+        self,
+        low: int,
+        high: int,
+        batch_size: int = 1024,
+        g: torch.Generator = torch.Generator()
+    ) -> Iterable[int]:
+        indices = torch.empty(batch_size, dtype=torch.long)
+        while True:
+            # record the generator states before sampling
+            self.rng_state = g.get_state()
+            indices = torch.randint(low, high, (batch_size,), out=indices, generator=g)
+            for i in indices[self.rand_id:].tolist():
+                self.rand_id += 1
+                yield i
+            self.rand_id = 0
+
+    def set_epoch(self, epoch):
+        self._epoch = epoch
+        if hasattr(self.dataset, "set_epoch"):
+            self.dataset.set_epoch(epoch)
+
+    def state_dict(self):
+        return {
+            'states': self.states,
+            'buffer': self.buffer.clone(),
+            'tokens': deepcopy(self.tokens),
+            'rand_id': self.rand_id,
+            'token_id': self.token_id,
+            'rng_state': self.rng_state,
+            'epoch': self._epoch
+        }
+
+    def load_state_dict(self, state_dict):
+        self.states = state_dict['states']
+        self.buffer = state_dict['buffer'].clone()
+        self.tokens = deepcopy(state_dict['tokens'])
+        self.rand_id = state_dict['rand_id']
+        self.token_id = state_dict['token_id']
+        self.rng_state = state_dict['rng_state'].clone() if state_dict['rng_state'] is not None else None
+        self._epoch = state_dict['epoch']
+
+
+@dataclass
+class DataCollatorForLanguageModeling:
+    """
+    Data collator used for language modeling.
+
+    Args:
+        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
+            The tokenizer used for encoding the data.
+        varlen (`bool`):
+            Whether to return sequences with variable lengths.
+            If `True`, the offsets indicating the start and end of each sequence will be returned.
+            For example, if the sequence lengths are `[4, 8, 12]`,
+            the returned `input_ids` will be a long flattened tensor of shape `[1, 24]`, with `offsets` being `[0, 4, 12, 24]`.
+            If `False`, the `input_ids` with shape `[batch_size, seq_len]` will be returned directly.
+        return_tensors (`str`):
+            The type of Tensor to return. Allowable values are "pt".
+    """
+
+    tokenizer: PreTrainedTokenizer
+    varlen: bool = False
+    return_tensors: str = "pt"
+
+    def __call__(
+        self,
+        examples: List[Union[List[int], Dict[str, Any]]]
+    ) -> Dict[str, Any]:
+        if not isinstance(examples[0], Dict):
+            examples = [{'input_ids': example} for example in examples]
+
+        def tensorize(example: Dict[str, Any]) -> Dict[str, Any]:
+            tensorized = {}
+            for key in ['input_ids', 'offsets']:
+                if key not in example:
+                    continue
+                if isinstance(example[key], List):
+                    tensorized[key] = torch.tensor(example[key], dtype=torch.long)
+                elif isinstance(example[key], np.ndarray):
+                    tensorized[key] = torch.from_numpy(example[key])
+                else:
+                    tensorized[key] = example[key]
+            return tensorized
+
+        examples = list(map(tensorize, examples))
+
+        if not self.varlen:
+            length_of_first = examples[0]['input_ids'].size(0)
+            # Check if padding is necessary.
+            if all(example['input_ids'].size(0) == length_of_first for example in examples):
+                batch = {
+                    'input_ids': torch.stack([example['input_ids'] for example in examples], dim=0),
+                }
+            else:
+                # If yes, check if we have a `pad_token`.
+                if self.tokenizer._pad_token is None:
+                    raise ValueError(
+                        f"You are attempting to pad samples but the tokenizer you are using "
+                        f"({self.tokenizer.__class__.__name__}) does not have a pad token."
+                    )
+                batch = self.tokenizer.pad(examples, return_tensors=self.return_tensors, return_attention_mask=False)
+        else:
+            if len(examples) > 1:
+                raise ValueError("The batch size must be 1 for variable length inputs.")
+            batch = {
+                'input_ids': torch.cat([example['input_ids'] for example in examples], dim=0).unsqueeze(0)
+            }
+            if 'offsets' in examples[0]:
+                batch['offsets'] = torch.cat([example['offsets'] for example in examples], dim=0).unsqueeze(0)
+            else:
+                # determine boundaries by bos/eos positions
+                if self.tokenizer.add_bos_token:
+                    offsets = []
+                    if batch['input_ids'][0, 0] != self.tokenizer.bos_token_id:
+                        offsets.append(torch.tensor([0], dtype=torch.long))
+                    offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.bos_token_id))[1])
+                    offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long))
+                    batch['offsets'] = torch.cat(offsets, dim=0)
+                elif self.tokenizer.add_eos_token:
+                    offsets = [torch.tensor([0], dtype=torch.long)]
+                    offsets.append(torch.where(batch['input_ids'].eq(self.tokenizer.eos_token_id))[1] + 1)
+                    if batch['input_ids'][0, -1] != self.tokenizer.eos_token_id:
+                        offsets.append(torch.tensor([len(batch['input_ids'][0])], dtype=torch.long))
+                    batch['offsets'] = torch.cat(offsets, dim=0)
+                else:
+                    raise ValueError("You must allow the tokenizer to add either a bos or eos token as separators.")
+
+        labels = batch['input_ids'].clone()
+        if self.tokenizer.pad_token_id is not None:
+            labels[labels == self.tokenizer.pad_token_id] = -100
+        batch["labels"] = labels
+        return batch

flame/parser.py

@@ -0,0 +1,94 @@
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Optional
+
+import transformers
+from transformers import HfArgumentParser, TrainingArguments
+
+from flame.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class TrainingArguments(TrainingArguments):
+
+    model_name_or_path: str = field(
+        default=None,
+        metadata={
+            "help": "Path to the model weight or identifier from huggingface.co/models or modelscope.cn/models."
+        },
+    )
+    tokenizer: str = field(
+        default="fla-hub/gla-1.3B-100B",
+        metadata={"help": "Name of the tokenizer to use."}
+    )
+    use_fast_tokenizer: bool = field(
+        default=False,
+        metadata={"help": "Whether or not to use one of the fast tokenizer (backed by the tokenizers library)."},
+    )
+    from_config: bool = field(
+        default=True,
+        metadata={"help": "Whether to initialize models from scratch."},
+    )
+    dataset: Optional[str] = field(
+        default=None,
+        metadata={"help": "The dataset(s) to use. Use commas to separate multiple datasets."},
+    )
+    dataset_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The name of provided dataset(s) to use."},
+    )
+    cache_dir: str = field(
+        default=None,
+        metadata={"help": "Path to the cached tokenized dataset."},
+    )
+    split: str = field(
+        default="train",
+        metadata={"help": "Which dataset split to use for training and evaluation."},
+    )
+    streaming: bool = field(
+        default=False,
+        metadata={"help": "Enable dataset streaming."},
+    )
+    hf_hub_token: Optional[str] = field(
+        default=None,
+        metadata={"help": "Auth token to log in with Hugging Face Hub."},
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the pre-processing."},
+    )
+    buffer_size: int = field(
+        default=2048,
+        metadata={"help": "Size of the buffer to randomly sample examples from in dataset streaming."},
+    )
+    context_length: int = field(
+        default=2048,
+        metadata={"help": "The context length of the tokenized inputs in the dataset."},
+    )
+    varlen: bool = field(
+        default=False,
+        metadata={"help": "Enable training with variable length inputs."},
+    )
+
+
+def get_train_args():
+    parser = HfArgumentParser(TrainingArguments)
+    args, unknown_args = parser.parse_args_into_dataclasses(return_remaining_strings=True)
+
+    if unknown_args:
+        print(parser.format_help())
+        print("Got unknown args, potentially deprecated arguments: {}".format(unknown_args))
+        raise ValueError("Some specified arguments are not used by the HfArgumentParser: {}".format(unknown_args))
+
+    if args.should_log:
+        transformers.utils.logging.set_verbosity(args.get_process_log_level())
+        transformers.utils.logging.enable_default_handler()
+        transformers.utils.logging.enable_explicit_format()
+    # set seeds manually
+    transformers.set_seed(args.seed)
+    return args