NVIDIA / NVIDIA/recsys-examples
[FEA] Slow unbucketize permute operation in SequenceEmbeddingsAllToAll for row-wise sharding
Open
@z52527 is already working on this.
Since Feb 6, 2026.
dynamicemb
enhancement
- Dominant language
- Python
- Stars
- 322
- Forks
- 83
- Avg merge
- 3d 3h
- Merged PRs (30d)
- 8
Description
Background
Currently, DynamicEmb has custom input_dist implementation (RwSparseFeaturesDist in input_dist.py) but still relies on TorchRec's original output_dist implementation. This causes:
- Performance issue: The
unbucketize_permuteoperation in TorchRec's output distribution is slow, especially for non-contiguous distribution patterns (e.g., round-robin) - Limited customization: Cannot optimize the output distribution without modifying TorchRec source code
Objective
Port TorchRec's output distribution classes to DynamicEmb library, enabling future performance optimizations.
Tasks
PR 1: Port output distribution classes to DynamicEmb
- Create
dynamicemb/output_dist.pywith:RwSequenceEmbeddingDistRwPooledEmbeddingDist
- Update
dynamicemb/planner/rw_sharding.pyto overridecreate_output_dist()methods - Verify with existing tests (
test_sequence_embedding_fw.py,test_pooled_embedding_fw.py)
PR 2: Optimize unbucketize permute with custom kernel
- Design optimized data format for permute tensor
- Implement CUDA kernel for efficient unbucketize operation
- Integrate with
output_dist.py - Benchmark and validate performance improvement
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.