[Perf][GLM-5.3-Flash] Shard long-context indexer prefill rows across TP - #54951

TL;DR

How it works / why it is faster: Instead of repeating sparse-indexer prefill MQA scoring and top-k on every TP rank, partition independent query rows into cost-balanced contiguous slices. Each rank scores its rows against the full available key history; all_gatherv reassembles the final int32 indices. This removes redundant computation in exchange for fixed-width index communication, so savings grow with context length. There is no attention approximation. Decode kernels are unchanged; smaller batches retain the existing path (activation requires at least 16,384 × TP scheduled prefill query rows).

Performance, baseline → optimized: main parent 5893426b88f7 → PR head 16934c6e76e5; GLM-5.3-Flash on 4× H200, TP4, C1, 256 output tokens, 65,536-token batch budget, legacy runner + explicit FA3 backend. Four fresh-process A/B pairs in ABBA/BAAB order; warmups excluded. Throughput below counts input + output tokens at C1, not saturated serving throughput.

Server-side TPOT at 1M: 6.631 → 6.633 ms/token. Across the full 1K–1M scan, absolute TPOT percentage-change point estimates are <0.04%; every paired 95% interval includes zero. This does not prove equivalence. The 2K TTFT point estimate is 2.84% slower (interval crosses zero); we do not claim blanket short-input non-regression.

Correctness / targeted accuracy, baseline → optimized: Across 33 distinct retrieval probes spanning 1K–1M, repeated over four processes per arm, expected-answer presence is 132/132 → 132/132; strict answer-only accuracy is 127/132 → 129/132. Exact generated-token agreement is 130/132 A/B (baseline A/A: 98/99). Both observed A/B differences contain the same correct code, with an additional explanatory sentence in baseline. These are targeted retrieval checks, not a broad quality benchmark or a claim of bitwise equivalence. The 54 focused PR tests also pass.

Full eleven-length results, paired confidence intervals, implementation scope, and reproducibility details follow.

Summary

Shard replicated sparse-indexer prefill query rows across tensor-parallel ranks. Each rank computes MQA logits and top-k for a cost-balanced contiguous slice, then one variable-size all_gatherv call per indexer layer restores the original row layout. This removes redundant scoring; it does not approximate attention or merge independent top-k candidates. The API call is not a claim of one underlying physical NCCL operation.

For GLM-5.3-Flash, the exchange includes the incomplete k-pool tail: index_topk + index_kpool - 1 = 2048 + 4 - 1 = 2,051 int32 columns. KV gathering remains unconditional because subsequent query chunks can reuse that workspace.

The avoided MQA scoring work grows with both query rows and the available pooled-key history, whereas the exchanged index width per query row stays fixed. Longer histories therefore offer more computation to save per byte exchanged. The cost-balanced partition also accounts for different causal-history lengths across rows. This targets prefill TTFT and prefill-dominated total-token throughput; the row partition and exchange are outside the decode branch, so a direct decode-kernel speedup is not expected. Indirect TPOT changes are measured and reported below.

Scope and activation

The diff against main is four files: the shared DSA indexer metadata/planner, the generic sparse indexer, the GLM k-pool indexer, and their focused tests. #53906 is already in the baseline; this PR does not reintroduce its model support or kernel changes.

• Requires CUDA, TP > 1, an MQA prefill, and at least 16,384 × TP scheduled prefill query rows in the current batch.
• DCP/PCP, non-CUDA platforms, disabled PyNCCL, NCCL symmetric-memory, batch-invariant mode, and configurations whose mixed CUDA-graph mode is FULL retain the existing path.
• Decode kernels, model weights, sampling, KV-cache format, and attention's mathematical definition are not changed. TPOT is nevertheless measured to check for indirect effects.
• The shared generic DSA path is also modified. H200 serving results below cover GLM-5.3-Flash only; they are not an end-to-end no-regression guarantee for other models, platforms, concurrency levels, or speculative decoding.

Important configuration constraint: at TP4, max_num_batched_tokens=8192 cannot activate this implementation, even with a 1M-token accumulated context. Both arms below use 65,536, so C1 inputs of 1K–32K exercise the fallback, and full 64K prefill chunks can exercise sharding. The gate is a batch-row count, not the request's context length; several shorter concurrent requests can also reach it.

Exact-head A/B rescan

These measurements replace the previous historical tables. Those tables used older PR/runtime overlays and an 8,192-token batch budget; their gains cannot be attributed to the current head and gate.

Both arms load their respective clean source trees, the same model snapshot (3f1971b7b5f7a528c9c4ef6212c8785298a8c24a), and identical compiled vLLM extensions from the official wheel for the baseline commit. Worker-level provenance checks verify the loaded vLLM paths for all four TP ranks. The historical vendor-vLLM Python overlay is not used. The PR range contains one commit.

Configuration and stability protocol

• 4 × NVIDIA H200, TP4, native FP8 weights, BF16 KV cache; text-only, legacy model runner (VLLM_USE_V2_MODEL_RUNNER=0) and explicit FLASH_ATTN_MLA_SPARSE attention backend, chunked prefill, max_model_len=1,048,576, max_num_batched_tokens=65,536, max_num_seqs=8, GPU-memory utilization 0.90.
• MTP/speculative decoding off; prefix caching off; FULL_AND_PIECEWISE CUDA-graph configuration with main's default breakable-graph behavior. The effective engine configuration reports CompilationMode.NONE and CUDA-graph capture sizes [1, 2, 4, 8, 16]; this is not a torch.compile-enabled prefill benchmark. Torch 2.13.0+cu130, Triton 3.7.1, FlashInfer 0.6.18, Transformers 5.12.1. Same process environment/cache policy in both arms; OMP/MKL threads fixed at 1. Synchronous CUDA debugging is off for timed runs. These explicit runner/backend settings avoid failures observed in the unmodified baseline; see limitations below.
• C1 sequential requests. Eight fresh serving processes in ABBA / BAAB order form four adjacent baseline/optimized restart pairs. Each process scans all eleven lengths, with three excluded warmups and eight measured requests per length. This is 32 measured requests per arm per length, not 32 independent restarts.
• Same exact-token prompts, seeds, and order in both arms; temperature 0, top-p 1, top-k 1, fixed 256 output tokens, ignore_eos=true. 1M means 1,000,000 input tokens, leaving room for generation within the server limit. Other K labels are powers of 1,024.
• GPU utilization, memory, clocks, temperature, and power are recorded at 1 Hz. Startup/JIT warmup and separate diagnostic instrumentation are excluded from measured samples. No outlier samples are silently removed.
• TTFT/TPOT use the server's per-request timestamps. TTFT starts at scheduling, excluding upload and queue time. TPOT = (last output timestamp − first output timestamp) / 255. Client-side latency is retained as a transport cross-check; SSE chunk arrival intervals are not substituted for decode time.
• Latencies below are geometric means of the four per-restart request medians. Throughput is total input + output tokens divided by the observed C1 batch interval (including client/request overhead), geometrically averaged across restarts. It is not saturated multi-client serving throughput. At a fixed input/output length, output-token throughput has the same percentage change.
• Positive improvement means lower latency (1 − optimized/baseline) or higher throughput (optimized/baseline − 1). Confidence intervals resample the four restart pairs, not individual requests (20,000 bootstrap draws). Four pairs give limited statistical resolution; an interval crossing zero is not proof of equivalence.

Performance

At 1M, measured TTFT improves by +25.42% and C1 total-token throughput by +32.59%. Across all lengths, TPOT improvement point estimates range from -0.03% to +0.01%. The fallback 1K–32K TTFT changes range from -2.84% to +0.43%. The intervals below and individual restart results, not just the sign of a small point estimate, determine how confidently an improvement or regression can be distinguished from run-to-run variation. These results do not prove zero overhead outside the tested configuration. Every TPOT interval includes zero: this scan does not establish a systematic TPOT change, nor does it prove exact equivalence.

Short-input caveat: the most negative fallback TTFT point is 2K: 151.39 → 155.69 ms (-2.84% improvement; paired 95% interval [-8.18, +2.24]%). This observation is retained, not discarded or declared harmless solely because the new collective is gated off. The scan does not establish blanket short-input non-regression or a causal explanation for this timing difference.

Correctness / targeted accuracy A/B

This is a targeted long-context retrieval and numerical-equivalence evaluation, not a general model-quality benchmark. Three distinct generated record probes per length place the answer at 10%, 50%, and 90% of the prompt. Each is repeated in the four independent processes per arm: 12 observations per arm per length, but only three distinct tasks. The prompt has a non-thinking assistant boundary; greedy decoding stops normally with at most 64 output tokens.

The table reports strict expected-answer accuracy, whether the expected answer is present, and exact generated-token equality for matched prompts. Answer containment alone is a weaker check: a response can contain the correct code yet fail strict answer-only formatting by adding an explanatory sentence. Baseline A/A compares adjacent independent baseline restarts, providing a control for process-to-process repeatability. HTTP success and valid token counts alone are not treated as accuracy.

Matched A/B output-token sequences agree in 130/132 comparisons; baseline A/A agrees in 98/99. The expected answer is present in 132/132 baseline and 132/132 optimized responses. Strict answer-only format passes 127/132 baseline and 129/132 optimized responses. Answer containment does not excuse extra or conflicting text; strict answer-only formatting and full token equality are reported separately. Repeated observations of the same three tasks are not independent evidence of broad model accuracy. Multimodal inputs, MTP, other DSA models, and general reasoning quality have not been evaluated by this scan.

Inspection of both A/B mismatches (the 90%-depth task at 512K and 1M) shows the same correct access code in both responses: baseline adds an explanatory sentence, while optimized returns only the code. The baseline A/A mismatch at 1M has the same formatting distinction; optimized's own restarts also differ in formatting at 512K. No different retrieved code was observed in these probes. This is an inspection of the observed mismatches, not a claim that all possible numerical differences are formatting-only.

Additional checks:

• 54 focused PR tests pass: row-cost balancing, gate predicates, metadata construction, and gather/reassembly coverage including 128K–1M. These tests include mocked collective/metadata checks; they are not a substitute for the distributed serving run.
• A separate instrumented TP4 pilot confirmed actual 64K/128K sharded execution and the [65,536, 2,051] gathered result. Its timings are excluded from the performance table.
• Real H200 DeepGEMM + CUDA top-k checks at nine lengths from 4K to 1M compare a 512-query-row batch against uneven row slices. Valid logits are exactly equal and selected top-k sets match. The top-k kernel's output ordering is not deterministic: different orders also occur in the baseline A/A control with identical sets. These checks do not establish bitwise equivalence of all attention outputs.
• Real Triton pool expansion plus incomplete-tail append agrees exactly with its direct reference at those nine lengths. Measurement-harness tests validate empty-text token accounting, bundled SSE responses, restart-level aggregation, and rejection of unmatched prompts or inconsistent token counts, fingerprints, and timing denominators.

Runtime qualification / limitations: all 704 measured performance requests complete with exactly 256 output tokens; all 264 separate retrieval requests complete. This is not an error-free-runtime claim: the common FA3 runtime emits TMA-descriptor diagnostics, which are preserved in the server logs. Two separate baseline failures were investigated before fixing the common benchmark configuration: (1) the automatically selected FlashInfer SM90 sparse MLA backend fails during 65,536-row startup/autotune; a standalone NoPE MLA probe also fails at 65,536 rows while passing 64 rows; (2) with FA3, the V2 runner fails at 256K in its slot-mapping kernel, reproduced with eager execution and CUDA_LAUNCH_BLOCKING=1. The legacy-runner/FA3 diagnostic then passes 256K, 512K, and 1M. Failed and synchronous diagnostic admissions are excluded from the performance table. No unrelated baseline source fix is overlaid. The common user environment also registers an unrelated SGLang operator at Python startup; neither pinned vLLM source tree has a callsite for that operator. These results qualify only the explicit legacy-runner/FA3 configuration, not default V2 or auto-selected FlashInfer SM90 MLA.

Reproduction

Use clean checkouts of the two pinned commits and one matched environment. Pin the compiled artifacts to the parent-main wheel (rather than allowing each editable installation to choose a different nightly):

export VLLM_USE_PRECOMPILED=1
export VLLM_PRECOMPILED_WHEEL_LOCATION='https://wheels.vllm.ai/5893426b88f7b3cd21101d194eb1c6f0a6f0e27b/vllm-0.28.1rc1.dev478%2Bg5893426b8-cp38-abi3-manylinux_2_28_x86_64.whl'
# In the selected arm's worktree, with common dependencies already provisioned:
uv pip install --no-deps -e .

Wheel SHA256: 93097b58f56b838b3d18928753e00b5e18624fb9084d14928e748e5fdc176dc3. The measurements used the same extracted wheel extensions in both source worktrees and checked the four workers' module paths and extension hashes. --no-deps above assumes the explicitly pinned common runtime is already installed; it is not a complete environment bootstrap.

Set MODEL_DIR to the local model snapshot, and PYTHON to the common environment's Python. Start only one arm at a time; wait for health before running requests:

export CUDA_VISIBLE_DEVICES=4,5,6,7
export VLLM_USE_V2_MODEL_RUNNER=0 CUDA_LAUNCH_BLOCKING=0
export OMP_NUM_THREADS=1 MKL_NUM_THREADS=1 TOKENIZERS_PARALLELISM=false
export VLLM_NO_USAGE_STATS=1
# Use the selected source tree's bundled DeepGEMM, not an unrelated installed copy.
export PYTHONPATH="${PWD}/vllm/third_party:${PWD}${PYTHONPATH:+:${PYTHONPATH}}"
"$PYTHON" -m vllm.entrypoints.cli.main serve "$MODEL_DIR" \
  --host 127.0.0.1 --port 8010 --served-model-name zai-org/GLM-5.3-Flash \
  --tensor-parallel-size 4 --kv-cache-dtype bfloat16 \
  --max-model-len 1048576 --max-num-batched-tokens 65536 --max-num-seqs 8 \
  --gpu-memory-utilization 0.90 --safetensors-load-strategy prefetch \
  --enable-chunked-prefill --no-enable-prefix-caching --language-model-only \
  --attention-backend FLASH_ATTN_MLA_SPARSE \
  --tool-call-parser glm47 --reasoning-parser glm45 --enable-auto-tool-choice \
  --enable-per-request-metrics \
  --compilation-config '{"cudagraph_mode":"FULL_AND_PIECEWISE"}'

Use fresh process tags 00-baseline, 01-optimized, 02-optimized, 03-baseline, 04-optimized, 05-baseline, 06-baseline, 07-optimized. For each, run performance and then correctness with the client below; stop that owned server before starting the next. Keep the same shared compilation caches and do not include server startup in timing.

# ARM is baseline or optimized; TAG identifies this independent restart.
"$PYTHON" pr54951_measure.py --arm "$ARM" --tag "$TAG-performance" \
  --kind performance --warmups 3 --requests 8 --output-len 256 --output results
"$PYTHON" pr54951_measure.py --arm "$ARM" --tag "$TAG-correctness" \
  --kind correctness --warmups 0 --requests 3 --output-len 64 --output results
# Focused tests, run in the optimized source worktree:
"$PYTHON" -m pytest -q tests/v1/attention/test_indexer_tp_row_shard.py

Pair adjacent fresh processes and aggregate as specified above. Do not merge any diagnostic pilot or failed baseline admission into these samples. The source trees are unchanged by the benchmark; benchmark helpers are not additional PR commits.

"""Exact-token serving measurements; raw samples are never silently discarded."""

from __future__ import annotations

import argparse
import os
import asyncio
import hashlib
import json
import statistics
import time
from pathlib import Path

import aiohttp
import numpy as np
from transformers import AutoTokenizer

ROOT = Path(__file__).resolve().parent
MODEL = Path(os.environ['MODEL_DIR'])
CONTEXTS = [1024, 2048, 4096, 8192, 16384, 32768, 65536,
            131072, 262144, 524288, 1000000]


def prompt_tokens(tokenizer, length: int, case: int, kind: str):
    """Construct exact lengths, preserving chat boundary and a complete needle.

    The synthetic retrieval probe is a smoke test, not a general quality eval.
    Needle positions rotate through 10%, 50%, and 90% of the input.
    Performance prompts use the same public, generated record workload.
    """
    seed = 54951 + case
    rng = np.random.default_rng(seed)
    key = f'record-{case:05d}-verification'
    answer = '-'.join(rng.choice(['amber', 'cedar', 'falcon', 'ivory',
                                  'maple', 'orbit', 'pearl', 'violet'], 4))
    prefix = tokenizer.encode(
        '[gMASK]<sop><|system|>You retrieve exact information from records. '
        'Answer only the requested access code, without explanation.'
        '<|user|>Read these archived records.\n', add_special_tokens=False)
    suffix = tokenizer.encode(
        f'\nEnd of records. What is the access code for {key}? '
        'Respond with only the access code.'
        '<|assistant|>\n<think>\n</think>\n', add_special_tokens=False)
    needle = tokenizer.encode(
        f'\nThe access code for {key} is {answer}.\n', add_special_tokens=False)
    # Generate ordinary-text background once; no special/random vocabulary IDs.
    paragraphs = []
    for i in range(512):
        number = int(rng.integers(100000, 999999))
        paragraphs.append(f'Archive entry {i}: shipment {number} was received '
                          'at the northern depot. The inspection confirmed '
                          'that the supplies were dry and accounted for.\n')
    background = tokenizer.encode(''.join(paragraphs), add_special_tokens=False)
    available = length - len(prefix) - len(suffix) - len(needle)
    assert available > 0
    background = (background * ((available // len(background)) + 1))[:available]
    depth = [0.1, 0.5, 0.9][case % 3]
    position = int(available * depth)
    tokens = prefix + background[:position] + needle + background[position:] + suffix
    assert len(tokens) == length
    digest = hashlib.sha256(np.asarray(tokens, dtype='<i4').tobytes()).hexdigest()
    return tokens, {'sha256': digest, 'case': case, 'seed': seed, 'kind': kind,
                    'needle_depth': depth, 'expected_answer': answer}


async def request(session, base_url, tokens, output_len, metadata, quality=False):
    body = {
        'model': 'zai-org/GLM-5.3-Flash', 'prompt': tokens,
        'max_tokens': output_len, 'temperature': 0, 'top_p': 1, 'top_k': 1,
        'seed': 54951, 'ignore_eos': not quality, 'stream': True,
        'stream_options': {'include_usage': True},
        # Streaming logprobs identify empty-text tokens without echoing a
        # million prompt IDs in the first response chunk.
        'logprobs': 5 if quality else 0, 'return_tokens_as_token_ids': True,
        'skip_special_tokens': False,
    }
    # Serialize outside the timed interval. HTTP upload remains in client TTFT.
    payload = json.dumps(body, separators=(',', ':')).encode()
    record = {**metadata, 'input_tokens': len(tokens), 'requested_output': output_len,
              'started_unix': time.time(), 'chunks': [], 'output_token_ids': [],
              'text': '', 'logprobs': [], 'server_metrics': None}
    started = time.perf_counter()
    async with session.post(base_url + '/v1/completions', data=payload,
                            headers={'Content-Type': 'application/json'}) as response:
        if response.status != 200:
            raise RuntimeError(f'HTTP {response.status}: {(await response.text())[:2000]}')
        async for raw_line in response.content:
            received = time.perf_counter() - started
            line = raw_line.decode().strip()
            if not line.startswith('data: '):
                continue
            if line == 'data: [DONE]':
                break
            event = json.loads(line[6:])
            if 'error' in event:
                raise RuntimeError(event['error'])
            if event.get('metrics'):
                record['server_metrics'] = event['metrics']
            if event.get('usage'):
                record['usage'] = event['usage']
            for choice in event.get('choices', []):
                logprobs = choice.get('logprobs') or {}
                ids = [int(t.removeprefix('token_id:'))
                       for t in logprobs.get('tokens', [])]
                if ids:
                    record['chunks'].append({'t_s': received, 'n': len(ids)})
                    record['output_token_ids'].extend(ids)
                    record['logprobs'].extend(logprobs.get('token_logprobs', []))
                record['text'] += choice.get('text', '')
                if choice.get('finish_reason'):
                    record['finish_reason'] = choice['finish_reason']
    record['client_elapsed_s'] = time.perf_counter() - started
    chunks = record['chunks']
    count = len(record['output_token_ids'])
    assert chunks and count > 0, 'No identifiable output tokens received'
    assert record.get('usage', {}).get('prompt_tokens') == len(tokens), record.get('usage')
    assert record['usage']['completion_tokens'] == count, record['usage']
    if not quality:
        assert count == output_len, (count, output_len)
    record['client_ttft_ms'] = chunks[0]['t_s'] * 1000
    record['client_tpot_ms'] = ((chunks[-1]['t_s'] - chunks[0]['t_s']) * 1000 /
                              (count - 1)) if count > 1 else None
    record['output_tokens'] = count
    record['token_sha256'] = hashlib.sha256(
        np.asarray(record['output_token_ids'], dtype='<i4').tobytes()).hexdigest()
    record['retrieval_contains_answer'] = metadata['expected_answer'] in record['text']
    # The forced non-thinking boundary makes this strict answer test meaningful.
    normalized = record['text'].replace('<|endoftext|>', '').replace('<|user|>', '').strip()
    record['retrieval_exact'] = normalized == metadata['expected_answer']
    return record


async def run(args):
    output = Path(args.output)
    output.mkdir(parents=True, exist_ok=True)
    path = output / f'{args.tag}.jsonl'
    if path.exists():
        raise FileExistsError(f'Refusing to overwrite {path}')
    tokenizer = AutoTokenizer.from_pretrained(str(MODEL), trust_remote_code=False)
    timeout = aiohttp.ClientTimeout(total=1800, connect=10)
    async with aiohttp.ClientSession(timeout=timeout, trust_env=False,
                                    read_bufsize=32 * 1024 * 1024) as session:
        async with session.get(args.base_url + '/health') as response:
            assert response.status == 200
        with path.open('x') as stream:
            for length in args.contexts:
                jobs = [('warmup', i) for i in range(args.warmups)]
                jobs += [('measured', i) for i in range(args.requests)]
                prepared = {}
                for phase, i in jobs:
                    case = 1000 + i if phase == 'warmup' else i
                    prepared[phase, i] = prompt_tokens(tokenizer, length, case, args.kind)
                batch_start = None
                batch_input = batch_output = 0
                for phase, i in jobs:
                    # Same measured prompts and order in every fresh process.
                    tokens, metadata = prepared[phase, i]
                    metadata.update(arm=args.arm, tag=args.tag, phase=phase,
                                    context=length, repetition=i)
                    if phase == 'measured' and batch_start is None:
                        batch_start = time.perf_counter()
                    record = await request(session, args.base_url, tokens,
                                           args.output_len, metadata,
                                           quality=args.kind == 'correctness')
                    if phase == 'measured':
                        batch_input += len(tokens)
                        batch_output += record['output_tokens']
                        batch_end = time.perf_counter()
                    stream.write(json.dumps(record, separators=(',', ':')) + '\n')
                    stream.flush()
                    metrics = record['server_metrics'] or {}
                    print(json.dumps({'tag': args.tag, 'context': length, 'phase': phase,
                                      'i': i, 'client_ttft': record['client_ttft_ms'],
                                      'server': metrics, 'tokens': record['output_tokens'],
                                      'answer': record['retrieval_contains_answer']}), flush=True)
                if batch_start is not None:
                    summary = {'record_type': 'batch_summary', 'context': length,
                               'arm': args.arm, 'tag': args.tag, 'kind': args.kind,
                               'duration_s': batch_end - batch_start,
                               'input_tokens': batch_input, 'output_tokens': batch_output,
                               'total_tok_s': (batch_input + batch_output) / (batch_end - batch_start),
                               'output_tok_s': batch_output / (batch_end - batch_start)}
                    stream.write(json.dumps(summary) + '\n')
                    stream.flush()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--arm', choices=['baseline', 'optimized'], required=True)
    parser.add_argument('--tag', required=True)
    parser.add_argument('--output', default=str(ROOT / 'results/pr54951-current-20260907'))
    parser.add_argument('--base-url', default='http://127.0.0.1:8010')
    parser.add_argument('--contexts', nargs='+', type=int, default=CONTEXTS)
    parser.add_argument('--warmups', type=int, default=3)
    parser.add_argument('--requests', type=int, default=8)
    parser.add_argument('--output-len', type=int, default=256)
    parser.add_argument('--kind', choices=['performance', 'correctness'], default='performance')
    asyncio.run(run(parser.parse_args()))

AI assistance and ownership

AI assistance was used for implementation and benchmark analysis. The author owns validation, review follow-up, and maintenance. This PR optimizes the already-merged GLM/DSA prefill path; it is not a duplicate model-support submission.