Skip to content

perf(metax): reduce RMSNorm blocks for small head dimensions - #999

Open
JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:masterfrom
JoeZhang-0x000:perf/metax-small-rms-norm-blocks
Open

JoeZhang-0x000 wants to merge 1 commit into
InfiniTensor:masterfrom
JoeZhang-0x000:perf/metax-small-rms-norm-blocks

Conversation

@JoeZhang-0x000

Copy link
Copy Markdown

MetaX's device-wide RMSNorm block size leaves idle reduction threads for small Q/K head dimensions. Select 128 threads when the normalized dimension is <=128 on MetaX. Preserve the current block-size selection for larger rows and all other CUDA backends; the computation kernel is unchanged.

Add packed-QKV-stride cases for Q/K head dimension 128 and a dimension-129 case for the unchanged branch.

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 Embedding/FusedAddRmsNorm changes; 301 focused operator tests passed, plus the smoke selection (2 tests).
  • All three patches also apply without fuzz to the plugin's locked InfiniOps 8c2f70a.
  • Prior BF16 Graph measurements of the same launch change on 8c2f70a: [64,16,128], 8.81 -> 6.06 us; [8192,16,128], 741.41 -> 369.40 us. These are per-call device times, not model throughput.

NVIDIA hardware was not available for testing; the MetaX constexpr guard preserves its current selection.

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