Add lookup operators with CMake build system and tests by aagalleg · Pull Request #127 · intel/torchlib-xpu

aagalleg · 2026-08-28T00:26:02Z

This PR adds Intel XPU implementations for FBGEMM dense embedding lookup and
split rowwise-Adagrad lookup operators, together with the build-time code
generation needed by the training extension.

Build and code generation

Jinja2 templates and Python generators are build inputs.
CMake runs the forward and backward generators automatically.
The build produces seven forward and five backward generated files in the
build tree.
Generated host sources, the static lookup operator files, and shared
utilities are compiled into fbgemm_xpu._C_training.
Generated files are neither committed nor installed as source files.
jinja2 is declared as an isolated build dependency.
fbgemm_xpu._C_training is imported after _C, so schemas exist before the
XPU and AutogradXPU implementations are registered.

Correctness fixes

reorder_batched_ad.cpp includes its dispatch macros directly instead of
relying on a transitive include.

Supported functionality

The plugin registers XPU implementations for:

fbgemm::dense_embedding_codegen_lookup_function
fbgemm::split_embedding_codegen_lookup_rowwise_adagrad_function_pt2

Runtime correctness is validated for:

no-bag, unweighted lookup (PoolingMode.NONE);
one- and two-table layouts with one or two batches per table and a uniform
embedding dimension;
FP32 and FP16 weight storage;
small (D=4) and general (D=36) forward kernels;
dense autograd;
split rowwise-Adagrad in-place updates;
warp and CTA backward/update paths, including 32 repeated indices targeting
one embedding row.

The PR does not claim support for pooled lookup, weighted lookup, VBE, global
weight decay, or cache-backed lookup.

Testing model

CI installs the built fbgemm-xpu wheel beside pristine
fbgemm-gpu-cpu==1.8.0.
Native extension loading, AutogradXPU registration, an existing XPU operator,
and pristine-runtime compatibility are tested before applying any FBGEMM
source patch.
The remaining FBGEMM patch is test-only and modifies only
fbgemm_gpu/test/**.
Lookup correctness tests live in packages/fbgemm-xpu/tests/ and call
torch.ops.fbgemm directly on XPU.
All twelve declared generated files are produced and compiled; the small and
general forward paths plus warp and CTA backward/update paths are also
executed by runtime tests.

Compatibility boundary

No files are copied into site-packages.
The installed FBGEMM Python frontend remains pristine.
This PR does not add ComputeDevice.XPU to FBGEMM.
Pristine high-level construction through
SplitTableBatchedEmbeddingBagsCodegen remains unavailable for XPU and is
tested in an isolated subprocess to fail without terminating or corrupting
the parent process.
Existing XPU operators and the direct lookup operators remain usable.

cc: @molintc, @mkrze

dvrogozh · 2026-09-03T14:02:27Z

Needs a rebase on top of recently merged 099a1c9.

dvrogozh · 2026-09-03T14:03:54Z

        pip install --no-cache-dir \
          "$(ls ./_wheels/${{ env.WHEELS }}-${{ env.PR_OR_SHA }}-${{ env.UNIQUE_ID }}/fbgemm_xpu*.whl)[test]" \
          --extra-index-url https://download.pytorch.org/whl/xpu
+    - name: 'Install patched FBGEMM Python frontend'


This frightens me :). Why do we need to do this? And this seems to do some magic with FBGEMM own sources rather than actually patching them, right?

Will be removed as discussed, this was a test-only workaround

dvrogozh · 2026-09-03T15:24:56Z

+++ b/fbgemm_gpu/fbgemm_gpu/split_table_batched_embeddings_ops_training.py
+@@ -359,6 +359,8 @@ def get_available_compute_device() -> ComputeDevice:
+         return ComputeDevice.CUDA
+     elif torch.mtia.is_available():


Can we try to further abstract this and use torch.accelerator API to collapse device specific code?

removed device specific code and using torch.accelerator instead

… generation Add jinja_environment.py module to support template-based code generation for FBGEMM-XPU kernels.

Add common.py module with CodeTemplate class that provides functionality for loading Jinja2 templates, rendering them with context variables, and writing generated files with appropriate headers.

Add torch_type_utils.py module with utilities for handling PyTorch data types in template-based code generation.

Add generate_forward_split.py and generate_backward_split.py scripts for generating SYCL embedding kernels from Jinja2 templates.

Set explicit Jinja autoescape policy in generator environments using select_autoescape with HTML/XML extensions only and non-HTML default. This keeps C++ template rendering behavior unchanged while satisfying Bandit B701 in codegen scripts. Replace regex assert checks in PTA parsing helpers with explicit validation that raises ValueError on malformed patterns. This preserves runtime safety under optimization and resolves Bandit B101 findings.

Add backward_utils.cpp and backward_utils.h with SYCL ports of FBGEMM backward pass utilities for embedding operations.

…mization Add vec4.h header implementing 4-element vectorized data structures for efficient memory access in embedding operations.

Add feature_gates module to enable/disable features at runtime via environment variables.

Add pt2_arg_utils.h header defining argument index enumerations for PyTorch 2 compiled embedding operations.

Add split_embeddings_cache_xpu.h header defining indices for UVM (Unified Virtual Memory) cache performance statistics.

…version Add stochastic_rounding.h header implementing stochastic rounding algorithms for float-to-half precision conversions in embedding operations.

Add weight_row.h header implementing abstractions for efficient access to embedding table weight rows with support for both direct table access and cache-resident data.

Add Jinja2 template for generating optimized SYCL forward kernels for small embedding dimensions (D <= 32).

Add embedding_forward_split_kernel_template.h as a Jinja2 template for generating the main SYCL forward kernel for no-bag (sequence) embeddings.

Add embedding_forward_nobag_unweighted_host_template.cpp as a Jinja2 template for generating the host-side dispatch function for no-bag (sequence) embedding forward passes.

…plit Add embedding_backward_split_kernel_templates.h as a Jinja2 template for generating SYCL backward kernels for no-bag unweighted embedding lookups. Generates both dense (gradient-only) and split (rowwise Adagrad optimizer) variants.

…ghted Add embedding_backward_nobag_unweighted_host_template.cpp as a Jinja2 template for generating the host-side backward dispatch function for no-bag (sequence) embedding passes. Generates both dense (gradient accumulation) and split (rowwise Adagrad) variants.

…ding Add embedding_forward_nobag_unweighted_pt2_wrapper_template.cpp as a Jinja2 template for generating PyTorch 2.0 compilation wrapper functions for split embedding nobag unweighted operations. The single template generates both forward and backward wrappers (controlled by the is_forward flag).

Add fbgemm_dense_lookups_ops.cpp implementing the dense embedding lookup operator for XPU with full autograd support. NOTE: Since the code in this file and the code in te split operator counterpart are from different templates, the code wasn't fused into a single template.

…ise Adagrad Add fbgemm_split_lookups_ops.cpp implementing the PT2-compatible autograd function for split embedding lookups with rowwise Adagrad optimizer on XPU. NOTE: Since the code in this file and the code in the dense operator counterpart are from different templates, the code wasn't fused into a single template.

Register the schemas required by the generated dense and split lookup implementations, including the PT2 wrappers and rowwise Adagrad backward paths. The rebase onto bc9c065 retains jagged_index_select_2d_forward from main in the shared TORCH_LIBRARY_FRAGMENT. Keeping both schema groups avoids reverting the newer jagged operator while preserving the complete training API added by this commit. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

dvrogozh · 2026-09-09T18:47:24Z

      run: |
        python -c "import torch; assert torch.xpu.is_available(), 'XPU (Intel GPU) is not available'"
-    - name: 'Test fbgemm-xpu'
+    - name: 'Test fbgemm-xpu with pristine FBGEMM'


What's the point of testing with pristine FBGEMM? As far as I understand this will simply fall to CPU.

Pristine in that case means unmodified/unpatched. It refers to the unmodified fbgemm-gpu-cpu==1.8.0 dependency, not to the execution backend. tests require an available XPU device, assert that the current accelerator is XPU, create the operator inputs on XPU, and exercise the XPU/AutogradXPU implementations registered by fbgemm_xpu. I'll just remove "pristine" to avoind missunderstanding

dvrogozh · 2026-09-09T18:55:55Z

+# Copyright (c) 2026 Intel Corporation. All Rights Reserved.
+# SPDX-License-Identifier: BSD-3-Clause
+
+cmake_minimum_required(VERSION 3.18)


This does not seem a stand-alone CMake file as you don't have project() in here. If so, there is no reason to repeat these statements here. I suggest to drop.

Agreed. This file is included through add_subdirectory() and inherits the minimum CMake version from the top-level project.

dvrogozh · 2026-09-09T18:56:38Z

+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+
+# --------------------------------------------------------------------------
+# Python / Torch / pybind11 — mirror the setup in src/fbgemm_xpu to ensure


This block is also a copy paste from the higher level CMake. Drop.

Agreed. I moved the shared Python, Torch, and pybind11 setup to the top-level CMake file and removed the duplicated blocks from both subdirectories.

dvrogozh · 2026-09-09T18:59:03Z

+Python3_add_library(_C_training MODULE WITH_SOABI ${training_sources})
+set_target_properties(_C_training PROPERTIES
+    PREFIX      ""
+    CXX_STANDARD 17


This tells to use same standard we already asked via CMAKE_CXX_STANDARD. Drop.

Agreed. Removed the redundant target-level CXX_STANDARD; _C_training inherits C++17 from CMAKE_CXX_STANDARD.

dvrogozh · 2026-09-09T19:01:56Z

+  return "UNKNOWN";
+}
+
+static bool env_check_key(const std::string& key) {


Where do we use that and for what? If we don't use - drop. If we do use - document environment variables in PR description.

env_check_key() is used by the feature-gate lookup. FBGEMM_TBE_ANNOTATE_KINETO_TRACE=1 enables Kineto annotations in the dense and split lookup paths. I will document this variable in the PR description. The retained TBE_V2 gate currently has no effective XPU behavior, so I will remove it.

dvrogozh · 2026-09-09T19:03:31Z

+  return !ten.has_value() || torch_tensor_on_sycl_xpu_check(ten.value());
+}
+
+#define TENSOR_ON_CPU_OR_MTIA(x)                                      \


Are we using this one?

This was a leftover from the upstream CPU/MTIA path copied into the XPU implementation. It is not needed here, so I removed the path together with the unused macro and helper.

dvrogozh · 2026-09-09T19:08:49Z

+  The lookup operators are currently supported through direct
+  `torch.ops.fbgemm` calls. The validated surface is:
+
+  - no-bag lookup (`PoolingMode.NONE`);


Can we have method names here listed?

I mapped each public lookup method to its validated functionality. Is this the method-level detail you were asking for?

dvrogozh · 2026-09-09T19:10:22Z

+The pristine FBGEMM 1.8.0 high-level training frontend does not expose
+`ComputeDevice.XPU`. Constructing
+`SplitTableBatchedEmbeddingBagsCodegen` for XPU is therefore unavailable in
+this release. Importing the plugin and calling the supported lookup operators


Ok... How FBGEMM tests work? Do they call torch.ops.fbgemm and we have this tested in CI or they do not?

The tests use torch.accelerator to select the current XPU device and create the input tensors on it. They then call the lookup methods directly through torch.ops.fbgemm. PyTorch dispatches these calls to the XPU/AutogradXPU implementations registered by fbgemm_xpu based on the tensor devices. The high-level FBGEMM TBE frontend is not involved. As discussed on sync - no ComputeDevice.XPU in scope of this PR

Remove local static assertions and exact auxiliary-list size checks. These checks are stricter than FBGEMM v1.8.0 and add review surface without an XPU-specific requirement; retain only the upstream argument layout. Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

mkrze · 2026-09-10T11:43:50Z

+                    const size_t grid_x = {%- if dense %}(total_B + kBlockDimY - 1) / kBlockDimY{%- else %}div_round_up(static_cast<size_t>(total_B), kBlockDimY){%- endif %};
+
+                    queue.submit([&](sycl::handler& cgh) {
+                        cgh.parallel_for<{{ mdesc | capitalize }}EmbeddingNobagCodegenForwardUnweightedSmallKernel<emb_t, cache_t, output_t, index_t, kSmallThreadGroupSize>>(
+                            sycl::nd_range<2>(
+                                sycl::range<2>(grid_x * kBlockDimY, sg_size),
+                                sycl::range<2>(kBlockDimY, sg_size)


Critical — forward launch fails at B*T = 2^26

This launch needs an effective grid cap. With the 512-item local range,
B*T = 2^26 produces 2^31 flattened work-items and fails on B60 with:

RuntimeError: Provided range and/or offset does not fit in int

I reproduced the same failure for D=4 (small path) and D=36 (general path).

The existing xpu_cap_grid_dim_x helper can cap both launch group counts. The
general kernel already has a grid-stride loop, but the small kernel currently
handles one b_t and returns, so it also needs a grid-stride loop before
capping is safe; otherwise valid entries beyond the capped grid are skipped.

Could we cap both launches, add the missing small-kernel grid stride, and add
large-grid regression cases for both D=4 and D=36? This shared template
generates both dense and split variants.

Thanks, this is a valid XPU launch-limit issue. The current general kernel already grid-strides, but neither launch is capped and the small kernel handles only one b_t. The same no-bag fix has also landed in newer upstream FBGEMM in PR #6071. I will apply the existing XPU grid cap to both launches, add the missing small-kernel grid-stride loop, and cover the D=4 and D=36 paths with large-grid regression tests.

mkrze · 2026-09-10T11:46:01Z

+
+bool check_feature_gate_key(const std::string& key) {
+  // Cache feature flags to avoid repeated environment lookups.
+  static std::map<std::string, bool> feature_flags_cache;


High — feature-gate cache has a C++ data race

This process-wide map is accessed without synchronization. Concurrent first
calls to the dense and split operators can execute find() and insert() on
the same std::map, which is undefined behavior.

The mutex added in 07e62c8 was reverted in 899d70d, so the race is present
again. The unused TBE_V2 gate has since been removed, but the retained
TBE_ANNOTATE_KINETO_TRACE gate still reaches this shared cache from four
dense/split call sites. Could we restore synchronized access or replace the map
with thread-safe one-time initialization?

Thanks, the race is technically possible during concurrent first calls. However, this cache behavior is inherited from FBGEMM v1.8.0 and is still present in upstream FeatureGate::lookup. I previously added a local mutex in 07e62c8, but intentionally reverted it in 899d70d to avoid introducing an XPU-only fix unrelated to this port. I propose addressing this upstream first and then syncing the upstream fix here. I do not want to fix upstream until directly asked by architect.

mkrze · 2026-09-10T11:46:44Z

+        return SplitNoBagLookupFunctionDenseOpXPU::apply(
+            output_dtype, dev_weights, weights_offsets, max_D, hash_size_cumsum,
+            total_hash_size_bits, indices, offsets)[0];


High — mixed dimensions silently return wrong rows

The no-bag path ignores D_offsets and total_D, then uses max_D as every
table's row stride. I reproduced incorrect output on B60 with two tables using
D=[4,8]: requesting row 1 of table 0 returned row 0 of table 1.

The README limits support to uniform dimensions, but this API does not reject
mixed dimensions, and the dense schema defaults mixed_D=True. Could we either
implement per-table strides or validate that every D_offsets delta equals
max_D and raise a clear unsupported-layout error? The split entry point needs
the same validation.

Thanks, mixed dimensions are unsupported by the upstream no-bag kernels, which intentionally receive a single D. The upstream split frontend rejects this layout before dispatch, but our direct torch.ops path bypasses that frontend validation. I will add a lightweight uniform-dimension check to both XPU entry points and regression tests. Implementing per-table strides is outside this PR’s scope.

mkrze · 2026-09-10T11:47:28Z

+TORCH_LIBRARY_IMPL(fbgemm, XPU, m) {
+{%- if is_forward %}
+    m.impl("{{ mdesc }}_embedding_nobag_codegen_forward_unweighted_pt2_wrapper",
+        &fbgemm_xpu::{{ mdesc }}_embedding_nobag_codegen_forward_unweighted_pt2_xpu_wrapper);


High — PT2 claim fails FakeTensor tracing

The split operator is tagged PT2-compliant, but a
torch.compile(..., backend="eager", fullgraph=True) probe fails during
FakeTensor tracing:

Unsupported: Operator does not support running with fake tensors fbgemm.split_embedding_nobag_codegen_forward_unweighted_pt2_wrapper.default

The dense path fails similarly on its internal forward helper. These internal
helpers have XPU implementations but no Meta/Fake behavior.

Could we add Meta/Fake implementations plus XPU torch.compile
forward/backward coverage? If compile support is outside this PR's scope, the
PT2-support claim should be narrowed to the eager functionality that is
actually validated.

Thanks. torch.compile/FakeTensor support is outside the currently validated scope. I removed the unsupported PT2-compliant tag and documented that the lookup paths currently support eager execution only.

mkrze · 2026-09-10T11:48:28Z

+                if (threadIdx_x == 0) {
+                    sycl::atomic_fence(sycl::memory_order::acq_rel, sycl::memory_scope::device);
+                    counter = xpuAtomicAdd(&grad_accum_counter_[really_long_run_id], -1);
+                }
+                counter = sycl::group_broadcast(sg, counter, 0);


High — multi-CTA completion lacks valid memory ordering

This completion counter does not establish the required SYCL memory ordering.
Other subgroup lanes atomically update temp_grad_accum_ptr, but there is no
subgroup barrier before lane 0 decrements the counter, and xpuAtomicAdd uses
relaxed ordering. The lane-0 fence does not order writes performed by the other
work-items or create a release/acquire hand-off between CTAs.

A nonuniform 1025-repeat case passed in this run, but that does not make the
missing memory-model edge valid; the result relies on lockstep and visibility
behavior SYCL does not guarantee. Could we synchronize the contributing
subgroup before publishing completion and use device-scope release/acquire
semantics before the final CTA loads the accumulation buffer?

Thanks, agreed. The current SYCL translation is not equivalent to the upstream CUDA completion protocol: the lane-0 fence does not order relaxed atomic writes from the other subgroup lanes. I will add subgroup synchronization before publishing completion, use a device-scope acquire-release counter operation, synchronize before the final accumulation-buffer reads, and add multi-CTA regression coverage.

Cap both no-bag forward launch paths and grid-stride the small kernel so B*T=2^26 stays within the SYCL launch range. Add D=4 and D=36 regressions. Co-authored-by: Cursor <cursoragent@cursor.com>

Validate the upstream no-bag uniform-dimension constraint at the direct XPU operator boundary to prevent max_D from selecting incorrect rows. Cover both dense and split entry points. Co-authored-by: Cursor <cursoragent@cursor.com>

Remove the unsupported PT2-compliant tag and document that FakeTensor tracing and torch.compile are outside the currently validated lookup surface. Co-authored-by: Cursor <cursoragent@cursor.com>

Synchronize subgroup accumulation before publishing completion and use a device-scope acquire-release counter handoff before the final CTA reads shared gradients. Add a 1025-repeat regression. Co-authored-by: Cursor <cursoragent@cursor.com>

Use TORCH_CHECK for the uniform-dimension predicate because TORCH_CHECK_EQ does not accept a custom diagnostic argument. Co-authored-by: Cursor <cursoragent@cursor.com>

Match the large-grid expected tensor rank and pass the required output dtype to the dense mixed-dimension rejection test. Co-authored-by: Cursor <cursoragent@cursor.com>

dvrogozh · 2026-09-10T15:09:59Z

  - [`permute_1D_sparse_data`][op-permute_1D_sparse_data]
  - [`permute_2D_sparse_data`][op-permute_2D_sparse_data]

+* Training lookup operators supported through direct calls:


Can we add link to FBGEMM documentation?

dvrogozh · 2026-09-10T16:33:32Z

+set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)


Move these 2 lines to the top level CMake script: https://github.com/intel/torchlib-xpu/blob/main/packages/fbgemm-xpu/CMakeLists.txt

dvrogozh · 2026-09-10T16:34:56Z

 set(CMAKE_CXX_STANDARD 17)
 set(CMAKE_CXX_STANDARD_REQUIRED ON)

-set(PYBIND11_FINDPYTHON ON)


Drop cmake_minimum_required and move set() directives to top level CMake script: https://github.com/intel/torchlib-xpu/blob/main/packages/fbgemm-xpu/CMakeLists.txt

aagalleg · 2026-09-10T18:25:30Z

+      run: |
+        rm -rf .git/rebase-apply
+        git am ../torchlib-xpu/packages/fbgemm-xpu/patches/0001-Add-XPU-support-to-fbgemm-tests.patch
+        unexpected=$(git diff --name-only v1.8.0 HEAD | grep -v '^fbgemm_gpu/test/' || true)


This pins the dependency to v1.8.0. Once v1.9.0 is released (FBGEMM is already at v1.9.0-rc1) CI may start testing against an outdated assumption, which could lead to unexpected failures or mask compatibility issues.

It is probably a good idea to have a separate file pointing to FBGEMM used version, so that we don't have to hard code the version each time and just read the used version from the file.

Handling separate file gives its own complications which I would better avoid. However, the rational part is to minimize places where we hardcode versions. In this sense, let's add fbgemm version to the workflow env block. And use it when we need refer fbgemm version. I.e. also modify this place:

torchlib-xpu/.github/workflows/ci.yml

Line 442 in 099a1c9

ref: v1.8.0

aagalleg · 2026-09-10T21:51:45Z

The values of AUX_TENSOR_SIZE, AUX_BOOL_SIZE, AUX_INT_SIZE, and AUX_FLOAT_SIZE have been changed from their original values. The current pt2_arg_utils.h is based on the FBGEMM file of the same name (produced from training/pt2/pt2_arg_utils_template.h after the codegen stage). Please restore the original values.

aagalleg · 2026-09-10T21:55:04Z

+namespace fbgemm_xpu::config {
+
+#define ENUMERATE_ALL_FEATURE_FLAGS X(TBE_ANNOTATE_KINETO_TRACE)
+// X(EXAMPLE_FEATURE_FLAG)


Remove commented code.

aagalleg · 2026-09-10T23:59:13Z

+                    {%- if not dense %}
+                    DISPATCH_KERNEL_FOR_CACHE_CASE(use_lxu_cache, [&] {
+                        {%- endif %}
+                        const size_t local_x = kThreadGroupSize;{%- if dense %} {%- endif %}


Remove "{%- if dense %} {%- endif %}" at the end of the line.

aagalleg · 2026-09-11T00:00:06Z

+                        {%- if dense %}
+
+                        {%- else %}
+
+                        {%- endif %}


Remove empty if else block.

aagalleg · 2026-09-11T01:44:33Z

Are this test based on FBGEMM TBE tests?

aagalleg · 2026-09-11T01:49:34Z

+            package_path.is_relative_to(environment_packages),
+            f"{package_path} is not installed under {environment_packages}",
+        )
+        self.assertEqual(version("fbgemm-gpu-cpu"), "1.8.0")


The version is hardcoded. This will fail when using v1.9.0 or subsequent versions.
It is probably a good idea to have a separate file pointing to FBGEMM used version, so that we don't have to hard code the version each time and just read the used version from the file.

aagalleg mentioned this pull request Aug 28, 2026

Add CMake build system and tests for the lookup operators #85

Closed

molintc force-pushed the feat/clean_lookup_operators_build branch from 8638f0e to 2c6fa32 Compare August 31, 2026 09:37

dvrogozh reviewed Sep 3, 2026

View reviewed changes

aagalleg and others added 25 commits September 3, 2026 18:51

feat(fbgemm_xpu): add Jinja environment and helper functions for code…

8d6e614

… generation Add jinja_environment.py module to support template-based code generation for FBGEMM-XPU kernels.

feat(fbgemm-xpu): add CodeTemplate class for template rendering

68ade45

Add common.py module with CodeTemplate class that provides functionality for loading Jinja2 templates, rendering them with context variables, and writing generated files with appropriate headers.

feat(fbgemm-xpu): add PyTorch type utilities for code generation

e7d0547

Add torch_type_utils.py module with utilities for handling PyTorch data types in template-based code generation.

feat(fbgemm-xpu): add argument parser for generation scripts

9fbaf3d

feat(fbgemm-xpu): fix license for jinja_environment.py

1d48fc5

feat(fbgemm-xpu): add forward and backward kernel generation scripts

28adfed

Add generate_forward_split.py and generate_backward_split.py scripts for generating SYCL embedding kernels from Jinja2 templates.

chore(codegen): add Intel copyright header to jinja environment script

88a1253

style(fbgemm-xpu): reorder imports in __init__.py

3b41230

feat(fbgemm-xpu): add SYCL backward pass utility functions

da77670

Add backward_utils.cpp and backward_utils.h with SYCL ports of FBGEMM backward pass utilities for embedding operations.

feat(fbgemm-xpu): add Vec4T vectorized data structure for memory opti…

f27818b

…mization Add vec4.h header implementing 4-element vectorized data structures for efficient memory access in embedding operations.

feat(fbgemm-xpu): add feature gates system for runtime configuration

e1b7db3

Add feature_gates module to enable/disable features at runtime via environment variables.

feat(fbgemm-xpu): add PT2 argument index enumerations

a5c0f15

Add pt2_arg_utils.h header defining argument index enumerations for PyTorch 2 compiled embedding operations.

feat(fbgemm-xpu): add UVM cache statistics index enumeration

d22e1d4

Add split_embeddings_cache_xpu.h header defining indices for UVM (Unified Virtual Memory) cache performance statistics.

feat(fbgemm-xpu): add stochastic rounding utilities for precision con…

33e426e

…version Add stochastic_rounding.h header implementing stochastic rounding algorithms for float-to-half precision conversions in embedding operations.

feat(fbgemm-xpu): add WeightRow abstraction for embedding table access

d8b19a5

Add weight_row.h header implementing abstractions for efficient access to embedding table weight rows with support for both direct table access and cache-resident data.

feat(fbgemm-xpu): add small embedding forward kernel template

e78eb11

Add Jinja2 template for generating optimized SYCL forward kernels for small embedding dimensions (D <= 32).

feat(fbgemm-xpu): add main embedding forward kernel Jinja2 template

19cc6ee

Add embedding_forward_split_kernel_template.h as a Jinja2 template for generating the main SYCL forward kernel for no-bag (sequence) embeddings.

feat(fbgemm-xpu): add host dispatch template for nobag forward pass

0403082

Add embedding_forward_nobag_unweighted_host_template.cpp as a Jinja2 template for generating the host-side dispatch function for no-bag (sequence) embedding forward passes.

molintc marked this pull request as ready for review September 9, 2026 13:34

molintc requested review from flezaalv, luis-real and mkrze as code owners September 9, 2026 13:34

dvrogozh reviewed Sep 9, 2026

View reviewed changes

molintc marked this pull request as draft September 10, 2026 08:47

molintc and others added 13 commits September 10, 2026 10:50

chore(codegen): add Intel copyrights

f433de9

Co-authored-by: Cursor <cursoragent@cursor.com>

ci(fbgemm-xpu): simplify test name

759e323

Co-authored-by: Cursor <cursoragent@cursor.com>

build(codegen): drop redundant CMake minimum

32bbacf

Co-authored-by: Cursor <cursoragent@cursor.com>

build(fbgemm-xpu): centralize Python setup

8cf3a70

Co-authored-by: Cursor <cursoragent@cursor.com>

build(fbgemm-xpu): centralize compiler check

9b9b2d6

Co-authored-by: Cursor <cursoragent@cursor.com>

build(fbgemm-xpu): centralize XPU configuration

88e9b31

Co-authored-by: Cursor <cursoragent@cursor.com>

build(codegen): drop redundant C++ standard

dffc9b1

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(fbgemm-xpu): drop unused TBE V2 gate

f7ff6e5

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(fbgemm-xpu): keep split lookup XPU-only

9b50055

Co-authored-by: Cursor <cursoragent@cursor.com>

docs(fbgemm-xpu): list lookup entry points

15a1351

Co-authored-by: Cursor <cursoragent@cursor.com>

docs(fbgemm-xpu): map lookup methods to coverage

816c4ab

Co-authored-by: Cursor <cursoragent@cursor.com>

docs(fbgemm-xpu): clarify lookup test path

970cd5e

Co-authored-by: Cursor <cursoragent@cursor.com>

molintc marked this pull request as ready for review September 10, 2026 10:20

mkrze requested changes Sep 10, 2026

View reviewed changes

molintc and others added 6 commits September 10, 2026 14:42

fix(fbgemm-xpu): cap lookup forward grids

8b38808

Cap both no-bag forward launch paths and grid-stride the small kernel so B*T=2^26 stays within the SYCL launch range. Add D=4 and D=36 regressions. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(fbgemm-xpu): reject mixed no-bag dimensions

a853925

Validate the upstream no-bag uniform-dimension constraint at the direct XPU operator boundary to prevent max_D from selecting incorrect rows. Cover both dense and split entry points. Co-authored-by: Cursor <cursoragent@cursor.com>

docs(fbgemm-xpu): scope lookups to eager mode

f4d0e62

Remove the unsupported PT2-compliant tag and document that FakeTensor tracing and torch.compile are outside the currently validated lookup surface. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(fbgemm-xpu): order multi-CTA completion

6757855

Synchronize subgroup accumulation before publishing completion and use a device-scope acquire-release counter handoff before the final CTA reads shared gradients. Add a 1025-repeat regression. Co-authored-by: Cursor <cursoragent@cursor.com>

fix(fbgemm-xpu): use message-capable dimension check

d12958e

Use TORCH_CHECK for the uniform-dimension predicate because TORCH_CHECK_EQ does not accept a custom diagnostic argument. Co-authored-by: Cursor <cursoragent@cursor.com>

test(fbgemm-xpu): fix lookup regression assertions

c05ae1f

Match the large-grid expected tensor rank and pass the required output dtype to the dense mixed-dimension rejection test. Co-authored-by: Cursor <cursoragent@cursor.com>

dvrogozh reviewed Sep 10, 2026

View reviewed changes

aagalleg commented Sep 11, 2026

View reviewed changes

		set(CMAKE_CXX_STANDARD 17)
		set(CMAKE_CXX_STANDARD_REQUIRED ON)

Uh oh!

Conversation

aagalleg commented Aug 28, 2026 • edited by molintc Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Build and code generation

Correctness fixes

Supported functionality

Testing model

Compatibility boundary

Uh oh!

dvrogozh commented Sep 3, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

molintc Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

molintc Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

molintc Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

molintc Sep 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

aagalleg commented Aug 28, 2026 •

edited by molintc

Loading

molintc Sep 10, 2026 •

edited

Loading

molintc Sep 10, 2026 •

edited

Loading

molintc Sep 10, 2026 •

edited

Loading

molintc Sep 10, 2026 •

edited

Loading