feat(metax): route basic Llama through canonical InfiniOps - #1498
Merged
Conversation
wooway777
approved these changes
Aug 13, 2026
wooway777
force-pushed
the
feat/metax-canonical-llama-infiniops
branch
from
August 13, 2026 09:29
8e17b04 to
4977420
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
RotaryEmbeddingAPI instead of the legacy InfiniOP RoPE C API.Argmaxprovider in the external InfiniOpsbuild on MetaX and use it for greedy sampling (
top_k=1).Migration
infiniopRoPEC APIRotaryEmbeddingMetaX providerrotary_embeddingtop_k=1)ArgmaxPyTorch providertorch.argmaxThe other basic Llama routes were already migrated to canonical InfiniOps by
#1483. This PR adds the missing MetaX selection for the two routes above.
Scope
This PR targets the basic Llama architecture path used by Qwen3 with paged
FlashAttention and greedy sampling on MetaX. Non-greedy sampling, other model
families, and other attention layouts remain out of scope.
This PR is stacked on the MetaX recurrent gated delta rule build fix because
that fix is required for a complete current-
mainMetaX build. The Llamaruntime changes themselves are independent.
Screenshots: N/A (backend integration only).
Validation
Run on
ssh metaxininfiniops-ci/metax:lateston one MetaX C550, based onmainata12eb953plus the stacked build fix:WITH_TORCH=ONand thegenerated
argmaxwrapper enabled.python3 test/infinicore/nn/rope.py --metax --debugpassed 6/6 cases.python3 test/infinicore/ops/random_sample.py --metax --debugpassed 32/32cases, including FP16/BF16
top_k=1return-value and explicit-output cases.python3 examples/test_infer.py \ --device=metax \ --model=/workspace/Qwen3-0.6B \ --enable-paged-attn \ --attn=flash-attn \ --disable-prefix-caching \ --top-k=1 \ --max-new-tokens=64LD_PRELOADtrap covering the legacy InfiniOP execution entries used bythe basic Transformer path did not fire during the 64-token prefill/decode
run. The same trap reproduced the previous
infiniopRoPEfallback beforethis change.
tokens matched exactly and token 55 diverged. This is an existing inference
precision difference, not introduced by switching the MetaX RoPE route; the
migration does not claim full token-for-token Transformers parity.
clang-format 16.0.6 --dry-run --Werrorandgit diff --checkpassed.