Skip to content

perf(metax): vectorize contiguous embedding lookups - #997

Open
JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:masterfrom
JoeZhang-0x000:perf/metax-vectorized-embedding
Open

JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:masterfrom
JoeZhang-0x000:perf/metax-vectorized-embedding

Conversation

@JoeZhang-0x000

Copy link
Copy Markdown

Contiguous MetaX embedding lookups at hidden size 4096 currently traverse the shared generic copy path. Add a MetaX override that copies each selected row with a 256-thread block and aligned vector loads/stores. Other shapes, layouts, unaligned pointers and max_norm continue through the existing implementation; other CUDA backends are unchanged.

Add coverage for the fast path, padded vocabulary rows, strided indices/weights/output, int32/int64 indices and max_norm fallback.

Validation on MetaX C550, MACA 3.8.0.23, PyTorch 2.10.0+metax3.8.0.7:

  • Built master 4c014ca with this change and the two separately proposed MetaX RMSNorm changes; 301 focused Embedding/FusedAddRmsNorm/RmsNorm tests passed, plus the smoke selection (2 tests).
  • All three patches also apply without fuzz to the plugin's locked InfiniOps 8c2f70a.
  • Prior CUDA Graph measurements of this kernel against the generic path on 8c2f70a: 64 tokens, BF16, hidden=4096, 215.72 -> 4.09 us; 8192 tokens, 656.32 -> 96.29 us. These are per-call device times, not model throughput.

NVIDIA hardware was not available for testing. The change is confined to the MetaX specialization.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant