noncausal_dense_attention_h32_d64

Noncausal self/cross attention with FP32 accumulation and BF16 output. Self attention uses Nq=Nk; cross attention allows unequal lengths. No dropout. Projection, normalization, rotary embeddings and output gates are outside the definition. No mask: every captured call's expanded visibility mask was all True (on one GPU the engine passes no mask, or an all-visible prompt mask), so KDA drops the dataset's mask input.

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

$O=\operatorname{softmax}(s Q K^T)V$. Every query can attend to every key. This is noncausal self/cross attention without a mask, with FP32 intermediate arithmetic and BF16 output.

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

Inputs

Outputs

Implementation and origins

FlashAttention 4 CuTe DSL without a mask. Converts the captured normalized audio-attention inputs to FP16 to reduce probability-rounding error, and returns BF16. All conversions are included in the call. Upstream: https://github.com/Dao-AILab/flash-attention (BSD-3-Clause).

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

• Lightricks/LTX-2.5 — revision 49b093a39c19169a8a81721dcc541e07e945232f

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

Workload generation

Three authored requests cover text-to-video, image-to-video and a longer video, each with 1, 2 and 4 outputs and seeds starting at 42. Requested grids are 768×448×49, 512×320×49 and 512×320×121 at 24 fps, with generated audio. The full pinned transformer and diffusion decoder execute eight steps with guidance scale 1. Transformer attention is sampled at steps 0, 4 and 7 in block 0; neighborhood attention retains complete decoder grids. Audio/video transformer attention uses math SDPA with BF16 inputs/outputs and FP32 intermediates. The optional auxiliary decoder keyframe stream is not exercised.

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: 108.
• Distinct input sets: 105.
• Exact runtime shapes: 30.
• Tensor snapshots: 417, 2.793 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-noncausal-dense-attention-h32-d64", 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/.