Skip to content

feat(metax): route basic Llama through canonical InfiniOps - #1498

Merged
wooway777 merged 1 commit into
mainfrom
feat/metax-canonical-llama-infiniops
Aug 13, 2026
Merged

feat(metax): route basic Llama through canonical InfiniOps#1498
wooway777 merged 1 commit into
mainfrom
feat/metax-canonical-llama-infiniops

Conversation

@voltjia

@voltjia voltjia commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

What

  • Route the MetaX basic Llama/Qwen RoPE path through the canonical InfiniOps
    RotaryEmbedding API instead of the legacy InfiniOP RoPE C API.
  • Enable the generated PyTorch Argmax provider in the external InfiniOps
    build on MetaX and use it for greedy sampling (top_k=1).
  • Keep public InfiniCore C++ and Python interfaces unchanged.

Migration

InfiniCore route Previous MetaX route Canonical InfiniOps route Alignment target
RoPE Legacy infiniopRoPE C API RotaryEmbedding MetaX provider vLLM rotary_embedding
Greedy sampling (top_k=1) Legacy random-sample path Generated Argmax PyTorch provider PyTorch torch.argmax

The 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-main MetaX build. The Llama
runtime changes themselves are independent.

Screenshots: N/A (backend integration only).

Validation

Run on ssh metax in infiniops-ci/metax:latest on one MetaX C550, based on
main at a12eb953 plus the stacked build fix:

  • Full InfiniCore MetaX release build passed with WITH_TORCH=ON and the
    generated argmax wrapper enabled.
  • python3 test/infinicore/nn/rope.py --metax --debug passed 6/6 cases.
  • python3 test/infinicore/ops/random_sample.py --metax --debug passed 32/32
    cases, including FP16/BF16 top_k=1 return-value and explicit-output cases.
  • Qwen3-0.6B generated 64 greedy tokens and exited 0 with:
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=64
  • An LD_PRELOAD trap covering the legacy InfiniOP execution entries used by
    the basic Transformer path did not fire during the 64-token prefill/decode
    run. The same trap reproduced the previous infiniopRoPE fallback before
    this change.
  • Source scans found no InfiniLM-suffixed InfiniOps references in InfiniCore.
  • The 64-token decoded output was unchanged from the pre-migration MetaX run.
  • Against PyTorch Transformers greedy generation, the first 54 generated
    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 --Werror and git diff --check passed.

Base automatically changed from fix/metax-recurrent-gated-delta-rule to main August 13, 2026 09:29
@wooway777
wooway777 requested a review from a team August 13, 2026 09:29
@wooway777
wooway777 force-pushed the feat/metax-canonical-llama-infiniops branch from 8e17b04 to 4977420 Compare August 13, 2026 09:29
@wooway777
wooway777 merged commit 7f9e55e into main Aug 13, 2026
10 checks passed
@wooway777
wooway777 deleted the feat/metax-canonical-llama-infiniops branch August 13, 2026 09:30
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.

2 participants