Skip to content

refactor(ops): route Llama paths through InfiniOps - #1508

Draft
voltjia wants to merge 1 commit into
mainfrom
refactor/use-infiniops-for-llama-routes
Draft

refactor(ops): route Llama paths through InfiniOps#1508
voltjia wants to merge 1 commit into
mainfrom
refactor/use-infiniops-for-llama-routes

Conversation

@voltjia

@voltjia voltjia commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Enable the canonical InfiniOps flash_attn_varlen_func and flash_attn_with_kvcache providers for MetaX prefill and decode.
  • Replace the legacy select_last_token_hidden provider on NVIDIA and MetaX with canonical InfiniOps add and index_select calls.
  • Select the required native, generated ATen, and linked implementations through the existing InfiniOps ops.json mechanism.

Motivation

The basic Llama/Qwen inference path still bypassed canonical InfiniOps for MetaX FlashAttention and last-token selection. This change keeps those paths in InfiniCore but routes their operator work through open-source-aligned InfiniOps interfaces, without calling any Infinilm-suffixed InfiniOps operator.

Dependency

Implementation

  • multi_head_attention_varlen accepts MetaX in its existing InfiniOps eligibility check.
  • mha_kvcache accepts MetaX in its existing InfiniOps eligibility check.
  • select_last_token_hidden computes input_offsets[1:] - 1 with canonical add, then gathers rows with canonical index_select.
  • The scalar one is initialized with a host-to-device copy during planning; no legacy InfiniOp operator is used by this route.
  • The external InfiniOps build writes a build-local ops.json: FlashAttention uses linked slot 16, argmax and index_select use generated ATen slot 8, and operators that also have unrelated linked providers remain on native slot 0.

Validation

  • Full infinicore_cpp_api and _infinicore builds passed on MetaX with the dependent InfiniOps branch.
  • The generated linked manifest contains only the MetaX flash_attn_varlen_func and flash_attn_with_kvcache source files.
  • Qwen3-0.6B paged FlashAttention inference passed after the ops.json migration with a 2-token smoke run.
  • Earlier runtime tracing of the same provider code hit:
    • Operator<FlashAttnVarlenFunc, Device::Type::kMetax, 16> through the linked Torch adapter.
    • select_last_token_hidden_impl::infiniops::run.
    • Operator<FlashAttnWithKvcache, Device::Type::kMetax, 16> through the linked Torch adapter.
  • A source scan over src/infinicore and include/infinicore returned no Infinilm or _infinilm references.
  • A full repository scan returned no INFINI_OPS_LINKED_OPS references.
  • All changed C++ files pass clang-format 16.0.6 --dry-run --Werror with the repository .clang-format.
  • git diff --check origin/main...HEAD passes.

Scope

This PR only changes the basic Llama/Qwen routes described above. Other models, layouts, and fallback behavior are intentionally unchanged.

@voltjia
voltjia force-pushed the refactor/use-infiniops-for-llama-routes branch from e02976d to 0bd84a6 Compare August 19, 2026 04:38
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