From e5bc280c722b7e531d1682befd4540ae2fd991f8 Mon Sep 17 00:00:00 2001 From: baishihao Date: Fri, 30 Oct 2026 18:30:10 +0800 Subject: [PATCH] perf(glm): batch K-pool gathers and reuse decode pages --- .../models/glm5_next/triton_kernel/kpool.py | 94 +++++++++---- test/benchmark_glm53_kpool.py | 129 ++++++++++++++++++ .../models/glm5_next/test_kpool_gather.py | 124 +++++++++++++++++ 3 files changed, 324 insertions(+), 23 deletions(-) create mode 100644 test/benchmark_glm53_kpool.py create mode 100644 unit_tests/models/glm5_next/test_kpool_gather.py diff --git a/lightllm/models/glm5_next/triton_kernel/kpool.py b/lightllm/models/glm5_next/triton_kernel/kpool.py index 87f67a798e..4b9ae46e70 100644 --- a/lightllm/models/glm5_next/triton_kernel/kpool.py +++ b/lightllm/models/glm5_next/triton_kernel/kpool.py @@ -219,7 +219,7 @@ def get_pool_ranges(lengths, cu_q_lens, max_q_len, max_pools): return starts, ends, pool_lengths -@triton.jit +@triton.jit(do_not_specialize=["POOLS"]) def _gather_pools( Packed, ReqTable, @@ -230,25 +230,33 @@ def _gather_pools( PACKED_STRIDE: tl.constexpr, REQ_STRIDE: tl.constexpr, POOLS, + BLOCK_POOLS: tl.constexpr, ): - pool, batch = tl.program_id(0), tl.program_id(1) + # Move multiple independent pools per CTA instead of launching one CTA for + # each 128-byte key. This is a byte-preserving gather, including FP32 scales. + batch = tl.program_id(1) + pool = tl.program_id(0) * BLOCK_POOLS + tl.arange(0, BLOCK_POOLS) length = tl.load(SeqLen + batch) req = tl.load(ReqIdx + batch) - valid = pool < length // 4 + valid = (pool < POOLS) & (pool < length // 4) loc = tl.load(ReqTable + req * REQ_STRIDE + pool * 4 + 3, valid, 0).to(tl.int64) cols = tl.arange(0, 128) - packed = tl.load(Packed + loc * PACKED_STRIDE + cols, valid, 0) + packed = tl.load(Packed + loc[:, None] * PACKED_STRIDE + cols[None, :], valid[:, None], 0) scale = tl.load((Packed + loc * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) row = batch.to(tl.int64) * POOLS + pool - tl.store(K + row * 128 + cols, packed.to(tl.float8e4nv, bitcast=True)) - tl.store(Scale + row, scale) + tl.store(K + row[:, None] * 128 + cols[None, :], packed.to(tl.float8e4nv, bitcast=True), (pool < POOLS)[:, None]) + tl.store(Scale + row, scale, pool < POOLS) def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): - keys = torch.empty((req_idx.numel() * max_pools, 128), device=packed_buffer.device, dtype=torch.float8_e4m3fn) + keys = torch.empty( + (req_idx.numel() * max_pools, 128), + device=packed_buffer.device, + dtype=torch.float8_e4m3fn, + ) scales = torch.empty((keys.shape[0],), device=keys.device, dtype=torch.float32) # Long contexts can exceed grid Y's 65535-block limit; put pools on X. - _gather_pools[(max_pools, req_idx.numel())]( + _gather_pools[(triton.cdiv(max_pools, 32), req_idx.numel())]( packed_buffer, req_table, req_idx, @@ -258,56 +266,96 @@ def gather_pools(packed_buffer, req_table, req_idx, seq_len, max_pools): packed_buffer.stride(0), req_table.stride(0), max_pools, + BLOCK_POOLS=32, num_warps=4, ) return keys, scales -@triton.jit +@triton.jit(do_not_specialize=["MAX_PAGES"]) def _gather_paged_pools( Packed, ReqTable, ReqIdx, PoolLengths, Pages, + BlockTable, PACKED_STRIDE: tl.constexpr, REQ_STRIDE: tl.constexpr, - MAX_PAGES: tl.constexpr, + MAX_PAGES, + BATCH: tl.constexpr, + REUSE_GROUP: tl.constexpr, ): batch = tl.program_id(1) req = tl.load(ReqIdx + batch) length = tl.load(PoolLengths + batch) + # MTP verification commonly has four adjacent queries for the same request. + # Share their packed keys, while DeepGEMM still receives each row's original + # causal length. Test IDs on device on every replay: arbitrary layouts, + # dynamic verification, HOLD rows and recycled slots remain valid. + owner = batch + if REUSE_GROUP > 1: + group_start = batch // REUSE_GROUP * REUSE_GROUP + members = group_start + tl.arange(0, REUSE_GROUP) + requests = tl.load(ReqIdx + members, members < BATCH, -1) + lengths = tl.load(PoolLengths + members, members < BATCH, 0) + same_request = ( + (group_start + REUSE_GROUP <= BATCH) + & (tl.min(requests, 0) == tl.max(requests, 0)) + & (tl.max(lengths, 0) >= 4096) + ) + owner = tl.where(same_request, group_start, batch) + length = tl.where(same_request, tl.max(lengths, 0), length) rows = tl.arange(0, 64) cols = tl.arange(0, 128) # A fixed grid is replayable at 1M; the GPU length bounds the work. - for page in range(tl.program_id(0), tl.cdiv(length, 64), tl.num_programs(0)): - pools = page * 64 + rows - valid = pools < length - locs = tl.load(ReqTable + req * REQ_STRIDE + pools * 4 + 3, valid, 0).to(tl.int64) - keys = tl.load(Packed + locs[:, None] * PACKED_STRIDE + cols[None, :], valid[:, None], 0) - scales = tl.load((Packed + locs * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), valid, 1.0) - dest = Pages + (batch.to(tl.int64) * MAX_PAGES + page) * (64 * 132) - # DeepGEMM stores 64 FP8 keys followed by their 64 FP32 scales. - tl.store(dest + rows[:, None] * 128 + cols[None, :], keys) - tl.store((dest + 64 * 128).to(tl.pointer_type(tl.float32)) + rows, scales) + for page in range(tl.program_id(0), tl.cdiv(tl.maximum(length, 1), 64), tl.num_programs(0)): + if batch != owner: + tl.store(BlockTable + batch * MAX_PAGES + page, owner * MAX_PAGES + page) + else: + pools = page * 64 + rows + valid = pools < length + locs = tl.load(ReqTable + req * REQ_STRIDE + pools * 4 + 3, valid, 0).to(tl.int64) + keys = tl.load(Packed + locs[:, None] * PACKED_STRIDE + cols[None, :], valid[:, None], 0) + scales = tl.load( + (Packed + locs * PACKED_STRIDE + 128).to(tl.pointer_type(tl.float32)), + valid, + 1.0, + ) + dest = Pages + (batch.to(tl.int64) * MAX_PAGES + page) * (64 * 132) + # DeepGEMM stores 64 FP8 keys followed by their 64 FP32 scales. + tl.store(dest + rows[:, None] * 128 + cols[None, :], keys) + tl.store((dest + 64 * 128).to(tl.pointer_type(tl.float32)) + rows, scales) + +def gather_paged_pools(packed_buffer, req_table, req_idx, pool_lengths, max_pools, reuse_group=4): + """Pack valid pools; adjacent rows with the same slot can share their keys. -def gather_paged_pools(packed_buffer, req_table, req_idx, pool_lengths, max_pools): - """Pack valid pools into DeepGEMM pages; unused pages remain unread.""" + Capacity stays at the original per-row upper bound so graph capture never + assumes a particular request layout. Reuse affects only bytes copied. + """ + if reuse_group not in (1, 2, 4, 8): + raise ValueError("reuse_group must be 1, 2, 4 or 8") batch = req_idx.numel() max_pages = triton.cdiv(max_pools, 64) pages = torch.empty((batch * max_pages, 64, 1, 132), device=packed_buffer.device, dtype=torch.uint8) block_table = torch.arange(batch * max_pages, device=pages.device, dtype=torch.int32).view(batch, max_pages) - blocks = min(max_pages, triton.cdiv(get_device_sm_count() * 4, batch)) + # Reused rows only update the page table. Keep enough CTAs for the rows + # actually copying keys, rather than leaving most SMs with no memory work. + active_groups = max(1, triton.cdiv(batch, reuse_group)) if max_pools >= 4096 else batch + blocks = min(max_pages, triton.cdiv(get_device_sm_count() * 4, active_groups)) _gather_paged_pools[(blocks, batch)]( packed_buffer, req_table, req_idx, pool_lengths, pages, + block_table, packed_buffer.stride(0), req_table.stride(0), max_pages, + BATCH=batch, + REUSE_GROUP=reuse_group if batch >= reuse_group else 1, num_warps=4, ) return pages, block_table diff --git a/test/benchmark_glm53_kpool.py b/test/benchmark_glm53_kpool.py new file mode 100644 index 0000000000..904ee9053c --- /dev/null +++ b/test/benchmark_glm53_kpool.py @@ -0,0 +1,129 @@ +"""Compare gather wrappers from two source files; no model or server required. + +Example: + git show BASE:lightllm/models/glm5_next/triton_kernel/kpool.py >/tmp/kpool_before.py + python test/benchmark_glm53_kpool.py --baseline /tmp/kpool_before.py + +The result includes allocation/page-table initialization, excludes compilation, +DeepGEMM, model computation and request scheduling, and is not an end-to-end TPS. +""" + +import argparse +import hashlib +import importlib.util +import json +from pathlib import Path +import statistics +import sys + +import torch +import triton.testing + + +def load(name, path): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +def measure(before, after, rounds): + # Compile/warm both implementations before measuring; alternate AB/BA. + before() + after() + torch.cuda.synchronize() + values = {"before": [], "after": []} + for index in range(rounds): + order = (("before", before), ("after", after)) + for key, function in order if index % 2 == 0 else reversed(order): + values[key].append(triton.testing.do_bench(function, warmup=10, rep=40, quantiles=[0.5, 0.2, 0.8])[0]) + old, new = (statistics.median(values[key]) for key in ("before", "after")) + return { + "before_us": old * 1000, + "after_us": new * 1000, + "speedup": old / new, + "latency_change_percent": (new / old - 1) * 100, + "before_rounds_ms": values["before"], + "after_rounds_ms": values["after"], + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--baseline", type=Path, required=True) + parser.add_argument( + "--candidate", + type=Path, + default=Path(__file__).resolve().parents[1] / "lightllm/models/glm5_next/triton_kernel/kpool.py", + ) + parser.add_argument("--rounds", type=int, default=3) + args = parser.parse_args() + if args.rounds < 1: + parser.error("rounds must be positive") + before = load("kpool_baseline_benchmark", args.baseline) + after = load("kpool_candidate_benchmark", args.candidate) + torch.manual_seed(31) + storage = torch.empty((262144, 1, 584), device="cuda", dtype=torch.bfloat16) + packed = storage.view(torch.uint8)[:, :, -132:] + packed[:, 0, :128] = torch.randn(262144, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + packed[:, 0, 128:] = torch.ones(262144, device="cuda").view(torch.uint8).view(-1, 4) + rows = [] + for pools in (640, 32768, 262144): + table = torch.full((8, pools * 4), -1, device="cuda", dtype=torch.int32) + table[:, 3::4] = torch.randint(0, 262144, (8, pools), device="cuda", dtype=torch.int32) + for batch in (1, 4): + req = torch.arange(batch, device="cuda", dtype=torch.int32) + lengths = torch.full((batch,), pools * 4, device="cuda", dtype=torch.int32) + rows.append( + dict( + operation="prefill_gather", + max_pools=pools, + batch=batch, + **measure( + lambda: before.gather_pools(packed, table, req, lengths, pools), + lambda: after.gather_pools(packed, table, req, lengths, pools), + args.rounds, + ), + ) + ) + for layout in ("shared_long", "mixed_long", "short_at_capacity", "empty_at_capacity"): + batch = 8 + ids = [1] * 4 + [2] * 4 if layout == "shared_long" else list(range(batch)) + req = torch.tensor(ids, device="cuda", dtype=torch.int32) + counts = [max(0, pools - i % 4) for i in range(batch)] + if layout == "short_at_capacity": + counts = [min(pools, 64 + i) for i in range(batch)] + elif layout == "empty_at_capacity": + counts = [0] * batch + lengths = torch.tensor(counts, device="cuda", dtype=torch.int32) + rows.append( + dict( + operation="decode_paged_gather", + layout=layout, + max_pools=pools, + batch=batch, + **measure( + lambda: before.gather_paged_pools(packed, table, req, lengths, pools), + lambda: after.gather_paged_pools(packed, table, req, lengths, pools), + args.rounds, + ), + ) + ) + report = { + "device": torch.cuda.get_device_name(), + "torch": torch.__version__, + "rounds": args.rounds, + "baseline_sha256": hashlib.sha256(args.baseline.read_bytes()).hexdigest(), + "candidate_sha256": hashlib.sha256(args.candidate.read_bytes()).hexdigest(), + "results": rows, + "scope": ( + "Synthetic gather wrapper latency including allocations/page-table initialization; " + "compilation excluded. Not model throughput or TTFT." + ), + } + print(json.dumps(report)) + + +if __name__ == "__main__": + main() diff --git a/unit_tests/models/glm5_next/test_kpool_gather.py b/unit_tests/models/glm5_next/test_kpool_gather.py new file mode 100644 index 0000000000..3c5189679a --- /dev/null +++ b/unit_tests/models/glm5_next/test_kpool_gather.py @@ -0,0 +1,124 @@ +"""Byte-exact gather and graph-replay coverage for GLM K-pool staging.""" + +import pytest +import torch +import triton + +from lightllm.models.glm5_next.triton_kernel.kpool import gather_paged_pools, gather_pools + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") + + +def _inputs(max_pools, requests=8): + # The packed view has both a nonzero byte offset and a larger row stride. + storage = torch.empty((128, 1, 584), dtype=torch.bfloat16, device="cuda") + packed = storage.view(torch.uint8)[:, :, -132:] + keys = torch.randn(128, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + scales = torch.pow(2.0, torch.randint(-4, 3, (128,), device="cuda")).float() + packed[:, 0, :128] = keys + packed[:, 0, 128:] = scales.view(torch.uint8).view(128, 4) + table_storage = torch.full((requests, max_pools * 4 + 16), -1, device="cuda", dtype=torch.int32) + table = table_storage[:, 4:-12] + locations = torch.randint(0, 128, (requests, max_pools), device="cuda", dtype=torch.int32) + table[:, 3::4] = locations + return packed, table, locations, keys, scales + + +@pytest.mark.parametrize("pools", [17, 33, 640, 65536, 262144]) +def test_gather_pools_bytes_scales_and_changing_graph_inputs(pools): + torch.manual_seed(29) + packed, table, locations, source_keys, source_scales = _inputs(pools) + req = torch.tensor([3, 1, 0, 2], device="cuda", dtype=torch.int32) + lengths = torch.zeros(4, device="cuda", dtype=torch.int32) + gather_pools(packed, table, req, lengths, pools) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + keys, scales = gather_pools(packed, table, req, lengths, pools) + for values in ([0, 3, pools * 4 - 3, pools * 4], [pools * 4, 0, 7, 65], [2, 2, 2, 2]): + lengths.copy_(torch.tensor(values, device="cuda", dtype=torch.int32).clamp_max(pools * 4)) + req.copy_(req.roll(1)) + graph.replay() + valid = torch.arange(pools, device="cuda")[None, :] < (lengths // 4)[:, None] + loc = locations[req.long()].long() + expected_keys = source_keys[loc].masked_fill(~valid[:, :, None], 0) + expected_scales = source_scales[loc].masked_fill(~valid, 1.0) + torch.testing.assert_close(keys.view(torch.uint8).view(4, pools, 128), expected_keys, rtol=0, atol=0) + torch.testing.assert_close(scales.view(4, pools), expected_scales, rtol=0, atol=0) + + +def _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales): + page_bytes = pages.view(pages.shape[0], -1) + for row, (slot, count) in enumerate(zip(req.tolist(), lengths.tolist())): + # Shared pages may contain the owner's later keys. Only the original + # per-query causal length is readable by the indexer. + page_ids = block_table[row, : triton.cdiv(count, 64)].long() + actual_keys = page_bytes[page_ids, : 64 * 128].reshape(-1, 128)[:count] + actual_scales = page_bytes[page_ids, 64 * 128 :].contiguous().view(torch.float32).flatten()[:count] + loc = locations[slot, :count].long() + torch.testing.assert_close(actual_keys, source_keys[loc], rtol=0, atol=0) + torch.testing.assert_close(actual_scales, source_scales[loc], rtol=0, atol=0) + + +@pytest.mark.parametrize("max_pools", [640, 4096, 262144]) +@pytest.mark.parametrize("reuse_group", [1, 2, 4, 8]) +def test_paged_reuse_mixed_requests_causal_lengths_and_graph_slot_recycling(max_pools, reuse_group): + torch.manual_seed(30) + packed, table, locations, source_keys, source_scales = _inputs(max_pools) + batch = reuse_group * 2 + 1 + shared = [1] * reuse_group + [2 + i % 2 for i in range(reuse_group)] + [4] + mixed = [1 + i % 3 for i in range(batch)] + req = torch.tensor(shared, device="cuda", dtype=torch.int32) + lengths = torch.tensor([max(0, max_pools - i * 63) for i in range(batch)], device="cuda", dtype=torch.int32) + gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + pages, block_table = gather_paged_pools(packed, table, req, lengths, max_pools, reuse_group=reuse_group) + identity = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) + scenarios = [ + (shared, [max(0, max_pools - i * 63) for i in range(batch)]), + (mixed, [max_pools if i % 2 else min(65, max_pools) for i in range(batch)]), + (shared, [min(i * 17, max_pools) for i in range(batch)]), + ([0] * batch, [0] * batch), + ([5] * reuse_group + mixed[reuse_group:], [max_pools] * batch), + ] + for slots, counts in scenarios: + req.copy_(torch.tensor(slots, device="cuda", dtype=torch.int32)) + lengths.copy_(torch.tensor(counts, device="cuda", dtype=torch.int32)) + # Recycle physical KV locations without changing any captured pointer. + locations.copy_(locations.roll(1, dims=1)) + table[:, 3::4] = locations + original_lengths = lengths.clone() + if not any(counts): + pages.fill_(127) + graph.replay() + torch.testing.assert_close(lengths, original_lengths, rtol=0, atol=0) + _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) + if max(counts) < 4096 or reuse_group == 1: + torch.testing.assert_close(block_table, identity, rtol=0, atol=0) + elif slots == shared and reuse_group > 1: + valid_pages = triton.cdiv(max(counts[:reuse_group]), 64) + for row in range(1, reuse_group): + torch.testing.assert_close(block_table[row, :valid_pages], block_table[0, :valid_pages]) + if not any(counts): + dummy = pages.view(pages.shape[0], -1)[block_table[:, 0].long()] + assert not dummy[:, : 64 * 128].any() + assert (dummy[:, 64 * 128 :].contiguous().view(torch.float32) == 1).all() + + +@pytest.mark.parametrize("batch", [1, 3, 5]) +def test_paged_partial_groups_keep_valid_per_query_pages(batch): + packed, table, locations, source_keys, source_scales = _inputs(4097) + req = torch.ones(batch, device="cuda", dtype=torch.int32) + lengths = torch.tensor([4097 - i for i in range(batch)], device="cuda", dtype=torch.int32) + pages, block_table = gather_paged_pools(packed, table, req, lengths, 4097, reuse_group=4) + _check_paged(pages, block_table, req, lengths, locations, source_keys, source_scales) + if batch < 4: + expected = torch.arange(block_table.numel(), device="cuda", dtype=torch.int32).view_as(block_table) + torch.testing.assert_close(block_table, expected, rtol=0, atol=0) + + +@pytest.mark.parametrize("reuse_group", [0, 3, 16]) +def test_reject_unsupported_reuse_group_before_allocating(reuse_group): + with pytest.raises(ValueError, match="reuse_group"): + gather_paged_pools(None, None, None, None, 4096, reuse_group=reuse_group)