test(torchwave): Barrier for a transposed leaf of the same kernel (#18863) - #18863
Open
oerling wants to merge 5 commits into
Open
test(torchwave): Barrier for a transposed leaf of the same kernel (#18863)#18863oerling wants to merge 5 commits into
oerling wants to merge 5 commits into
Conversation
✅ Deploy Preview for meta-velox canceled.
|
meta-cla
Bot
added
the
CLA Signed
Contributor
|
@oerling has exported this pull request. If you are a Meta employee, you can view the originating Diff in D118061299. |
CI Failure Analysis
🟡 Presto Fuzzer — FUZZER Failure View logsFuzzer crash: 2 of 4 instances failed. Instance 1 (seed=31842952) — Result verification mismatch: Instance 3 (seed=279323501) — F14 hash table assertion failure during rehash: 🟡 Expression Fuzzer with Presto SOT — FUZZER Failure View logsFuzzer crash: 1 of 4 instances failed. Instance 2 (seed=635187378) — Velox vs Presto result mismatch: The mismatch shows Velox returned Correlation with PR changes: The PR ( Known issues:
Reproduce locally: # Presto Fuzzer — Instance 1 (result verification mismatch)
./_build/debug/velox/expression/fuzzer/velox_expression_fuzzer_test \
--seed 31842952 \
--enable_variadic_signatures \
--velox_fuzzer_enable_complex_types \
--velox_fuzzer_enable_decimal_type \
--lazy_vector_generation_ratio 0.2 \
--common_dictionary_wraps_generation_ratio=0.3 \
--velox_fuzzer_enable_column_reuse \
--velox_fuzzer_enable_expression_reuse \
--max_expression_trees_per_step 2 \
--retry_with_try \
--special_forms="and,or,cast,coalesce" \
--enable_dereference \
--duration_sec 300
# Presto Fuzzer — Instance 3 (F14Table rehash assertion)
./_build/debug/velox/expression/fuzzer/velox_expression_fuzzer_test \
--seed 279323501 \
--enable_variadic_signatures \
--velox_fuzzer_enable_complex_types \
--velox_fuzzer_enable_decimal_type \
--lazy_vector_generation_ratio 0.2 \
--common_dictionary_wraps_generation_ratio=0.3 \
--velox_fuzzer_enable_column_reuse \
--velox_fuzzer_enable_expression_reuse \
--max_expression_trees_per_step 2 \
--retry_with_try \
--special_forms="and,or,cast,coalesce" \
--enable_dereference \
--duration_sec 300
# Expression Fuzzer with Presto SOT — Instance 2 (Velox vs Presto mismatch)
./_build/debug/velox/expression/fuzzer/velox_expression_fuzzer_test \
--seed 635187378 \
--enable_variadic_signatures \
--velox_fuzzer_enable_complex_types \
--lazy_vector_generation_ratio 0.2 \
--common_dictionary_wraps_generation_ratio=0.3 \
--velox_fuzzer_enable_column_reuse \
--velox_fuzzer_enable_expression_reuse \
--enable_dereference \
--special_forms="cast,coalesce,if" \
--velox_fuzzer_max_level_of_nesting=1 \
--presto_url=http://127.0.0.1:8080 \
--duration_sec 300Recommended fix: No action needed on this PR. These are pre-existing flaky fuzzer failures unrelated to the torchwave changes. Safe to re-run CI or merge if other checks pass. |
Selective Build Plan
Selective build plan |
Yuhta
approved these changes
Sep 5, 2026
Summary:
Hooks the IG-ranking ROO preproc benchmark to the TorchWave GPU engine via a new
`--torchwave` arm, using the whole-graph AOT path (the AOTInductor alternative
from D107956042): the entire preproc graph is exported once and handed to
TorchWave as a single unit. TorchWave fuses what it can and falls back to
nativert kernels for the rest, so there is no `torch.compile` and no automatic
graph breaks.
- `roo_wave_runner.py` (new): `WholeGraphWaveModel` exports the eager preproc
with `export_preproc_pt2`, packages it whole via `compile_and_package(...,
decompose=False)` (the export already ran the ROO-safe
`run_decompositions(decomp_table={})`), and loads it with `torchwave.load`.
Inputs are flattened with `fx_pytree.tree_flatten_spec` against the export
in_spec (so KJT string keys are absorbed into the spec, not passed as leaves)
and moved to CUDA; outputs are compared against eager position-by-position via
the out_spec. `set_kernel_cache_dir` points the compiler at an on-disk cubin
cache.
- `roo_train_profile_benchmark.py`: adds `--torchwave` (checks parity vs eager,
then times warmup + repeats; requires `--on-gpu`) and `--kernel-cache-dir` /
`--kernel_cache_dir`.
- `pybind.cpp`: exposes `WaveConfig.kernel_cache_dir` to Python. The on-disk
cubin cache (`KernelFsCache`) is already wired in C++; this is the missing
setter so a repeated `--torchwave` run skips NVRTC compilation.
- `tests/test_roo_preproc_wave_parity.py` (new): GPU `python_unittest` that
builds the eager module from the model config (no Hive) and asserts
TorchWave-vs-eager parity on a sample batch. Skips when CUDA or the batch is
absent; the batch path can be overridden with `ROO_SAMPLE_BATCH`.
- `Executor.h` / `Executor.cpp`: `runInputsReuse(inputs)` holds ONE execution frame across calls with weights and constants resident, refilling only the user inputs, instead of the per-call getFrame / clear / returnFrame cycle. That is how the benchmark's wave arm is timed, and it matches what the C++ executor test measures. Non-persistent values are still cleared each call, since successive batches have data-dependent shapes and a stale-sized intermediate would otherwise be reused. `TW_TIME_EXEC=1` prints the per-call phase breakdown (clear/getFrame, fillUserInputs, exec+sync).
- `Executor.h` / `Executor.cpp` also add TU-local setters -- `setWaveTrace`, `setFreeIntermediates`, `setAutoAdjustCost`, `setIsCg`, `setKernelCacheDir`, `setAllStandalone`, `setEnableReuse`, `setBlockSize` -- so a caller in another translation unit (the pybind extensions) reaches the exact `WaveConfig` singleton the executor reads. A plain `wave_config()` assignment from another TU can land on a duplicated inline-static instance and silently not take effect, which is the failure this removes.
- `torchwave_meta/WaveExecutorPy.cpp` and the two `.pyi` stubs expose those setters plus `last_perf_report()` on the fused `_torchwave_meta` module, so the benchmark can configure the same engine instance it runs.
- `roo_train_pt2_export_benchmark.py`: a `--torchwave` arm that validates TorchWave against the SIGMOID INTERPRETER (and the PT2 module) rather than eager. Both execute the same exported graph, so this is the apples-to-apples engine comparison; a failure in one executor is reported without aborting the others.
- `GraphPrep`: `rewriteMergeAndDedupToTw` retargets the sparsenn merge-and-dedup nodes to their fused TorchWave `_tw` CUDA equivalents, but only when those ops are actually registered in the dispatcher. When they are absent it is a no-op and the base op runs as a nativert standalone, so one `load()` path serves both the base engine and the fused-ops build.
Differential Revision: D113334660
Summary: Explore Triton as a torchwave codegen target by porting the `tw.masked_select_jagged` op to standalone Triton kernels. This does NOT touch torchwave proper (no WaveGraph, no op registration); it adds a new self-contained package under `torchwave_meta/tests/triton/`. Two variants parallel the two CUDA forms in `minimal_viable_ai/torchwave_meta/Select.cuh`: - multikernel: staged launches (per-chunk count -> scan of partials -> write exclusive prefix -> scatter -> new segment lengths), mirroring torchwave's multi-kernel execution mode. - cg: a single persistent kernel carrying the same three stages across device-wide atomic barriers -- the Triton analogue of the CUDA `opBarrier` / grid.sync (`num_programs` kept co-resident so the spin-barrier cannot deadlock; programs grid-stride the chunks). The standalone test parallels the CUDA gtest `minimal_viable_ai/torchwave_meta/tests/KernelTest.cu` (SelectKernelTest): same input generation, the same per-segment CPU reference, and the same sizes; it validates the Triton kernels against the CPU reference and also covers the grouped mixed-dtype path. Authored with Claude. Differential Revision: D114379433
Summary: Fuses `torch.ops.fb.batch_flip_and_truncate_sparse` on wave by splitting it, in `maybeReplace`, into `tw.flip_and_truncate_head` and `tw.flip_and_truncate_final` (device functions in `Flip.cuh`). The head processes one feature per block: it computes each row's post-truncation length, inclusive-scans those into new offsets and the original lengths into old offsets, and after an op barrier accumulates the per-feature output totals. The final partitions the concatenated output into fixed 8192-element tiles and grid-strides over them, binary-searching once per tile for the owning feature and row, then walking rows to the end of the tile. Fixed tiles keep every block's work equal even though rows span roughly 300 to 10000 elements after truncation. Each thread resolves the same per-row copy descriptor, which hoists the flip-mode branch out of the element loop and leaves coalesced writes and forward or reversed coalesced reads. It matches the eager kernel across the four flip modes (none, flip_before, flip_after, flip_before_and_after), per-feature `max_lengths`, `pad_value`, and `adaptive_max_len` with a per-row `adaptive_mask`. On the four flip instances of an internal preproc graph the final reaches 75-76% of an A100's nominal bandwidth, against 89% for a `cudaMemcpy` of the same size and 53-66% for the eager kernel's thread mapping. The final stages its per-feature constants (the offset, value and output pointers, row count, cap and mode) in shared memory rather than holding them in registers across the row loop. That is worth an occupancy step: the register version needs 64 registers and fits 4 blocks per SM, the shared version needs 48 and fits 5, and on a bandwidth-bound copy the extra block is worth about 3%. Sizing that array needs the feature count as a compile-time constant, which is where `Metadata::templateParamFuncs` comes in: a list of functions run on the node at codegen time to append extra template arguments to the device function call, empty for every other op. `Metadata::gridSizeSumsInputs` is the second new hook. An op whose work spans a whole tensor list was sized by its largest member, so the head got `ceil(768 / blockSize)` blocks no matter how many features it had, and most features queued onto three blocks; the flag switches that op to the sum instead. Neither hook changes any op that does not set it. Three wave-core fixes in `Compile.cpp` support this and correct general bugs. First, an elementwise producer of a `TensorList` element that a fused op reads is now materialized to that element's memory buffer in all modes, so a list built inside the op (for example values from an add) is written where the consumer reads it. Second, a fused op that reads such a list emits a barrier when a same-op element producer is not already behind one, independent of the input's `randomAccess` flag. Third, subgraph deduplication now hashes int-list (ScalarList) attributes, whose element values are baked into the generated code; without this two otherwise-identical ops that differ only in such an attribute (for example `adaptive_mask`) would share a wrong kernel. Differential Revision: D111315788
…rnel wrote
Summary:
An elementwise op's operands are all marked `isRegister` by
`registerElementwise`, and `callNeedsBarrier` skips register operands on
the grounds that they flow inline rather than through memory. That holds
for a value fused into the expression tree; it does not hold for a leaf,
which is a load like any other. The test was that the two are the same
thing, and they are not, so no pure-elementwise op ever asked for a
barrier.
What that costs is visible whenever the leaf does not map index to index.
The ROO preproc graph fuses
(%3274) = where.self(...)
(%11769) = transpose.int(%3274, dim0=0, dim1=1)
(%11800) = mul(%11771, %11769)
into one kernel with nothing but `__syncthreads()` between the write and
the read. Through the transpose the element thread i of the `mul` wants
was written by a different block, so `%11800` came out part stale and
part uninitialized -- and only in the multi-block modes, which is what
made it look like a miscompile rather than a race.
The operand test `callNeedsBarrier` already had is now
`valueNeedsBarrier`, and `generateElementwise` runs it over each
subgraph's own memory leaves.
Differential Revision: D118061298
…cebookincubator#18863) Summary: Pull Request resolved: facebookincubator#18863 A fused pair whose second op reads a transposed view of what the first wrote has to carry an intra-kernel barrier. The contrast case -- the same pair with the intermediate consumed straight through, which keeps it in a register -- must not, so the test still fails if every fused elementwise pair starts carrying one. Reviewed By: Yuhta Differential Revision: D118061299
meta-codesync
Bot
changed the title
oerling
force-pushed
the
export-D118061299
branch
from
September 5, 2026 18:09
818c925 to
e4391d2
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary:
A fused pair whose second op reads a transposed view of what the first
wrote has to carry an intra-kernel barrier. The contrast case -- the same
pair with the intermediate consumed straight through, which keeps it in a
register -- must not, so the test still fails if every fused elementwise
pair starts carrying one.
Reviewed By: Yuhta
Differential Revision: D118061299