[WS1][kernels] Batch-invariant attention (standard softmax--cuda version) - #240
Conversation
…ign#147) Forward: QK → masked softmax+LSE → PV (FP32 intermediate, fixed reduction order) Backward: dP → softmax_bwd → dQ/dK/dV (§4.1 fixed GQA accumulation order) - 3 forward kernels + 5 backward kernels, all batch-invariant (no split-KV, no atomics) - autograd.Function wrapper for differentiable forward - RL-Align#108 harness integration (OP_SPECS, tolerance_contract, operator_inputs) - 634 tests: full sweep, scale, padding, LSE, batch invariance, chunked-prefill, prefill/decode handoff, KV-cache cat, FP64 gradient, GQA order validation - docs/operators/attention.md: fixed reduction order + shared contract - Benchmark for Qwen3-8B representative shapes
|
Warning Review limit reached
Next review available in: 18 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (5)
📝 WalkthroughWalkthroughAdds a deterministic standard-softmax CUDA attention backend with fixed reduction order, masking, LSE output, backward gradients, registry dispatch, benchmarks, documentation, and CUDA correctness and invariance tests. ChangesDeterministic attention backend
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant DeterministicAttentionOp
participant CUDAExtension
participant CUDAKernels
Caller->>DeterministicAttentionOp: Request attention output
DeterministicAttentionOp->>CUDAExtension: Submit q, k, v, mask, causal, scale
CUDAExtension->>CUDAKernels: Run fixed-order QK, softmax/LSE, and PV
CUDAKernels-->>CUDAExtension: Return output, LSE, and saved probabilities
CUDAExtension-->>DeterministicAttentionOp: Return output and LSE
DeterministicAttentionOp-->>Caller: Return attention output
Possibly related issues
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (5)
benchmarks/benchmark_deterministic_attention.py (2)
74-85: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueAvoid mutating the global configuration dictionary and clear the cache on OOM.
The current loop mutates the
QWEN3_8B_SHAPESdictionaries in-place by popping and restoring thelabelkey. If an unhandled exception (likeKeyboardInterrupt) interrupts the loop, the dictionary is left in a mutated state.Additionally, if a configuration triggers an Out Of Memory (OOM) error, explicitly emptying the CUDA cache in the
exceptblock can reduce memory fragmentation, improving the chances that subsequent smaller shapes succeed.♻️ Proposed refactor
for shape in QWEN3_8B_SHAPES: - label = shape.pop("label") + kwargs = shape.copy() + label = kwargs.pop("label") try: - result = benchmark_attention(**shape, dtype=dtype, warmup=args.warmup, iters=args.iters) + result = benchmark_attention(**kwargs, dtype=dtype, warmup=args.warmup, iters=args.iters) print( f"{label:<25} {result['latency_ms']:>12.3f} " f"{result['peak_memory_mb']:>12.1f} " f"{result['scores_materialization_mb']:>11.1f}" ) except RuntimeError as exc: print(f"{label:<25} {'OOM' if 'out of memory' in str(exc) else 'ERROR':>12}") - shape["label"] = label + if "out of memory" in str(exc): + torch.cuda.empty_cache()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/benchmark_deterministic_attention.py` around lines 74 - 85, Update the loop over QWEN3_8B_SHAPES to avoid mutating each global configuration: read label without pop and pass a separate shape copy without the label to benchmark_attention. In the RuntimeError handler, detect OOM as currently done and empty the CUDA cache before continuing so later configurations can run.
30-37: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueWrap the benchmark execution in
torch.no_grad().While the input tensors do not have
requires_grad=True(meaning PyTorch won't build an autograd graph and memory won't leak), it is best practice to wrap inference benchmarking loops in atorch.no_grad()context. This explicitly communicates intent and ensures that no unnecessary autograd dispatch overhead interferes with the timing.♻️ Proposed refactor
- for _ in range(warmup): - op.forward(q, k, v, causal=True) - torch.cuda.synchronize() - - start = time.perf_counter() - for _ in range(iters): - op.forward(q, k, v, causal=True) - torch.cuda.synchronize() + with torch.no_grad(): + for _ in range(warmup): + op.forward(q, k, v, causal=True) + torch.cuda.synchronize() + + start = time.perf_counter() + for _ in range(iters): + op.forward(q, k, v, causal=True) + torch.cuda.synchronize()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/benchmark_deterministic_attention.py` around lines 30 - 37, Wrap both the warmup and timed operation loops in the benchmark around op.forward calls within a torch.no_grad() context, while preserving the existing CUDA synchronization and timing boundaries.csrc/cuda/attention/deterministic_attention.cu (1)
289-342: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd launch error checks after each kernel launch.
None of the
<<<...>>>launches are followed by an error check, so a bad launch config or in-kernel fault surfaces later at an unrelated CUDA call, making failures hard to attribute. AppendC10_CUDA_KERNEL_LAUNCH_CHECK();after each launch here and in the backward wrapper.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@csrc/cuda/attention/deterministic_attention.cu` around lines 289 - 342, Add C10_CUDA_KERNEL_LAUNCH_CHECK() immediately after each qk_kernel, masked_softmax_lse_kernel, and pv_kernel launch in the forward wrapper, and after every corresponding CUDA kernel launch in the backward wrapper. Keep each check directly adjacent to its launch so configuration and execution errors are attributed to the correct kernel.rl_engine/kernels/gtest/operator_inputs.py (1)
96-101: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win
operator_shape_nameno longer reflects actual generated shapes for"attention".Now that
_make_attention_inputsderivesskv,n_heads,n_kv_headsfromargs, the unchangedoperator_shape_nameentry (f"{batch}x{DEFAULT_N_HEADS}x{seq}x{DEFAULT_HEAD_DIM}") still hardcodesDEFAULT_N_HEADSand reusesseqfor what is actuallySkv, so harness reports will show a misleading shape whenever these args diverge from defaults."attention": ( f"{batch}x{_arg_int(args, 'n_heads', DEFAULT_N_HEADS)}x{seq}x" f"{_arg_int(args, 'skv', seq)}x{DEFAULT_HEAD_DIM}" ),🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/kernels/gtest/operator_inputs.py` around lines 96 - 101, Update the "attention" entry in operator_shape_name to derive n_heads and skv from args via _arg_int, matching _make_attention_inputs instead of hardcoding DEFAULT_N_HEADS or reusing seq. Preserve the existing batch and DEFAULT_HEAD_DIM components while reporting the actual generated shape.rl_engine/kernels/ops/cuda/attention/deterministic_attn.py (1)
25-67: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSave
key_padding_maskviasave_for_backward; consider@once_differentiable.Line 47 stores
mask_cdirectly onctxinstead of throughctx.save_for_backward. Per PyTorch's docs, this is discouraged: "All tensors intended to be used in the backward pass should be saved with save_for_backward (as opposed to directly on ctx) to prevent incorrect gradients and memory leaks." Also,backwardisn't decorated with@once_differentiable; since dQ/dK/dV come from an opaque C++ call with no further autograd tracking, acreate_graph=Truebackward would fail with a confusing autograd error rather than an explicit one.♻️ Proposed fix
+from torch.autograd.function import once_differentiable + class _DeterministicAttentionFn(Function): `@staticmethod` def forward(...): ... - ctx.save_for_backward(q_c, k_c, v_c, P) + ctx.save_for_backward(q_c, k_c, v_c, P, mask_c if mask_c is not None else torch.tensor([])) ctx.causal = causal ctx.scale = scale - ctx.key_padding_mask = mask_c + ctx.has_mask = mask_c is not None ctx.mark_non_differentiable(lse) return out, lse `@staticmethod` + `@once_differentiable` def backward(ctx, grad_out, grad_lse): - q_c, k_c, v_c, P = ctx.saved_tensors + q_c, k_c, v_c, P, mask_saved = ctx.saved_tensors + mask_c = mask_saved if ctx.has_mask else None ...🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@rl_engine/kernels/ops/cuda/attention/deterministic_attn.py` around lines 25 - 67, Update _DeterministicAttentionFn.forward and backward to save mask_c through ctx.save_for_backward and retrieve it from ctx.saved_tensors, keeping the existing non-tensor context fields unchanged. Decorate _DeterministicAttentionFn.backward with PyTorch’s once-differentiable guard so higher-order gradient requests fail explicitly for the opaque C++ backward.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@csrc/cuda/attention/deterministic_attention.cu`:
- Line 284: Add an at::cuda::OptionalCUDAGuard using device_of(q) in both
deterministic_attention_forward at
csrc/cuda/attention/deterministic_attention.cu:284, after
check_deterministic_attention_inputs(...) and before getCurrentCUDAStream(), and
deterministic_attention_backward at
csrc/cuda/attention/deterministic_attention.cu:560 before acquiring the stream.
This pins kernel execution to q’s device in both wrappers.
In `@docs/operators/attention.md`:
- Around line 174-179: Update the fenced pseudocode block near the
attention-loop example to specify the text language identifier, changing the
opening fence to use text while preserving the pseudocode content unchanged.
In `@tests/test_deterministic_attention_cuda.py`:
- Line 42: Defer CUDA tensor creation in module-level HARNESS_CASES and
_make_case until tests execute, so importing the module never calls torch.randn
with DEVICE="cuda". Preserve the existing pytestmark skipif behavior so CUDA
tests remain skipped on non-CUDA machines while collection, including pytest
--collect-only, succeeds.
---
Nitpick comments:
In `@benchmarks/benchmark_deterministic_attention.py`:
- Around line 74-85: Update the loop over QWEN3_8B_SHAPES to avoid mutating each
global configuration: read label without pop and pass a separate shape copy
without the label to benchmark_attention. In the RuntimeError handler, detect
OOM as currently done and empty the CUDA cache before continuing so later
configurations can run.
- Around line 30-37: Wrap both the warmup and timed operation loops in the
benchmark around op.forward calls within a torch.no_grad() context, while
preserving the existing CUDA synchronization and timing boundaries.
In `@csrc/cuda/attention/deterministic_attention.cu`:
- Around line 289-342: Add C10_CUDA_KERNEL_LAUNCH_CHECK() immediately after each
qk_kernel, masked_softmax_lse_kernel, and pv_kernel launch in the forward
wrapper, and after every corresponding CUDA kernel launch in the backward
wrapper. Keep each check directly adjacent to its launch so configuration and
execution errors are attributed to the correct kernel.
In `@rl_engine/kernels/gtest/operator_inputs.py`:
- Around line 96-101: Update the "attention" entry in operator_shape_name to
derive n_heads and skv from args via _arg_int, matching _make_attention_inputs
instead of hardcoding DEFAULT_N_HEADS or reusing seq. Preserve the existing
batch and DEFAULT_HEAD_DIM components while reporting the actual generated
shape.
In `@rl_engine/kernels/ops/cuda/attention/deterministic_attn.py`:
- Around line 25-67: Update _DeterministicAttentionFn.forward and backward to
save mask_c through ctx.save_for_backward and retrieve it from
ctx.saved_tensors, keeping the existing non-tensor context fields unchanged.
Decorate _DeterministicAttentionFn.backward with PyTorch’s once-differentiable
guard so higher-order gradient requests fail explicitly for the opaque C++
backward.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 1b0d81f1-354d-4a6f-a12b-900f5e53b36e
📒 Files selected for processing (15)
benchmarks/benchmark_deterministic_attention.pycsrc/cuda/attention/deterministic_attention.cucsrc/ops.cppdocs/operators/attention.mdrl_engine/_C.pyirl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/gtest/tolerance_contract.jsonrl_engine/kernels/ops/cuda/attention/__init__.pyrl_engine/kernels/ops/cuda/attention/deterministic_attn.pyrl_engine/kernels/registry.pysetup.pytests/test_attention.pytests/test_deterministic_attention_cuda.pytests/test_tolerance_contract.py
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
Left a few comments, otherwise LGTM.
| ) | ||
|
|
||
|
|
||
| HARNESS_CASES = [ |
There was a problem hiding this comment.
Defer the creation of CUDA tensors in HARNESS_CASES until runtime to prevent pytest collection failures on CPU-only CI nodes.
| const int64_t Hkv = k_contig.size(1); | ||
| const int64_t Skv = k_contig.size(2); | ||
|
|
||
| auto stream = at::cuda::getCurrentCUDAStream(); |
There was a problem hiding this comment.
Insert at::cuda::OptionalCUDAGuard before acquiring the CUDA stream to ensure kernels launch on the correct tensor device and prevent illegal memory accesses.
| const int64_t Hkv = k_c.size(1); | ||
| const int64_t Skv = k_c.size(2); | ||
|
|
||
| auto stream = at::cuda::getCurrentCUDAStream(); |
| results = _C.deterministic_attention_forward(q_c, k_c, v_c, causal, float(scale), mask_c) | ||
| out, lse, P = results[0], results[1], results[2] | ||
|
|
||
| ctx.save_for_backward(q_c, k_c, v_c, P) |
There was a problem hiding this comment.
Save all required tensors using ctx.save_for_backward to prevent memory leaks, and add the @once_differentiable decorator to the backward function to explicitly fail unsupported higher-order gradient requests.
|
please resolve the code conflicts first, Thank you. |
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@rl_engine/kernels/gtest/operator_specs.py`:
- Line 73: Close the preceding attention OperatorSpec before the
"batch_invariant_logp" entry by adding the missing terminator after
grad_input_names=("q", "k", "v"),, preserving the dictionary entry separator so
OP_SPECS remains valid Python.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: e2a510e0-b0cb-4967-9fc9-e30892819ac6
📒 Files selected for processing (3)
rl_engine/kernels/gtest/operator_inputs.pyrl_engine/kernels/gtest/operator_specs.pyrl_engine/kernels/registry.py
🚧 Files skipped from review as they are similar to previous changes (2)
- rl_engine/kernels/registry.py
- rl_engine/kernels/gtest/operator_inputs.py
@Flink-ddd, thank you for the review. I’ve resolved the merge conflicts and addressed all four comments in commit 040acd9:
I also ran Black, Flake8, and the CPU-only collection checks successfully. The updated changes have been pushed. Could you please take another look when convenient? |
Flink-ddd
left a comment
Flink-ddd
left a comment
There was a problem hiding this comment.
LGTM now, Thank you for update. also cc @KJLdefeated @inaniloquentee @frank-2077 PTAL.
3 participants
Summary
Adds a deterministic CUDA standard-softmax attention implementation for #147.
The implementation provides batch-invariant forward and backward paths with a fixed accumulation order. It supports causal masking, key-padding masks, GQA/MQA, variable sequence lengths, chunked prefill, decode, and KV-cache handoff.
The CUDA path matches the existing FP32 golden reference NativeAttentionOp within the shared #108 tolerance contract. Outputs and gradients are bitwise invariant when the same sequence is evaluated at different batch sizes, batch positions, or prefill chunk sizes.
Closes #147.
Implementation
The forward path is decomposed into three deterministic CUDA kernels:
The backward path implements deterministic:
For GQA, dK and dV accumulate query heads and query positions in a fixed order. There are no cross-CTA atomics, split-KV paths, or launch-order-dependent reductions.
Path Status
━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
CUDA forward Deterministic QK → masked softmax/LSE → PV
───────────────── ───────────────────────────────────────────────────────
CUDA backward Deterministic dP → softmax backward → dQ/dK/dV
───────────────── ───────────────────────────────────────────────────────
Accumulation FP32 intermediates and statistics
───────────────── ───────────────────────────────────────────────────────
Dtypes FP16 and BF16
───────────────── ───────────────────────────────────────────────────────
Head dimension D=128
───────────────── ───────────────────────────────────────────────────────
Attention modes Causal and non-causal
───────────────── ───────────────────────────────────────────────────────
Head layouts MHA, MQA, and GQA
───────────────── ───────────────────────────────────────────────────────
Padding key_padding_mask support, including fully masked rows
───────────────── ───────────────────────────────────────────────────────
Prefill/decode Shared kernel and causal-offset contract
───────────────── ───────────────────────────────────────────────────────
Registry CUDA_DETERMINISTIC_ATTENTION with native fallback
───────────────── ───────────────────────────────────────────────────────
Harness Integrated with the #108 operator suite
───────────────── ───────────────────────────────────────────────────────
Benchmark Qwen3-8B representative prefill/decode shapes
Fixed Reduction Contract
QK
Each score has exactly one writer. The head dimension is accumulated in ascending order:
for d = 0 .. D-1:
score += float(q[d]) * float(k[d])
Softmax and LSE
Each (batch, query_head, query_position) row is processed by one CTA.
Max and sum-exp use a fixed power-of-two shared-memory tree reduction. The reduction topology does not depend on batch size, sequence length, SM count, or runtime split count.
PV
Each output element has exactly one writer:
for k = 0 .. Skv-1:
out += probability[k] * float(value[k])
GQA backward
Each dK and dV element is accumulated by one thread:
for local_head = 0 .. group_size-1:
for query_position = 0 .. Sq-1:
grad += ...
No split-KV or cross-CTA atomic accumulation is used.
Prefill / Decode / KV-cache Contract
All modes use the same deterministic attention kernels:
The causal condition is:
key_index <= Skv - Sq + query_index
This keeps full prefill, chunked prefill, and decode on the same effective reduction path.
The Python API exposes:
Acceptance Criteria
Requirement Result
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Batch-size invariance ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Batch-position invariance ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Batch chunking invariance ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Chunked-prefill on/off ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Chunked-prefill with padding ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Softmax and LSE invariance ✅ Bitwise for matching physical shapes
───────────────────────────────────── ─────────────────────────────────────────
Fixed and variable sequence lengths ✅
───────────────────────────────────── ─────────────────────────────────────────
Causal and padding masks ✅
───────────────────────────────────── ─────────────────────────────────────────
Prefill/decode handoff ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
KV-cache concatenation handoff ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Backward correctness ✅ Within #108 tolerance
───────────────────────────────────── ─────────────────────────────────────────
Gradient batch invariance ✅ Bitwise
───────────────────────────────────── ─────────────────────────────────────────
Fixed accumulation order ✅ Documented and tested
───────────────────────────────────── ─────────────────────────────────────────
No variable split-KV path ✅
A valid-only sequence and a physically padded sequence use different reduction widths. That comparison is therefore checked within the shared accuracy tolerance rather than bitwise. Batch placement and chunking with the same effective physical shape remain bitwise invariant.
Validation Environment
Item Value
━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
GPU NVIDIA H20
──────────────────── ───────────────────────────────────────────────────────────────────────────
Compute capability 9.0
──────────────────── ───────────────────────────────────────────────────────────────────────────
GPU memory 97,871 MiB
──────────────────── ───────────────────────────────────────────────────────────────────────────
Driver 580.65.06
──────────────────── ───────────────────────────────────────────────────────────────────────────
PyTorch 2.8.0+cu128
──────────────────── ───────────────────────────────────────────────────────────────────────────
CUDA runtime 12.8
──────────────────── ───────────────────────────────────────────────────────────────────────────
Python 3.12.3
──────────────────── ───────────────────────────────────────────────────────────────────────────
Extension Editable CUDA build with deterministic attention forward/backward symbols
Correctness / Tests
Issue #147 CUDA suite
python -m pytest tests/test_deterministic_attention_cuda.py -q -ra
Result:
609 passed in 3.38s
Coverage includes:
Attention integration tests
python -m pytest
tests/test_attention.py
tests/test_operator_inputs.py
tests/test_kernel_registry.py
-q -ra
Result:
50 passed in 2.54s
Full repository suite
python -m pytest -q -ra
Result:
1165 passed, 102 skipped in 19.42s
The skipped tests are environment-specific FlashAttention, ROCm, CI-only extension enforcement, and optional SM90 fused-logp tests. No deterministic-attention test was skipped.
Pre-commit
pre-commit run --all-files
All hooks pass:
Benchmark
Command:
python benchmarks/benchmark_deterministic_attention.py
--dtype bf16
--warmup 1
--iters 2
Environment: NVIDIA H20. These numbers are smoke-test measurements rather than a formal performance study.
Shape Latency (ms) Peak memory (MiB) Scores (MiB)
━━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━━━━━━ ━━━━━━━━━━━━━━
decode-128 0.057 0.5 0.0
──────────────────── ────────────── ─────────────────── ──────────────
decode-1k 0.121 4.1 0.1
──────────────────── ────────────── ─────────────────── ──────────────
prefill-128 0.344 4.5 2.0
──────────────────── ────────────── ─────────────────── ──────────────
prefill-512 4.821 42.1 32.0
──────────────────── ────────────── ─────────────────── ──────────────
prefill-1k 19.150 148.1 128.0
──────────────────── ────────────── ─────────────────── ──────────────
batch4-prefill-128 1.287 18.1 8.0
──────────────────── ────────────── ─────────────────── ──────────────
batch8-prefill-64 0.687 14.1 4.0
Memory / Performance Tradeoff
The first deterministic baseline materializes full FP32 tensors:
scores: [B, Hq, Sq, Skv]
P: [B, Hq, Sq, Skv]
This gives a simple, auditable fixed reduction order but has quadratic sequence-length memory usage and is not intended to replace FlashAttention for maximum throughput.
A future implementation may use a fixed-order dual-kernel design, provided it preserves the same batch-invariance contract.
Limitations
Summary by CodeRabbit
New Features
Documentation
Tests
Chores