Repository navigation
Conversation
The workspace cache currently ignores the calling CUDA stream. Overlapping GEMMs can therefore overwrite the same scratch buffer and return incorrect results. Include the stream in the cache key, keep reuse on each stream, and add ownership and FP32/BF16 numerical regression tests. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
github-actions
Bot
added
the
community-contribution
 [Medium impact] The PR appears safe to merge; no new blocking issues were found. SummaryThe PR replaces cached cuBLAS scratch with memory owned by PyTorch’s CUDA allocator for each call.
Diagram%%{init: {'theme': 'neutral'}}%%
flowchart TD
A[GEMM call] --> B{Execution path}
B -->|Ordinary native| C[Allocate scratch in C++]
B -->|List-based grouped| D[Allocate separate aligned slices]
B -->|Custom recipe or communication overlap| E[Allocate explicit Tensor workspace]
C --> F[Submit GEMM on caller stream]
D --> G[Run grouped GEMMs]
G --> H[Join helper streams to caller]
E --> I[Run existing explicit-workspace path]
F --> J[Release local owner]
H --> J
I --> J
J --> K[PyTorch controls reuse and graph-pool lifetime]
Reviews (22) · Last reviewed commit: "Merge branch 'main' into fix/cublas-work..." · Reviewed by Greptile |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Keep scratch local to each invocation so graph captures use their own private pools and temporary streams do not retain tensors in Python. Remove Linear's preallocation workaround for the global workspace cache. Cover concurrent graph replays, graph generations, scratch reclamation, grouped GEMM stream joins, and compiled Linear forward and backward. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
Vtmpas
changed the title
Cover NN router dgrad with FP32 and BF16 operands in concurrent eager execution and independent CUDA graphs. Preserve NT wgrad coverage and exact references while retaining only one pair of large eager outputs. All 22 workspace cases pass on native SM90 H200 with CUDA 13 and PyTorch 2.13, including all four new NN cases. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
|
@Oleg-Goncharov , @ksivaman , hello! can i have your assistance here? |
|
@ksivaman kindly |
|
/te-ci pytorch L0 |
This comment has been minimized.
This comment has been minimized.
|
Want your agent to iterate on Greptile's feedback? Try greploops. |
Replace the cold-start test's cache_info assertions with checks for capture-time allocation and release of workspace tensors. Keep output and gradient comparisons across CUDA, CPU and meta initialization. Remove BasicLinear's obsolete cache preallocation and use the existing opaque-object capability guard for compiled execution. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
|
Fixed in f01b35e, on top of the latest main merge. The cold-start test now checks capture-time workspace allocation and release of Python tensor references instead of On H200 with CUDA 13 and PyTorch 2.13, the original |
|
@timmoon10 @ksivaman conflicts have been solved |
|
/te-ci pytorch L0 |
|
@ksivaman thanks for kicking off CI again. The last run stopped at |
|
@timmoon10 @ksivaman @ptrendx , sorry guys, are there any updated or feedback? Could you guide me here what's blocking, please? |
|
Kindly |
|
@ksivaman are there any updates? |
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
|
/te-ci pytorch L0 L1 |
|
For the future we should add the same change to the general grouped GEMM to mitigate the risk for that as well. |
I can open a separate PR for this if you’re okay with it. |
|
I am worried about re-adding the very high CPU overhead associated with allocated workspaces per gemm. Some of the work in #2377 was done to mitigate this. Although this is higher priority due to correctness, we should make a note of the added overhead via some benchmarking. |
Good point. Reusing storage through the CUDA allocator still leaves tensor creation and allocator overhead on every eager call. An earlier H200 microbenchmark measured about 1.31 µs for the intermediate per-stream cache versus 2.62 µs for invocation-local allocation. That only measured the workspace factory, though, and neither side represents a current main-vs-PR comparison. The module timings I collected were also too noisy to establish a reliable regression bound. |
|
@ksivaman I ran the current PR against main on H200. The CPU cost is real, particularly for small eager GEMMs. Compared main The warmed workspace factory, including releasing the returned reference, measured:
These are thread CPU times with no GEMM enqueued, so they isolate the factory cost from GPU execution. The older 1.31/2.62 µs figures in my previous reply are not the baseline for this comparison. For actual eager execution, the small CUDA graph replay is reported separately below; capture is excluded. It does not execute the Python workspace factory on every replay. I would not claim a blanket zero-overhead result: large BF16 Linear and large grouped block-FP8 GEMM had substantial pair-to-pair variation, including changes of both signs. Full latency results: BF16 and block FP8, eager and graph replayAll latencies are µs per GEMM call or per module forward+backward. Values are medians over six independent process pairs. Changes are the median of paired PR/main ratios, with their interquartile range in brackets; they are not calculated from the two separately aggregated latency medians. Positive means slower. IQR describes the observations, not a confidence interval. Grouped cases have four groups, with M rows per group. Shapes: small BF16
FP8 block scaling
Targeted recheck of the six warmup-flagged casesThe initial warmup screen flagged a >15% range over its last three batches. I retained those measurements and repeated only the six affected configurations in another six alternating process pairs, with 10×100 warmup calls and 7×100 timed calls per process. These are a separate series; the table above has not been overwritten.
No warmup-variance flags remained in this recheck, and outputs/gradients matched bitwise in all 36 additional paired case comparisons. Environment, timing and correctness checks
|
Use an invocation-local DataPtr from the PyTorch CUDA allocator instead of creating Python workspace tensors for ordinary GEMMs. Allocate one buffer with aligned, disjoint slices for the grouped multistream path. Keep stream-ordered reuse and graph-pool ownership without adding a persistent cache. Preserve explicit workspaces for custom recipes and communication overlap, and observe native requests in the compiled graph-generation regression. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
|
@ksivaman I pushed 10f6369 to address the host overhead without bringing back a workspace cache. Ordinary GEMMs now get an invocation-local I reran the comparison against main
Small is For host enqueue time, the small BF16 GEMM adds 0.211 µs, the small block-FP8 GEMM 0.518 µs, and the small BF16 grouped call 0.056 µs (paired medians). These measure the whole submission path and can include runtime backpressure; they are not isolated allocator CPU timings. The old Python factory microbenchmark no longer represents the ordinary GEMM path. Correctness:
Measurement details and limitsOne H200 for timing, Xeon Platinum 8570 with the host process pinned to one CPU, default GPU clocks, PyTorch 2.13.0+cu130, CUDA toolkit 13.0, cuBLASLt 13.5.1, driver 615.71.09, native PyTorch caching allocator. Six AB/BA pairs; each case has 10×100 warmup calls and 7×100 timed calls. Shapes were 128×256×256, 128×4096×128 and 4096³, covering raw GEMM, list/multistream grouped GEMM, Linear and GroupedLinear forward/backward in BF16 and block FP8, both eager and graph replay. Completion time synchronizes at batch boundaries. Raw GEMMs use TN layout and preallocated BF16 output; raw FP8 quantization is outside timing, while module quantization is included. Modules use BF16 parameters, no bias, optimizer or FP8 weight cache, and reset gradients each call. Grouped cases use four groups and explicitly select the list/multistream backend. CUDA graph capture is excluded from replay timing. Median paired replay changes were within ±1% across the measured cases; this is an observation, not a zero-overhead guarantee. Module graph timings include TE's graph/autograd wrapper. The existing AccumulateGrad stream warning occurred during setup in both variants; capture and replay completed. Three initial cases had warmup-variance flags: small FP8 eager GEMM, large FP8 eager GEMM and small FP8 Linear graph replay. Those were retained and separately repeated in six pairs with the same warmup/timing counts; no flags remained. The table uses the rechecks for the two affected eager GEMMs. IQR describes the observed pairs, not a confidence interval. A separate warmed Nsight Systems capture used the same BF16 GEMM kernel and showed no cudaMalloc/cudaFree calls. Profiling was disabled for latency measurements. Final source changes after the GPU runs were Black/clang-format formatting only; pre-commit, scoped pylint/cpplint and the license check pass. This qualifies the tested native-allocator paths on H200, not full-model throughput, distributed overlap or every allocator backend. Full L0 still needs CI. |
Observe generic_gemm requests during capture instead of the Python workspace factory, which ordinary GEMMs no longer call. Retain only scalar metadata and require the native None/zero workspace contract. Keep the cold-start initialization check by rejecting Python workspace allocation, and preserve the existing fullgraph, capture, and numerical assertions for every initialization and training case. Signed-off-by: Matvey Saprykin <mtvey.s@gmail.com>
|
@ksivaman could you rerun PyTorch L0 on |
Description
Concurrent GEMMs on independent CUDA streams receive the same cuBLAS workspace from the device/mode cache and can overwrite each other's partial sums. On one H200, two FP32
NTGEMMs with a reduction dimension of 16384 produced incorrect results in 35 of 48 calls; separate scratch produced zero mismatches against serial execution.Replace the Python workspace cache with invocation-local memory owned by PyTorch's CUDA allocator. Ordinary GEMMs allocate a native
DataPtrin the PyTorch binding, avoiding per-call Python Tensor creation. The list/multistream grouped path uses one allocation with disjoint, 256-byte-aligned workspace slices. Its helper streams join the caller before the owner is released. Eager reuse remains stream-ordered; captured allocations belong to their graph's private pool. No persistent per-stream cache or new device synchronization is added.Custom recipes and communication overlap retain explicit Tensor workspaces. Existing explicit-workspace callers remain supported; the private bindings additionally accept
Nonewith size zero to request native allocation. Usable workspace sizes, GEMM dtypes and algorithm selection are unchanged. The grouped-tensor backend is unchanged.Type of change
Changes
weight[128,4096],grad_output[8192,128], output[8192,4096], and reduction 128 against exact analytic references in FP32 and BF16. Preserve NT coverage.Validation
H200, PyTorch 2.13.0+cu130, CUDA toolkit 13.0, cuBLASLt 13.5.1, native PyTorch CUDA allocator. The candidate binding was rebuilt with the same flags as main
9ea05cf2e8dc, using the unchanged common library.tests/pytorch/test_gemm_workspace.pypass, including the compiled opaque-GEMM capture-generation case, with no skips.48a48c91896a: update both Linear compile-test observers to the native binding. Both edited observers pass H200 capture checks through the compiled-GEMM fixture; the full Linear tests remain capability-skipped (11 cases) on the available PyTorch 2.13 build, pending the compatible opaque-object CI environment.Checklist