sparse_mla_prefill_h64_d512_k640

Sparse Multi-head Latent Attention for prefill, with BF16 Q and a BF16 latent workspace. K and V are the same latent vector. Each query has its own sparse index list and length; visibility is fully specified by these inputs, with no additional causal mask. One softmax includes all valid selected slots and a per-head zero-value sink. Queries with no selected keys return zero. Cache dequantization, index selection, Q normalization/rotary embedding and output projection are outside this operator.

This kernel task packages the flashmla_prefill_k640 baseline, a complete captured workload subset, and an independent-reference correctness and timing judge. The public callable is run with the definition input order.

Computation

For each query, gather the valid latent vectors specified by indices and topk_length. K and V are the same latent vector. Compute scaled softmax attention with an additional per-head zero-value sink logit in the denominator. Inputs use a BF16 workspace; dequantization and sparse index selection are outside the operator.

Fixed dimensions: H=64, D=512, KVH=1, TopK=640. The exact reference, input constraints and output types are in definition.json.

Inputs

Outputs

Implementation and origins

FlashMLA Blackwell BF16 sparse prefill with 640 index slots, per-query sparse lengths and attention sinks. Upstream: https://github.com/deepseek-ai/FlashMLA/tree/0f11e178da1f67becb087a324b400854931ec3e3 (MIT).

The embedded baseline solution is extracted byte-for-byte into torch-ext/sparse_mla_prefill_h64_d512_k640/. The ready-to-load noarch CUDA variant lives in build/torch-cuda/. The wrapper imports the preserved backend implementation when run is called.

Captured checkpoint:

• deepseek-ai/DeepSeek-V4.1-Flash — revision 5254448103b6254fc646dacbca1cbbf20e9216c1

Source dataset: humanfia-lab/inference-kernels, GitHub source commit 008b39387b852875c62e2aa8b0bc7f2c6232fea1. Captures use the pinned SGLang build 684fea8af3b11a5b51ec9966c02959a32ad3421c on B300 hardware.

Workload generation

Input lengths and shared-prefix structure were selected from the pinned AgentX trace semianalysisai/cc-traces-weka-062126. Original prompt text is unavailable: an authored coding conversation is tokenized and repeated in deterministic 64-token spans. These are real full-model inference activations for the replacement text, not original production prompt activations. Cached cases execute their predecessor before the target request. Concurrent groups are submitted together without replaying original arrival times. Generation uses temperature zero and the recorded output length. Capture uses 4096-token prefill chunks and selected early/middle/final decode steps. Layer coverage is concentrated in the first matching layers; see each capture record for available layer IDs. Qwen GDN prefill sequences are separated from the packed engine call; DeepSeek shards are gathered into complete logical operators. Tensor strides are contiguous, and KV snapshots retain all referenced pages and page tails while remapping physical IDs. Weights retain the full model dimensions and all experts. Text-only router captures do not cover image-token routing.

The trace folder retains complete workload inputs and the independent definition. Capture provenance and replay records are available in the source dataset. Paths are relative to flashinfer_trace/; all snapshots are raw safetensors under blobs/.

Workload statistics

• Captured workloads: 20.
• Distinct input sets: 20.
• Exact runtime shapes: 20.
• Tensor snapshots: 94, 11.126 GB materialized.

Each bar counts complete matching runtime shapes. These counts are not serving invocation frequencies.

Install and load

Install the appropriate CUDA-enabled Torch build and the pinned baseline dependencies:

python -m pip install -r requirements.txt
git lfs pull

from kernels import get_kernel
kernel = get_kernel("humanfia-lab/dev-sparse-mla-prefill-h64-d512-k640", revision="main", trust_remote_code=True)
# kernel.run(...) accepts the inputs listed above.

The project follows the HF kernel source layout with build.toml, a pinned flake.nix, torch-ext/, a noarch build variant, and benchmarks/benchmark.py. The runtime backend packages remain explicit requirements because the current implementation is the original baseline.

SGLang baseline and task acceptance

The SGLang baseline and task config come from the same source revision. Config paths are adapted to this standalone layout; tolerances, development UUIDs and timing values are unchanged. SGLang adapters for layouts, metadata, rank-local calls and output order remain explicit. A standalone correctness result does not establish an in-engine replacement.

Run the SGLang implementation in the pinned capture image lmsysorg/sglang@sha256:ddae0e399d3d94f7827f16183301cf8dac4cf0eae0ef25590b7dac5881d713d8, with the kernel SDK installed; requirements-sglang.txt lists its source dependency.

python benchmarks/benchmark.py --implementation sglang --dev --validation-only
python benchmarks/benchmark.py --implementation /path/to/candidate.py:run --compare-baseline

Baseline comparison defaults to SGLang; --comparison-baseline baseline selects the packaged open-source implementation. --dev selects at most ten source development cases; final evaluation without it uses every workload. Warmup, iterations and trials default to the task config. SGLang prepare runs once per case and only launch is timed.

Judge an implementation

flashinfer_trace/task.json is the evaluation entry point. Its paths are relative to the repository root. dev contains UUIDs of existing workloads for quick iteration; the default full evaluation uses every workload. Capture provenance and replay records are available in the source dataset.

benchmarks/benchmark.py reads the local trace; no model checkpoint or serving engine is required.

# Packaged baseline, every captured workload
python benchmarks/benchmark.py --report results/baseline.jsonl

# Local candidate with the same run signature, compared to the baseline
python benchmarks/benchmark.py --implementation /path/to/candidate.py:run \
  --compare-baseline --report results/candidate.jsonl

# Inspect UUIDs and shapes, or run correctness only
python benchmarks/benchmark.py --list
python benchmarks/benchmark.py --validation-only --limit 1

A candidate can also be a local HF kernel repository directory or org/repo@revision. --workload UUID is repeatable. Optional --gpu GPU-UUID --lock-directory /path/to/locks selects and exclusively locks a shared GPU.

The judge freezes the complete independent reference before candidate calls, evaluates all elements under the applicable tolerance, and rejects input mutation and stale fixed-output/data-pointer caches using a valid alternate input at the same addresses. The task config explicitly selects the definition, workloads, baseline, tolerances, development UUIDs and timing. Floating-point acceptance uses its tolerance policy; integer outputs require exact equality. All outputs must be finite. Outputs must retain their specified shape, dtype and device and contain finite floating values.

Timing uses CUDA events over eager calls. SGLang baselines prepare once per workload and time only launch; other implementations time the complete call. Loading, reference execution and checking are excluded. Warmup and iteration counts are configurable. Timing results are checked again; failures exit nonzero. With --compare-baseline, speedup uses a fresh baseline measurement on the same device and inputs. These functional checks are not an adversarial security sandbox.

The official HF benchmark CLI discovers this file and its TraceBenchmark subclass:

kernels benchmark . --warmup 3 --iterations 20

That adapter runs strict reference/dependency checks in workload setup before timing each UUID. Use the direct script above for per-output errors, post-timing checks, failure summaries and baseline speedups; the HF CLI does not provide this task-specific judge report.

Coverage and validation limits

• This is a coverage sample without production frequency weights. Duplicate input sets can occur; split evaluation by request and input signature.
• Fixed dimensions specialize the operator. Capture adaptations and engine-output replay records are documented in the source dataset.
• The baseline source is preserved, with external backend dependencies pinned in requirements.txt. This is a source/JIT baseline package, not a bundled native multi-ABI distribution; install those backends before calling run.
• The source dataset passes complete input validation. SGLang tasks retain engine-output replay records; imported GDN tasks have only reference shape/finiteness checks because the contest publishes no engine outputs. No serving-speedup claim is made here.

Licenses

See LICENSES.md and the retained upstream/model license texts in licenses/.