prefix_cached_block_causal_h32_d128

Prefix-cached block-causal attention: allow j <= Np+i or matching nonnegative image IDs, then intersect key_valid. Text and different image blocks are causal; same-image tokens are bidirectional. Np=0 covers initial full evaluation. Prefix K/V are already computed; changed target-image K/V must be supplied as current inputs.

This kernel task packages the fa4_prefix_h32_d128 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

Concatenate cached and current K/V. Query $i$ sees key $j$ when $j\leq N_p+i$ or both tokens share a nonnegative image ID, intersected with key_valid. Apply scaled softmax attention over exactly those keys. Tokens within the same image block are bidirectional.

Fixed dimensions: H=32, D=128. The exact reference, input constraints and output types are in definition.json.

Inputs

Outputs

Implementation and origins

FlashAttention 4 with causal-or-same-image visibility and runtime key validity. Prefix concatenation is part of the solution call. Upstream: https://github.com/Dao-AILab/flash-attention (BSD-3-Clause).

The embedded baseline solution is extracted byte-for-byte into torch-ext/prefix_cached_block_causal_h32_d128/. 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:

• Qwen/Qwen-Image-2.1-Turbo — revision fc13ba1986fcb905854dc8de2f52eee395f49b5b

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

Workload generation

Four authored requests cover square text-to-image generation, landscape generation, image editing and transparent-background generation, with 1, 2 and 4 requested outputs and seeds starting at 42. Resolutions are 1024×1024 and 1536×1024. The full pinned model executes its original eight-step Turbo schedule with guidance scale 1. Selected diffusion steps and the first transformer block are captured alongside matching input/output projections. Each saved tensor preserves every element; layouts are normalized to contiguous storage.

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: 28.
• Distinct input sets: 12.
• Exact runtime shapes: 4.
• Tensor snapshots: 62, 1.814 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-prefix-cached-block-causal-h32-d128", 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/.