Skip to content

refactor(ops): route basic Llama through canonical InfiniOps - #1483

Merged
voltjia merged 10 commits into
mainfrom
refactor/remove-infinilm-infiniops-adapters
Aug 11, 2026
Merged

refactor(ops): route basic Llama through canonical InfiniOps#1483
voltjia merged 10 commits into
mainfrom
refactor/remove-infinilm-infiniops-adapters

Conversation

@voltjia

@voltjiavoltjia commented Aug 10, 2026

Copy link
Copy Markdown
Collaborator

What

  • Route the NVIDIA basic Llama prefill, decode, RoPE, KV-cache update,
    greedy-sampling, residual RMSNorm, and SwiGLU paths through canonical
    InfiniOps C++ APIs.
  • Include the paged-attention decode and prefill changes formerly reviewed in
    refactor(ops): migrate paged attention to canonical InfiniOps #1481. This PR now targets main; refactor(ops): migrate paged attention to canonical InfiniOps #1481 is closed as superseded.
  • Remove InfiniLM-suffixed InfiniOps references and legacy InfiniOP C API
    adapters from the validated InfiniCore path.
  • Use the generated PyTorch Argmax provider on NVIDIA. The InfiniOps external
    build enables only argmax, InfiniCore selects provider implementation 8,
    and the gitlink remains at 47c1c496, before InfiniOps issue/867 fix cpu malloc #918.
  • Keep the public InfiniCore C++ and Python APIs unchanged. Compatibility
    adaptation remains inside InfiniCore providers.

Migration

InfiniCore routePrevious dependencyCanonical InfiniOps routeProvider adaptation
KV-cache insertionKvCachingInfinilm / legacy InfiniOP adapterReshapeAndCacheFlashMap InfiniCore cache metadata to the vLLM-style cache update contract.
Paged prefillPagedAttentionPrefillInfinilm / legacy InfiniOP adapterFlashAttnVarlenFuncMap the basic Llama varlen layout to the FlashAttention public contract.
Paged decodePagedAttentionInfinilm fallbackFlashAttnWithKvcacheUse the canonical path for the validated basic Llama cache layout.
RoPERotaryEmbeddingInfinilm / legacy InfiniOP adapterRotaryEmbeddingPreserve the public InfiniCore argument order while calling the vLLM-aligned operator.
Greedy sampling (top_k=1)RandomSampleInfinilm / legacy InfiniOP adapterArgmaxReduce final logits along the vocabulary dimension with the generated PyTorch provider.
Residual RMSNormdeprecated AddRmsNormFusedAddRmsNormCopy only for non-aliasing public outputs, then use the canonical in-place vLLM contract. The basic Llama path aliases and adds no extra copies.
SwiGLUdeprecated Swiglu(a, b)Copy + SiluAndMulPack [gate=b, up=a], then evaluate vLLM's silu(gate) * up contract.

Alignment Evidence

The selected targets follow InfiniOps' alignment policy: prefer a public
PyTorch Python API, then a serving-framework wrapper, then ONNX or a
library-owned public interface. InfiniOps applies the accepted C++ ordering
adaptation: tensor inputs, attributes, then writable outputs.

InfiniOps APIUpstream contractAlignment target
Embeddingembedding(input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)PyTorch torch.nn.functional.embedding
RmsNormrms_norm(out, input, weight, epsilon); InfiniOps moves out to the output groupvLLM _custom_ops.rms_norm
GemmGemm(A, B, C?, alpha, beta, transA, transB) -> YONNX Gemm
RotaryEmbeddingrotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox, rope_dim_offset=0, inverse=False)vLLM _custom_ops.rotary_embedding
ReshapeAndCacheFlashreshape_and_cache_flash(key, value, key_cache, value_cache, slot_mapping, kv_cache_dtype, k_scale, v_scale)vLLM _custom_ops.reshape_and_cache_flash
FlashAttnVarlenFuncflash_attn_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_k, max_seqlen_q, max_seqlen_k, ..., return_attn_probs=False)FlashAttention public interface
FlashAttnWithKvcacheflash_attn_with_kvcache(q, k_cache, v_cache, k=None, v=None, ..., return_softmax_lse=False)FlashAttention public interface
FusedAddRmsNormfused_add_rms_norm(input, residual, weight, epsilon)vLLM _custom_ops.fused_add_rms_norm
CopyTensor.copy_(src, non_blocking=False); InfiniOps represents the receiver as trailing outPyTorch Tensor.copy_
SiluAndMulSiluAndMul.forward(x), where silu(x[..., :d]) * x[..., d:]vLLM SiluAndMul
Argmaxtorch.argmax(input, dim=None, keepdim=False)PyTorch torch.argmax

Scope

This PR targets the NVIDIA basic Llama architecture flow used by InfiniLM with
paged FlashAttention and greedy sampling. Qwen3 uses this same validated path.
Non-greedy sampling, unsupported attention layouts, other model families, and
other devices remain out of scope.

InfiniOps #918 is not a dependency of this PR. It remains an optional native
NVIDIA Argmax implementation; this PR uses the generated PyTorch provider.

Screenshots: N/A (backend integration only).

Validation

Run on ssh nvidia in accelerator-dev/nvidia:latest on one NVIDIA
A100-SXM4-80GB:

  • git diff --check origin/main...HEAD passed.
  • clang-format 16.0.6 --dry-run --Werror passed for every modified C++ file.
  • ruff 0.15.20 check and ruff 0.15.20 format --check passed for the modified
    Python test.
  • The focused InfiniOps/InfiniCore build passed with WITH_TORCH=ON and
    INFINI_OPS_TORCH_OPS=argmax. The built library references
    at::argmax_out; no native NVIDIA Argmax source from issue/867 fix cpu malloc #918 is present.
  • Source scans found no InfiniLM-suffixed InfiniOps references in the validated
    InfiniCore path.
  • An LD_PRELOAD trap covering 607 legacy infiniop* C API symbols across 154
    descriptor families did not fire during inference.
  • TinyLlama-1.1B-Chat inference exited 0 on the same canonical path.
  • Qwen3-0.6B inference exited 0 and produced 64 coherent greedy tokens. It
    matched PyTorch for the first 6 generated tokens, then selected " just"
    where PyTorch selected " said". At that first divergence, PyTorch BF16
    logits are exactly tied at 22.875 for both tokens, so exact sequence
    equality is not a stable precision criterion for this prompt.
  • Qwen3-4B inference exited 0 and its 64 generated tokens matched PyTorch
    Transformers exactly (64/64), including the decoded text.

Commands used for both Qwen3 models:

python3 examples/test_infer.py \
--device=nvidia \
--model=/tmp/standard-llama/Qwen3-0.6B \
--enable-paged-attn \
--attn=flash-attn \
--top-k=1 \
--top-p=1.0 \
--temperature=1.0 \
--max-new-tokens=64 \
--prompt=Hello
python3 examples/test_infer.py \
--device=nvidia \
--model=/tmp/standard-llama/Qwen3-4B \
--enable-paged-attn \
--attn=flash-attn \
--top-k=1 \
--top-p=1.0 \
--temperature=1.0 \
--max-new-tokens=64 \
--prompt=Hello

@voltjiavoltjia changed the title refactor(ops): remove InfiniLM-suffixed InfiniOps callsrefactor(ops): route basic Llama through canonical InfiniOpsAug 10, 2026
@voltjia
voltjiaforce-pushed the refactor/remove-infinilm-infiniops-adapters branch from 2221fb1 to a69708cCompareAugust 11, 2026 00:49
@voltjia
voltjia changed the base branch from refactor/migrate-paged-attention-infiniops to mainAugust 11, 2026 00:50
@voltjia
voltjia requested a review from a teamAugust 11, 2026 00:50
@voltjia
voltjiaforce-pushed the refactor/remove-infinilm-infiniops-adapters branch from 04bb6ac to a69708cCompareAugust 11, 2026 02:33
@voltjia
voltjia merged commit 2c60352 into mainAug 11, 2026
16 checks passed
@voltjia
voltjia deleted the refactor/remove-infinilm-infiniops-adapters branch August 11, 2026 05:03
Sign up for freeto 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

@voltjia@Ziminli