security: enforce dense current-stream contract

[malaiwah]

Problem

The dense GEMM APIs accept a caller-selected CUDA stream, but the complete composite operation includes PyTorch allocation, fill, reduction, and epilogue work around raw CuTe launches. A non-current explicit stream therefore needs both execution dependencies and allocator-lifetime ownership at one boundary.

Implementation

• Preserve the existing stream=None path. It enters the original implementation after one branch and performs no CUDA stream lookup, synchronization, or record_stream work.
• Run the complete dense, fused-quant, or grouped fused-quant operation under the selected CUDA device and stream.
• For a distinct explicit stream, make the launch stream wait for the caller's current stream, record each unique external storage once, and make the caller stream wait for completion before returning.
• Avoid allocator recording and stream waits when the explicit stream is already current.
• Accept Torch streams and raw handles, validating raw-handle device ownership before wrapping them.
• Keep custom-op and raw-launch layers free of duplicate stream validation or ownership bookkeeping.

Cost model

• stream=None: one required Python branch; no stream resolution, events, or allocator bookkeeping.
• Explicit current stream: stream/device resolution and a device context; no events or allocator bookkeeping.
• Explicit non-current stream: two stream dependencies and one deduplicated record_stream pass over caller-owned storage.

Validation

• CPU/source contract tests on head 7b7420662f32871e8890edec78ec0c16dd40ca7d: 6 passed.
• Python compilation, test collection (8 tests), and git diff --check: passed.
• Two target-GPU regressions cover producer/launch/consumer ordering with allocator churn and explicit capture-stream replay. They remain pending on this exact head because all shared GPUs are occupied.
• Representative dense GEMM, weight-only, and block-FP8 eager/CUDA-graph measurements are required before merge.

Closes #153

[coderabbitai]

Warning

Review limit reached

@lukealonso, you've reached your PR review limit, so we couldn't start this review.

Next review available in: 27 minutes

You've used all free OSS reviews for now. Wait for the free limit to reset to keep reviewing this public repository.

How can I continue?

After more reviews become available, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

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 configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 8dc950d2-aee5-4c74-acb4-374308a6529d

📥 Commits

Reviewing files that changed from the base of the PR and between 954fd01 and 7b74206.

📒 Files selected for processing (3)

• b12x/_lib/dense_gemm.py
• b12x/_lib/utils.py
• tests/gemm/test_dense_stream_validation.py

---

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.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[malaiwah]

Aggregate integration qualification evidence

This PR's patch was included in the combined security qualification branch malaiwah/sparkinfer:security/integrated-qualification.

• Qualified integration commit: 01d74106b9efbff9bce1ecd66893b4cbadc0c919
• Integration baseline/merge-base: 195e26c5b67eb9c162f00fd8907d06e0fe27c569
• Environment: PyTorch 2.13.0+cu130, CUDA 13.0, CUTLASS DSL 4.6.0
• Hardware: 2× NVIDIA RTX PRO 6000 Blackwell Server Edition, 97,887 MiB each
• Benchmark/test GPU UUID: GPU-554f7aeb-578c-187f-4f29-af06f2170507

The combined changed-path GPU run exercised compiler-cache, attention, page/index validation, stream handling, PCIe IPC/DCP, GEMM, MoE routing, FP6 artifact/export/dequantization, profiler, packaging, and release paths:

1011 passed, 4 skipped, 1 warning in 125.69s

The artifact/profiler/release subset independently completed:

403 passed, 2 skipped in 38.48s

With CUDA hidden, the CPU-compatible security set completed:

525 passed, 59 skipped in 39.75s

Correctness gates passed for the benchmarked varlen attention paths: baseline and integrated arms produced identical checksums and maximum magnitude, finite tensors, and nonzero output. Balanced ABBA timings on the GPU above measured:

Case	Baseline run means	Integrated run means	Direction
Varlen (128,256,64,512), QH=4, KVH=2, D=64	20.888, 21.024 µs	66.070, 64.094 µs	integrated 3.106× slower, +44.13 µs
Varlen 4×512, QH=8, KVH=2, D=128	24.787, 24.813 µs	67.469, 63.840 µs	integrated 2.647× slower, +40.85 µs

The fixed eager-launch cost comes from graph-safe validation/sanitization kernels. Standalone MoE route sanitization measured 7.387/7.384/7.322 µs for 8/512/4096 routes. Invalid-row scrub measured 6.967/6.909/7.031 µs for 8/64/512 rows at K=4096.

Memory deltas observed in qualification:

• short varlen scratch: 506,880 → 508,932 bytes, +2,052 bytes (+0.405%)
• prefill varlen scratch: 4,259,840 → 4,261,892 bytes, +2,052 bytes (+0.048%)
• MoE route sanitizer: 8 bytes per route; 32 KiB at the default 4096-route capacity

Scope: this is evidence for the exact aggregate integration commit, not a claim that a later PR-head update was independently rerun. The full repository test tree is not CPU-compatible, and no end-to-end vLLM/SGLang model-serving run was performed. The temporary GPU runner was destroyed after qualification.

[lukealonso]

The stream-lifetime hardening is directionally useful, but this implementation records the same tensors at the wrapper, custom-op, and launch layers. That puts repeated record_stream calls and allocator bookkeeping on the dense GEMM eager hot path, while also narrowing the API to the ambient current stream.

Please consolidate ownership recording at one boundary, preserve an explicit and testable stream contract, and provide target-GPU measurements for representative dense GEMM, weight-only, and block-FP8 shapes in both eager and CUDA-graph modes. Rebase the result with #194/#213 so the overlapping admission and buffer-contract changes are evaluated as one path.

[lukealonso]

Closing rather than merging. The stated threat is same-process misuse of an explicitly selected non-current stream. This patch changes the central dense launch path and adds stream dependencies plus allocator bookkeeping, while its required GPU correctness and performance checks remain pending. That is outside the zero-kernel-path-change boundary and introduces unmeasured hot-path cost. The issue remains open for an API-contract decision.