Deepseek_sparse_attention_topk_benchmark

# GLM-5.2 Agent DSA TopK benchmark inputs Real FP32 score inputs for benchmarking TopK kernels at long context lengths. The primary operator contract is: ```text input: N FP32 scores output: indices of the largest 2048 scores ``` Each file contains 256 independent query rows. `N` is the context length: 65,536, 131,072, 262,144, 524,288, or 992,757. Scores are captured from the GLM-5.2 DSA indexer immediately before SGLang's TopK-2048 transform. This data is intended for kernel latency/throughput benchmarking, tuning, and correctness comparison using model-distributed inputs instead of synthetic random scores. ## Collection status Completed batches: **10/10**. Uploaded records: **105/105**. | Batch | Context | Layers | Records | Status | |---:|---:|---|---:|---| | 1 | 65,536 | 0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34 | 11 | complete | | 2 | 65,536 | 38, 42, 46, 50, 54, 58, 62, 66, 70, 74 | 10 | complete | | 3 | 131,072 | 0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34 | 11 | complete | | 4 | 131,072 | 38, 42, 46, 50, 54, 58, 62, 66, 70, 74 | 10 | complete | | 5 | 262,144 | 0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34 | 11 | complete | | 6 | 262,144 | 38, 42, 46, 50, 54, 58, 62, 66, 70, 74 | 10 | complete | | 7 | 524,288 | 0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34 | 11 | complete | | 8 | 524,288 | 38, 42, 46, 50, 54, 58, 62, 66, 70, 74 | 10 | complete | | 9 | 992,757 | 0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34 | 11 | complete | | 10 | 992,757 | 38, 42, 46, 50, 54, 58, 62, 66, 70, 74 | 10 | complete | The 105 layer/context records are uploaded in 10 atomic batches. Each context point is split into an 11-layer batch and a 10-layer batch. A batch is complete only after every safetensors/JSON pair passes local schema, size, and SHA256 checks and the corresponding remote objects pass size and SHA256 verification. ## Data source - Base dataset: `aisa-group/PostTrainBench-Trajectories`, pinned at revision `9310b65ef86b385d4f12fd7fbfb8b2febdd0b0b8`. - Natural prefix: 424,370 tokens generated by `glm-x-preview[1m]` in the PostTrainBench agent environment. - Continuation: genuine `zai-org/GLM-5.2` agent/tool turns against a deterministic archive of authentic PostTrainBench tool data. - Final request: 992,757 tokens, 1,115 messages, and 78 continuation events. - Replay: official GLM-5.2-FP8 weights on one NVIDIA H20; index TopK is 2048. This is a transparently labelled hybrid trace, not an untouched production request. Historical reasoning, assistant messages, tool calls, and tool results are retained with `clear_thinking=False`; compaction count is zero. The pipeline does not remove chain-of-thought, summarize history, pad tokens, concatenate unrelated conversations, or generate output tokens during capture. ## Collection method 1. Reconstruct the complete conversation with the official GLM-5.2 tokenizer and chat template. 2. Replay the request in prefill-only mode with `max_new_tokens=0`. Routed expert weights are streamed layer by layer. 3. Hook the FP32 output of `deep_gemm.fp8_mqa_logits` immediately before the SGLang TopK transform. 4. Store 256 consecutive query rows ending at every configured context point. 5. Capture only the 21 `full` indexer layers. Shared layers reuse prior TopK indices and do not invoke a new TopK operator. ## Record format Each safetensors file contains: - `topk_input`: FP32 `[256, context_length]` raw indexer scores; - `valid_lengths`: int32 exclusive causal end position for each query row; - `row_starts`: int32 inclusive causal start position for each query row. The full-indexer layers are `0, 1, 2, 6, 10, 14, 18, 22, 26, 30, 34, 38, 42, 46, 50, 54, 58, 62, 66, 70, 74`. ## Layout ```text data/<trace_id>/L<context_length>/layer_<layer>.safetensors data/<trace_id>/L<context_length>/layer_<layer>.json provenance/<trace_id>/manifest.json provenance/<trace_id>/input_ids.safetensors provenance/<trace_id>/final_request.txt.gz manifest/progress.json ``` ## Main use The main use is to select any `n` real query rows whose score lengths are near one million, then run a TopK operation on each row: ```text input per query: N FP32 scores, where N is close to 1,000,000 output per query: indices of the largest 2048 scores ``` Use records under `L992757` for the near-1M workload. Each file provides 256 queries from one real model layer, so users can choose any desired `n` between 1 and 256. `row_starts` and `valid_lengths` give the actual score range for each query. The returned index is relative to the full context row. ## Example ```python import torch from safetensors import safe_open path = "data/<trace_id>/L992757/layer_74.safetensors" n = 8 with safe_open(path, framework="pt", device="cpu") as f: scores = f.get_slice("topk_input")[256 - n : 256] valid_lengths = f.get_tensor("valid_lengths")[256 - n : 256] row_starts = f.get_tensor("row_starts")[256 - n : 256] topk_indices = [] for row, start, end in zip(scores, row_starts, valid_lengths): start = int(start) end = int(end) indices = torch.topk( row[start:end], k=2048, sorted=False ).indices + start topk_indices.append(indices) # n outputs; every output contains 2048 indices. topk_indices = torch.stack(topk_indices) print(topk_indices.shape) # [n, 2048] ``` Replace `torch.topk` with the candidate TopK kernel for performance testing. The adjacent JSON file records SHA256, shape, trace hashes, source/model revisions, layer, and score-stage identifier. ## Provenance and limitations The source trajectories are Apache-2.0 licensed and remain available from `aisa-group/PostTrainBench-Trajectories`. GLM-5.2 weights are MIT licensed. This repository stores derived operator inputs and source hashes, not model outputs, Q/K tensors, or TopK outputs. The manifest records the prefix model, continuation model, source revision, request hashes, token boundary, and zero-compaction assertion. The hybrid trace must not be represented as an untouched production request.