perf(kda): parallelize diagonal solves in Prefill

lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/kda.py

@@ -1122,6 +1122,7 @@ def _chunk_kda_fwd_with_cumulative_g(
         A=Akk,
         cu_seqlens=cu_seqlens,
         output_dtype=k.dtype,
+        split_diagonal=True,
     )
     del Akk
     # 3. Prepare U/W/Kg for the state update without reading the input state S_in.

lightllm/common/basemodel/triton_kernel/linear_att/fla/ops/solve_tril.py

@@ -199,6 +199,33 @@ def merge_16x16_to_32x32_inverse_kernel(
         desc_o.store([i_t * BT + 16, 16], b_Ai_22.to(desc_o.dtype, fp_downcast_rounding="rtne"))
 
 
+@triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
+@triton.jit(do_not_specialize=["T"])
+def _invert_64_diagonal_tiles(A, Diag, cu_seqlens, chunk_indices, T, H: tl.constexpr, IS_VARLEN: tl.constexpr):
+    """Invert the four independent 16x16 blocks before the full inverse merge."""
+    i_t, i_bh, d = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_b, h = i_bh // H, i_bh % H
+    if IS_VARLEN:
+        i_n = tl.load(chunk_indices + i_t * 2).to(tl.int32)
+        chunk = tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32)
+        bos = tl.load(cu_seqlens + i_n).to(tl.int32)
+        eos = tl.load(cu_seqlens + i_n + 1).to(tl.int32)
+    else:
+        bos, eos = i_b * T, (i_b + 1) * T
+        chunk = i_t
+    row = chunk * 64 + d * 16
+    i = tl.arange(0, 16)
+    base = A + ((bos + row).to(tl.int64) * H + h) * 64 + d * 16
+    tile = tl.load(base + i[:, None] * H * 64 + i[None, :], mask=(bos + row + i[:, None]) < eos, other=0).to(tl.float32)
+    inv = -tl.where(i[:, None] > i[None, :], tile, 0)
+    for j in range(2, min(16, eos - bos - row)):
+        a = -tl.load(base + j * H * 64 + i)
+        a += tl.sum(a[:, None] * inv, 0)
+        inv = tl.where((i == j)[:, None], a, inv)
+    inv += i[:, None] == i[None, :]
+    tl.store(Diag + ((i_t.to(tl.int64) * tl.num_programs(1) + i_bh) * 4 + d) * 256 + i[:, None] * 16 + i[None, :], inv)
+
+
 @triton.heuristics({"IS_VARLEN": lambda args: args["cu_seqlens"] is not None})
 @triton.jit(do_not_specialize=["T"])
 def merge_16x16_to_64x64_inverse_kernel(
@@ -212,6 +239,8 @@ def merge_16x16_to_64x64_inverse_kernel(
     USE_TMA: tl.constexpr,
     IS_VARLEN: tl.constexpr,
     DOT_PRECISION: tl.constexpr,
+    Diag=None,
+    PREINVERTED: tl.constexpr = False,
 ):
     i_t, i_bh = tl.program_id(0), tl.program_id(1)
     i_b, i_h = i_bh // H, i_bh % H
@@ -234,49 +263,59 @@ def merge_16x16_to_64x64_inverse_kernel(
     A += (bos * H + i_h) * BT
     Ai += (bos * H + i_h) * BT
 
-    if not USE_TMA:
-        p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0))
-        p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0))
-        p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0))
-        p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0))
-        b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
-        b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
-        b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32)
-        b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32)
-    else:
+    if USE_TMA:
         desc = make_tensor_descriptor(A, [T, BT], [H * BT, 1], [16, 16])
         desc_o = make_tensor_descriptor(Ai, [T, BT], [H * BT, 1], [16, 16])
-        b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
-        b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
-        b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32)
-        b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32)
-
-    # [16, 16]
-    b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
-    b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
-    b_Ai_33 = -tl.where(m_A, b_Ai_33, 0)
-    b_Ai_44 = -tl.where(m_A, b_Ai_44, 0)
+    if PREINVERTED:
+        base = (
+            Diag + (tl.program_id(0).to(tl.int64) * tl.num_programs(1) + i_bh) * 1024 + o_i[:, None] * 16 + o_i[None, :]
+        )
+        b_Ai_11 = tl.load(base)
+        b_Ai_22 = tl.load(base + 256)
+        b_Ai_33 = tl.load(base + 512)
+        b_Ai_44 = tl.load(base + 768)
+    else:
+        if not USE_TMA:
+            p_A_11 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT, 0), (16, 16), (1, 0))
+            p_A_22 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 16), (16, 16), (1, 0))
+            p_A_33 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 32, 32), (16, 16), (1, 0))
+            p_A_44 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 48, 48), (16, 16), (1, 0))
+            b_Ai_11 = tl.load(p_A_11, boundary_check=(0, 1)).to(tl.float32)
+            b_Ai_22 = tl.load(p_A_22, boundary_check=(0, 1)).to(tl.float32)
+            b_Ai_33 = tl.load(p_A_33, boundary_check=(0, 1)).to(tl.float32)
+            b_Ai_44 = tl.load(p_A_44, boundary_check=(0, 1)).to(tl.float32)
+        else:
+            b_Ai_11 = desc.load([i_t * BT + 0, 0]).to(tl.float32)
+            b_Ai_22 = desc.load([i_t * BT + 16, 16]).to(tl.float32)
+            b_Ai_33 = desc.load([i_t * BT + 32, 32]).to(tl.float32)
+            b_Ai_44 = desc.load([i_t * BT + 48, 48]).to(tl.float32)
 
-    for i in range(2, min(16, T - i_t * BT)):
-        b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
-        b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
-        b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
-    for i in range(16 + 2, min(32, T - i_t * BT)):
-        b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
-        b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
-        b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
-    for i in range(32 + 2, min(48, T - i_t * BT)):
-        b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32)
-        b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0)
-        b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33)
-    for i in range(48 + 2, min(64, T - i_t * BT)):
-        b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48)
-        b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0)
-        b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44)
-    b_Ai_11 += m_I
-    b_Ai_22 += m_I
-    b_Ai_33 += m_I
-    b_Ai_44 += m_I
+        # [16, 16]
+        b_Ai_11 = -tl.where(m_A, b_Ai_11, 0)
+        b_Ai_22 = -tl.where(m_A, b_Ai_22, 0)
+        b_Ai_33 = -tl.where(m_A, b_Ai_33, 0)
+        b_Ai_44 = -tl.where(m_A, b_Ai_44, 0)
+
+        for i in range(2, min(16, T - i_t * BT)):
+            b_a_11 = -tl.load(A + (i_t * BT + i) * H * BT + o_i)
+            b_a_11 += tl.sum(b_a_11[:, None] * b_Ai_11, 0)
+            b_Ai_11 = tl.where((o_i == i)[:, None], b_a_11, b_Ai_11)
+        for i in range(16 + 2, min(32, T - i_t * BT)):
+            b_a_22 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 16)
+            b_a_22 += tl.sum(b_a_22[:, None] * b_Ai_22, 0)
+            b_Ai_22 = tl.where((o_i == i - 16)[:, None], b_a_22, b_Ai_22)
+        for i in range(32 + 2, min(48, T - i_t * BT)):
+            b_a_33 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 32)
+            b_a_33 += tl.sum(b_a_33[:, None] * b_Ai_33, 0)
+            b_Ai_33 = tl.where((o_i == i - 32)[:, None], b_a_33, b_Ai_33)
+        for i in range(48 + 2, min(64, T - i_t * BT)):
+            b_a_44 = -tl.load(A + (i_t * BT + i) * H * BT + o_i + 48)
+            b_a_44 += tl.sum(b_a_44[:, None] * b_Ai_44, 0)
+            b_Ai_44 = tl.where((o_i == i - 48)[:, None], b_a_44, b_Ai_44)
+        b_Ai_11 += m_I
+        b_Ai_22 += m_I
+        b_Ai_33 += m_I
+        b_Ai_44 += m_I
 
     if not USE_TMA:
         p_A_21 = tl.make_block_ptr(A, (T, BT), (H * BT, 1), (i_t * BT + 16, 0), (16, 16), (1, 0))
@@ -412,6 +451,8 @@ def solve_tril(
     A: torch.Tensor,
     cu_seqlens: torch.Tensor | None = None,
     output_dtype: torch.dtype = torch.float,
+    *,
+    split_diagonal: bool = False,
 ) -> torch.Tensor:
     """
     Compute the inverse of the matrix I + A
@@ -425,6 +466,9 @@ def solve_tril(
         output_dtype (torch.dtype):
             The dtype of the output tensor. Default: `torch.float`.
             If `None`, the output dtype will be the same as the input dtype.
+        split_diagonal (bool):
+            Split the four diagonal solves for 64-token chunks. Small inputs
+            and batches with excessive FP32 workspace use the original solver.
 
     Returns:
         (I + A)^-1 with the same shape as A
@@ -437,16 +481,24 @@ def solve_tril(
     NT = len(chunk_indices) if cu_seqlens is not None else triton.cdiv(T, BT)
 
     Ai = torch.zeros_like(A, dtype=output_dtype)
+    # Ensure Triton allocator is set for TMA kernels that require scratch memory
+    if is_tma_supported:
+        _ensure_triton_allocator()
+
+    split_args = {}
     if BT == 16:
         merge_fn = solve_tril_16x16_kernel
     elif BT == 32:
         merge_fn = merge_16x16_to_32x32_inverse_kernel
     elif BT == 64:
         merge_fn = merge_16x16_to_64x64_inverse_kernel
-
-    # Ensure Triton allocator is set for TMA kernels that require scratch memory
-    if is_tma_supported:
-        _ensure_triton_allocator()
+        # Bound the FP32 workspace by the FP32 input size and skip sub-chunk inputs.
+        diag_numel = NT * B * H * 4 * 16 * 16
+        if split_diagonal and T >= BT and diag_numel <= A.numel():
+            # Keep diagonal inverses in FP32 until the final output conversion.
+            diag = torch.empty((NT, B * H, 4, 16, 16), dtype=torch.float32, device=A.device)
+            _invert_64_diagonal_tiles[NT, B * H, 4](A, diag, cu_seqlens, chunk_indices, T, H=H, num_warps=4)
+            split_args = {"Diag": diag, "PREINVERTED": True}
 
     merge_fn[NT, B * H](
         A=A,
@@ -458,5 +510,6 @@ def solve_tril(
         BT=BT,
         USE_TMA=is_tma_supported,
         DOT_PRECISION=FLA_TRIL_PRECISION,
+        **split_args,
     )
     return Ai

unit_tests/common/basemodel/triton_kernel/linear_att/test_kda_autotune.py

@@ -1,5 +1,6 @@
 import collections
 import functools
+import importlib
 import inspect
 import json
 from itertools import accumulate
@@ -194,3 +195,93 @@ def test_kda_requires_packed_inputs(fused_gate, invalid_input):
     expected_message = "packed q" if invalid_input == "batch-dimension" else "cu_seqlens is required"
     with pytest.raises(AssertionError, match=expected_message):
         kernel(**inputs)
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "heads,input_dtype,output_dtype,precision",
+    [
+        (2, torch.float32, torch.bfloat16, "ieee"),
+        (32, torch.bfloat16, torch.bfloat16, "ieee"),
+        (64, torch.float16, torch.float16, "ieee"),
+        (64, torch.float32, torch.float32, "ieee"),
+        (64, torch.float32, torch.bfloat16, "tf32"),
+        (64, torch.float32, torch.bfloat16, "tf32x3"),
+    ],
+)
+def test_split_diagonal_matches_unsplit(monkeypatch, heads, input_dtype, output_dtype, precision):
+    if precision != "ieee" and torch.cuda.get_device_capability()[0] < 8:
+        pytest.skip("TF32 requires SM80+")
+    solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril")
+    monkeypatch.setattr(solver, "FLA_TRIL_PRECISION", precision)
+    lengths = (1, 16, 1519)
+    cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32)
+    rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths])
+    matrix = torch.randn(1, sum(lengths), heads, 64, device="cuda", dtype=input_dtype) * 0.025
+    matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :])
+    expected = solver.solve_tril(matrix, cu_seqlens, output_dtype)
+    launches = []
+    original = solver._invert_64_diagonal_tiles.run
+
+    def observe(*args, **kwargs):
+        launches.append(True)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe)
+    actual = solver.solve_tril(matrix, cu_seqlens, output_dtype, split_diagonal=True)
+    assert len(launches) == 1
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize("batch,tokens,heads", [(1, 64, 2), (2, 65, 3), (3, 129, 2)])
+def test_split_diagonal_matches_unsplit_dense(monkeypatch, batch, tokens, heads):
+    solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril")
+    rows = torch.arange(tokens, device="cuda") % 64
+    matrix = torch.randn(batch, tokens, heads, 64, device="cuda") * 0.025
+    matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :])
+    expected = solver.solve_tril(matrix, output_dtype=torch.bfloat16)
+    launches = []
+    original = solver._invert_64_diagonal_tiles.run
+
+    def observe(*args, **kwargs):
+        launches.append(True)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe)
+    actual = solver.solve_tril(matrix, output_dtype=torch.bfloat16, split_diagonal=True)
+    assert len(launches) == 1
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))
+
+
+@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+@pytest.mark.parametrize(
+    "lengths,expect_split",
+    [
+        ((512,) * 16, True),
+        ((64,) * 128, True),
+        ((16,) * 16, True),
+        ((128,) * 128, True),
+        ((1,) * 256, False),
+        ((32,), False),
+    ],
+    ids=["16-requests", "128-requests", "quarter-full-chunks", "16k-tokens", "mostly-empty-chunks", "tiny-batch"],
+)
+def test_split_diagonal_uses_chunk_occupancy(monkeypatch, lengths, expect_split):
+    solver = importlib.import_module("lightllm.common.basemodel.triton_kernel.linear_att.fla.ops.solve_tril")
+    cu_seqlens = torch.tensor([0, *accumulate(lengths)], device="cuda", dtype=torch.int32)
+    rows = torch.cat([torch.arange(length, device="cuda") % 64 for length in lengths])
+    matrix = torch.randn(1, sum(lengths), 2, 64, device="cuda") * 0.025
+    matrix.mul_((rows[:, None] > torch.arange(64, device="cuda"))[None, :, None, :])
+    expected = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16)
+    launches = []
+    original = solver._invert_64_diagonal_tiles.run
+
+    def observe(*args, **kwargs):
+        launches.append(True)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(solver._invert_64_diagonal_tiles, "run", observe)
+    actual = solver.solve_tril(matrix, cu_seqlens, torch.bfloat16, split_diagonal=True)
+    assert len(launches) == int(expect_split)
+    assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8))