NVIDIA / NVIDIA/recsys-examples

[FEA] Slow unbucketize permute operation in SequenceEmbeddingsAllToAll for row-wise sharding

Open
#296 0 comments 0 reactions 1 assignee View on GitHub

@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:

  1. Performance issue: The unbucketize_permute operation in TorchRec's output distribution is slow, especially for non-contiguous distribution patterns (e.g., round-robin)
  2. 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.py with:
    • RwSequenceEmbeddingDist
    • RwPooledEmbeddingDist
  • Update dynamicemb/planner/rw_sharding.py to override create_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

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.