perf(glm): reuse Triton kernels across query and pool lengths

lightllm/models/glm5_next/triton_kernel/context_flashattention_nopad_fp8.py

@@ -249,14 +249,14 @@ def _context_attention_fwd_fp8_direct_smallq(
     return output
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_tokens"])
 def _build_prefill_row_table_kernel(PrefillMemIndex, RowTable, num_tokens, BLOCK: tl.constexpr):
     row = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
     mem_index = tl.load(PrefillMemIndex + row, row < num_tokens, 0)
     tl.store(RowTable + mem_index, row, row < num_tokens)
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["num_tokens"])
 def _gather_prefill_kv_kernel(
     PackedKV,
     MemIndex,
@@ -302,7 +302,7 @@ def _gather_prefill_kv_kernel(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["chunk_start", "chunk_end"])
 def _compact_topk_indices_kernel(
     TopkIndices,
     KVStarts,
@@ -331,13 +331,13 @@ def _compact_topk_indices_kernel(
         )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["NUM_Q_TOKENS", "chunk_lse_stride_s", "chunk_lse_stride_h"])
 def _merge_att_output_kernel(
     Output,
     Lse,
     ChunkOutput,
     ChunkLse,
-    chunk_lse_stride_s: tl.constexpr,
+    chunk_lse_stride_s,
     chunk_lse_stride_h,
     NUM_HEADS: tl.constexpr,
     NUM_Q_TOKENS,
@@ -372,7 +372,7 @@ def _merge_att_output_kernel(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["ragged_n"])
 def _selected_prefill_compact_kernel(
     Topk,
     KVStarts,

lightllm/models/glm5_next/triton_kernel/kpool.py

@@ -178,7 +178,7 @@ def compress_pools(
     )
 
 
-@triton.jit
+@triton.jit(do_not_specialize=["TOKENS", "MAX_POOLS"])
 def _get_pool_ranges(
     Lengths,
     CuQLens,

unit_tests/models/glm5_next/test_fp8_compile_cache.py

@@ -0,0 +1,151 @@
+import pytest
+import torch
+import triton
+
+from lightllm.models.glm5_next.triton_kernel import (
+    context_flashattention_nopad_fp8 as fp8,
+)
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def record_kernel_hashes(monkeypatch, kernel):
+    hashes = []
+    original = kernel.run
+
+    def run(*args, **kwargs):
+        compiled = original(*args, **kwargs)
+        hashes.append(compiled.hash)
+        return compiled
+
+    monkeypatch.setattr(kernel, "run", run)
+    return hashes
+
+
+def test_selected_prefill_compaction_reuses_kernel_across_context_lengths(monkeypatch):
+    hashes = record_kernel_hashes(monkeypatch, fp8._selected_prefill_compact_kernel)
+    topk = torch.full((2, 2176), -1, dtype=torch.int32, device="cuda")
+    topk[:, :32] = torch.arange(32, dtype=torch.int32, device="cuda")
+    compact = torch.empty((2, 2176), dtype=torch.int64, device="cuda")
+    lengths = torch.empty(2, dtype=torch.int32, device="cuda")
+    pages = torch.empty((2, 2176), dtype=torch.int32, device="cuda")
+    cu_q = torch.empty(3, dtype=torch.int32, device="cuda")
+    for size in [65536, 65537, 98304, 131071]:
+        mapping = torch.arange(size, dtype=torch.int32, device="cuda") % 8192
+        starts = torch.tensor([0, size // 4], dtype=torch.int32, device="cuda")
+        fp8._selected_prefill_compact_kernel[(2,)](
+            topk,
+            starts,
+            mapping,
+            compact,
+            lengths,
+            pages,
+            cu_q,
+            size,
+            8192,
+            topk.stride(0),
+            topk.stride(1),
+            BATCH=2,
+            BLOCK=4096,
+            num_warps=8,
+        )
+        torch.testing.assert_close(lengths, torch.full_like(lengths, 32))
+        torch.testing.assert_close(compact[:, :32], mapping[(starts[:, None] + topk[:, :32]).long()].long())
+        torch.testing.assert_close(cu_q, torch.arange(3, dtype=torch.int32, device="cuda"))
+    assert len(hashes) == 4 and len(set(hashes)) == 1
+
+
+def test_online_attention_merge_reuses_kernel_across_query_counts(monkeypatch):
+    hashes = record_kernel_hashes(monkeypatch, fp8._merge_att_output_kernel)
+    for num_q_tokens in [3, 8, 17, 33, 64]:
+        output = torch.randn(num_q_tokens, 4, 512, dtype=torch.float32, device="cuda")
+        lse = torch.randn(4, num_q_tokens, dtype=torch.float32, device="cuda")
+        chunk_output = torch.randn_like(output)
+        chunk_lse = torch.randn_like(lse)
+        expected_lse = torch.logaddexp(lse, chunk_lse)
+        expected = output * torch.exp(lse - expected_lse).T.unsqueeze(-1) + chunk_output * torch.exp(
+            chunk_lse - expected_lse
+        ).T.unsqueeze(-1)
+        fp8._merge_att_output_kernel[(num_q_tokens, 1)](
+            output,
+            lse,
+            chunk_output,
+            chunk_lse,
+            chunk_lse.stride(1),
+            chunk_lse.stride(0),
+            NUM_HEADS=4,
+            NUM_Q_TOKENS=num_q_tokens,
+            BLOCK_HEADS=4,
+            num_warps=4,
+        )
+        torch.testing.assert_close(output, expected, rtol=2e-5, atol=2e-6)
+        torch.testing.assert_close(lse, expected_lse, rtol=2e-5, atol=2e-6)
+    assert len(hashes) == 5 and len(set(hashes)) == 1
+
+
+def test_prefill_row_table_and_kv_gather_reuse_kernels_across_token_counts(monkeypatch):
+    row_table_hashes = record_kernel_hashes(monkeypatch, fp8._build_prefill_row_table_kernel)
+    gather_hashes = record_kernel_hashes(monkeypatch, fp8._gather_prefill_kv_kernel)
+    cached_kv = torch.randn(64, 512, dtype=torch.bfloat16, device="cuda")
+    packed_kv = torch.zeros(64, 1, 584, dtype=torch.uint8, device="cuda")
+    packed_kv[:, 0, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).view(torch.uint8)
+    packed_kv[:, 0, 448:576] = cached_kv[:, 448:].contiguous().view(torch.uint8)
+    packed_kv[:, 0, 576:583] = 127
+    cached_kv[:, :448] = cached_kv[:, :448].to(torch.float8_e4m3fn).to(torch.bfloat16)
+    for num_tokens in [3, 8, 17, 33]:
+        prefill_mem_index = torch.arange(0, num_tokens, 2, dtype=torch.int32, device="cuda")
+        prefill_row_table = torch.full((64,), -1, dtype=torch.int32, device="cuda")
+        prefill_kv = torch.randn(prefill_mem_index.numel(), 1, 512, dtype=torch.bfloat16, device="cuda")
+        fp8._build_prefill_row_table_kernel[(1,)](
+            prefill_mem_index, prefill_row_table, prefill_mem_index.numel(), BLOCK=256
+        )
+        mem_index = torch.arange(num_tokens, dtype=torch.int32, device="cuda")
+        output = torch.empty(num_tokens, 1, 512, dtype=torch.bfloat16, device="cuda")
+        fp8._gather_prefill_kv_kernel[(triton.cdiv(num_tokens, 4),)](
+            packed_kv,
+            mem_index,
+            prefill_row_table,
+            prefill_kv,
+            output,
+            num_tokens,
+            packed_kv.stride(0),
+            prefill_kv.stride(0),
+            prefill_kv.stride(2),
+            BLOCK_ROWS=4,
+        )
+        expected = cached_kv[:num_tokens].clone()
+        expected[prefill_mem_index.long()] = prefill_kv[:, 0]
+        torch.testing.assert_close(output[:, 0], expected, rtol=0, atol=0)
+    assert len(set(row_table_hashes)) == len(set(gather_hashes)) == 1
+
+
+@pytest.mark.parametrize("compact", [False, True])
+def test_topk_compaction_reuses_kernel_across_chunk_boundaries(monkeypatch, compact):
+    hashes = record_kernel_hashes(monkeypatch, fp8._compact_topk_indices_kernel)
+    topk_indices = torch.tensor([[-1, 0, 2, 3, 8, 17, 32]] * 2, dtype=torch.int32, device="cuda")
+    kv_starts = torch.tensor([0, 64], dtype=torch.int32, device="cuda")
+    ragged_indices = topk_indices + kv_starts[:, None]
+    for chunk_start, chunk_end in [(0, 3), (3, 19), (19, 128)]:
+        chunk_indices = torch.full_like(topk_indices, -2)
+        chunk_lengths = torch.empty(2, dtype=torch.int32, device="cuda")
+        fp8._compact_topk_indices_kernel[(2,)](
+            topk_indices,
+            kv_starts,
+            chunk_indices,
+            chunk_lengths,
+            chunk_start,
+            chunk_end,
+            TOPK=7,
+            BLOCK=8,
+            COMPACT=compact,
+        )
+        valid = (topk_indices >= 0) & (ragged_indices >= chunk_start) & (ragged_indices < chunk_end)
+        expected = ragged_indices - chunk_start
+        if compact:
+            torch.testing.assert_close(chunk_lengths, valid.sum(dim=1).int())
+            for row in range(2):
+                selected = expected[row, valid[row]]
+                torch.testing.assert_close(chunk_indices[row, : selected.numel()], selected, rtol=0, atol=0)
+        else:
+            torch.testing.assert_close(chunk_indices, expected.masked_fill(~valid, -1), rtol=0, atol=0)
+    assert len(set(hashes)) == 1

unit_tests/models/glm5_next/test_kpool_compile_cache.py

@@ -0,0 +1,49 @@
+import pytest
+import torch
+from lightllm.models.glm5_next.triton_kernel import kpool
+
+
+pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")
+
+
+def record_kernels(monkeypatch, kernel):
+    hashes = []
+    original = kernel.run
+
+    def run(*args, **kwargs):
+        compiled = original(*args, **kwargs)
+        hashes.append(compiled.hash)
+        return compiled
+
+    monkeypatch.setattr(kernel, "run", run)
+    return hashes
+
+
+def test_pool_ranges_reuses_binary_across_context_and_chunk_lengths(monkeypatch):
+    hashes = record_kernels(monkeypatch, kpool._get_pool_ranges)
+    for tokens, max_pools in [
+        (5, 7),
+        (16, 8),
+        (17, 31),
+        (255, 32),
+        (256, 33),
+        (257, 1025),
+        (513, 262144),
+    ]:
+        split = tokens // 2
+        lengths = torch.arange(tokens, device="cuda", dtype=torch.int32) % (max_pools * 4) + 1
+        cu = torch.tensor([0, split, tokens], device="cuda", dtype=torch.int32)
+        expected_starts = torch.cat(
+            (
+                torch.zeros_like(lengths[:split]),
+                torch.full_like(lengths[split:], max_pools),
+            )
+        )
+        expected_lengths = lengths // 4
+        actual = kpool.get_pool_ranges(lengths, cu, tokens - split, max_pools)
+        for output, expected in zip(
+            actual,
+            (expected_starts, expected_starts + expected_lengths, expected_lengths),
+        ):
+            assert torch.equal(output, expected)
+    assert len(set(hashes)) == 1