Backends rápidos

Portabilidade do TileLang

Os cinco kernels de laya/tl_kernels.py podem ser executados no alvo de CPU (c) do TileLang via compile_cpu. É uma especialização escalar fp32 explícita, não um backend de aceleração por CPU para Agent.accelerate. Requer o TileLang e um compilador C++ local. Não introduz nenhuma dependência de serviço em tempo de execução.

import torch
from laya import tl_kernels as K

A = torch.randn(17, 67)
W = torch.randn(70, 67)
b = torch.randn(70)
C = torch.empty(17, 70)
kernel = K.compile_cpu(K.gemm_kernel, 70, 67, bias=True, act="gelu")
kernel(A, W, b, C)

compile_cpu(factory, *args, **kwargs) aceita as mesmas opções de forma e de operação que as cinco factories de GPU. Seleciona cpu=True, dtype="float32", o alvo c e desativa a vetorização. Todas as entradas e saídas de vírgula flutuante devem ser tensores float32 de CPU contíguos; os comprimentos de atenção continuam int32. Converte as ativações de 16 bits com .cpu().float().contiguous() antes de o chamares. O LayerNorm continua a atualizar o seu tensor residual no lugar, e o RoPE continua a atualizar as colunas Q/K no lugar. Retém o kernel compilado para reutilizar as suas dimensões dinâmicas.

A especialização de CPU substitui as alocações de fragment e shared por buffers locais, usa ciclos T.grid em série e omite a anotação swizzle de GPU da atenção. Os GEMMs usam a implementação escalar de CPU do TileLang e as reduções usam buffers locais. Os algoritmos em blocos originais, os predicados de padding, as máscaras de atenção e o softmax online permanecem partilhados com a implementação de GPU. Os intermediários de CPU permanecem em fp32, incluindo as probabilidades de atenção; os intermediários de GPU mantêm o seu dtype original. Nenhuma aceleração de CPU ou backend de CPU para o modelo completo é reivindicado.

Nova sondagem no TileLang 0.1.14

Observado no Linux, Python 3.12.13, torch 2.11.0+cu130, TileLang 0.1.14 e uma RTX 4070 Ti SUPER. A preocupação de portabilidade anterior aplica-se a compilar a especialização de GPU inalterada, não à possibilidade de gerar código para CPU. No commit base fa9a2a7, esta página de documentação estava ausente do checkout.

O teste compila factory.get_tir(...) com target="c" e a configuração da passagem FAST da GPU. Estes são os textos de diagnóstico exatos (locais de código e stack traces omitidos). Tanto bf16 como fp16 foram sondados para cada kernel.

Os quatro primeiros erros de fragment/redução são tvm.error.InternalError. A falha de geração de código bf16 também é um InternalError. O RoPE em FP16 lança RuntimeError: Compilation Failed! seguido da invocação do compilador e do código-fonte; o diagnóstico acima é emitido em stderr. As suas conversões vetoriais geradas também falham ao converter vetores float de volta para vetores half.

Para isolar o suporte a dtype do suporte a fragment, os testes compilam cada kernel novamente com cpu=True (buffers locais e ciclos em série), mantêm bf16 e desativam a vetorização. Os cinco falham então exatamente com:

Cannot convert type bfloat16 to C type

Assim, os fragments são o primeiro bloqueio para GEMM, GEGLU e atenção, as reduções de fragment para o LayerNorm, e bf16 é independentemente um bloqueio para os cinco. O RoPE não tem alocação de fragment nem redução; GEMM e GEGLU não têm operação T.reduce_* explícita. As reduções da atenção são inicialmente mascaradas pela sua falha de layout. Sondagens separadas de fragment apenas fp32 reproduzem o erro de preenchimento e o erro de redução tanto para reduce_sum como para reduce_max, sem nenhum GEMM ou bf16. Substituir apenas os scopes também produziu este erro semântico para o buffer do GEMM (os outros buffers afetados eram Ci, x e s):

[Tilelang Semantic Check] Local buffer `C_l` is indexed by T.Parallel loop variable `i`. Local buffers are thread-private and do not participate in parallel layout inference. Use T.serial/T.vectorized/T.unroll for per-thread local indexing, or T.alloc_fragment when the indexed dimension should be distributed across threads.

A especialização fp32 em série remove estes bloqueios. As asserções de diagnóstico estão fixadas na versão 0.1.14 e são ignoradas noutra versão, onde os erros devem ser sondados de novo. Os testes numéricos de CPU continuam a correr noutras versões.

Verificações numéricas

Executa python -m pytest tests/test_fast_cpu.py -q -s. No ambiente acima: 58 passaram. As verificações apenas de CPU não exigem CUDA; só as comparações de GPU são ignoradas sem ela. A cobertura inclui tiles GEMM M/N/K irregulares, todos os epílogos de GEMM, GEGLU, todas as combinações de residual/bias, valores residuais grandes, wraparound de posição do RoPE e colunas V intactas, formas de atenção estáticas/dinâmicas, janelas deslizantes, tiles parciais, comprimentos desiguais, sequências vazias e saídas de padding finitas.

A semente é 1234. As comparações de GPU usam valores de entrada idênticos arredondados para bf16 ou fp16 e depois promovidos para fp32 para execução na CPU. São tolerâncias absolutas para os fixtures limitados, não uma garantia para magnitudes arbitrárias ou profundidade de modelo. As comparações de atenção usam linhas de query válidas, como em tests/test_fast.py.

A tolerância CPU/referência é 2e-5 (2e-6 para o RoPE). As atualizações do fluxo residual são exatamente iguais, incluindo o caso sem residual.

Evidência de preservação da GPU

A palavra-chave cpu tem valor padrão false. Os padrões existentes de bf16/fp16, os scopes de alocação de GPU, os ciclos paralelos, os swizzles e as opções FAST não mudam. As chamadas de GPU em fp32 padrão ainda lançam ValueError.

As execuções antes/depois usaram o módulo original extraído com git show fa9a2a7:laya/tl_kernels.py e o módulo alterado, respetivamente. O original foi carregado como laya.tl_kernels via importlib para as execuções de linha de base; todo o resto do código do Laya e o ambiente Python permaneceram iguais.

• python -m pytest tests/test_fast.py -q, com os testes full-forward configurados para carregar o checkpoint inglês em cache identificado abaixo: 13 passaram antes; 13 passaram depois. No início, sem checkpoint, eram 11 passaram / 2 ignorados. Ambas as execuções completas emitiram o aviso existente de limitação de temperatura do checkpoint (choice:11+=0.10058280825614929 -> 0.5).
• Captura com semente dos testes de kernel existentes: 24 tensores de saída idênticos bit a bit, incluindo as atualizações residuais; diferença máxima antes/depois 0.
• Código-fonte CUDA gerado para as cinco formas de sondagem acima, em ambos os dtypes: 10/10 idênticos byte a byte. Isto também cobre o RoPE, ausente da suíte fast original.
• python benchmarks/parity_fast.py --model "$MODEL" --dtype bf16 --json ... e o comando fp16 equivalente: 288 perguntas em 60 estados por dtype. Todos os registos JSON antes/depois (probabilidades fp32, stock e fast) comparam exatamente iguais; diferença máxima de probabilidade antes/depois 0.

MODEL era o snapshot inglês em cache convaiinnovations/laya 55cf4c4ebb4ebe31b2550e8bdf3bd21b99753851; as execuções usaram HF_HUB_OFFLINE=1. Nenhuma comparação com checkpoint multilingue ou de decisões tipadas é reivindicada aqui.

Verificações do repositório

Todos os comandos usaram /home/ckl/projects/S/laya/.venv/bin/python; ruff e zensical vieram do diretório bin desse virtualenv.

O novo conjunto de testes está registado como exigindo o extra opcional do TileLang e um compilador C++ nas isenções existentes do teste de empacotamento. O teste de contrato de API fixa o seletor de CPU aditivo apenas por palavra-chave, o dtype padrão de GPU inalterado e a assinatura de compile_cpu, sem importar o TileLang no ambiente de CI base.

AOTInductor

DecisionModel pode ser exportado, compilado num pacote .pt2 e carregado com torch._inductor.aoti_load_package. O pacote devolve tanto os logits de decisão como os logits de ação. A tokenização, o padding, a calibração de temperatura e a formatação das respostas continuam a ser tarefa de quem chama; isto não adiciona um backend de Agent nem altera o seu caminho de execução padrão.

O PR #472 fechado documentou o bloqueio de dtype anterior. O estado agrupado e as características de confiança da cabeça de ação são calculados em fp32, mesmo quando os pesos de um modelo são explicitamente bf16. Sem autocast, o primeiro linear de ação, portanto, recebia uma entrada fp32 e pesos bf16. Agora o forward converte a entrada concatenada para o dtype dos pesos da cabeça apenas fora do autocast. Softmax, entropia e os logits de decisão devolvidos mantêm os seus cálculos em fp32. A inferência eager existente com parâmetros fp32, incluindo AMP fp16/bf16, mantém o seu comportamento numérico; dtypes mistos de peso/autocast também mantêm a conversão original do autocast.

Exporta uma cópia de avaliação separada convertida para bf16, fora do autocast. Declara o eixo de tokens como 16 * Dim("tokens16", ...) para satisfazer as guardas de alinhamento da atenção. Linhas e contagens de marcadores podem ser dinâmicas de forma independente. Não presumas que uma exportação capturada sob autocast possa ser empacotada fora desse contexto: usa dtypes de peso explícitos para esta receita.

Reproduzir offline

A verificação baseada em asserções usa, por padrão, um ModernBERT minúsculo inicializado aleatoriamente, sem acesso à rede. Passa um diretório de checkpoint local para uma medição com um modelo real. A ausência de CUDA, das APIs do AOTInductor ou de um compilador C++ produz um SKIP explícito; falhas de compilação e de paridade numa instalação compatível fazem a verificação falhar.

python scripts/check_aoti.py --output-dir /tmp/laya-aoti-smoke
HF_HUB_OFFLINE=1 TORCHINDUCTOR_CACHE_DIR=/tmp/laya-aoti/cache \
  python scripts/check_aoti.py --model /path/to/local/multilingual \
  --output-dir /tmp/laya-aoti

O diretório de saída contém decision.pt2, results.json, inputs.pt e eager.pt. O JSON inclui snapshots completos do nvidia-smi, tempos de exportação/empacotamento/carregamento, bytes do pacote, latência e deltas máximos absolutos de logit/probabilidade para ambas as cabeças. Artefactos binários e caches do compilador são deliberadamente mantidos fora do repositório.

Medido numa RTX 4070 Ti SUPER

2026-10-03, Linux x86-64, Python 3.12.13, torch 2.11.0+cu130, transformers 5.17.0, driver NVIDIA 615.71.09, 16.376 MiB de VRAM. Checkpoint: convaiinnovations/laya, subdiretório multilingual na revisão 1c5edc17a7acd8701df6fc341c0d179f1c62c982. A linha de base é a main fa9a2a7. Snapshots completos da máquina e medições não arredondadas estão em aoti_multilingual_rtx4070.json.

A falha reproduzida ocorre no empacotamento, após a exportação bem-sucedida: mat1 and mat2 must have the same dtype, but got Float and BFloat16. Cada execução usou um cache do Inductor separado e inicialmente vazio; a compilação AMP correu antes do empacotamento, por isso o tempo do pacote não é uma medição de arranque a frio de um interpretador novo. A verificação no processo novo carregou apenas o pacote e as entradas guardadas (sem checkpoint), com torch.compile e os pontos de entrada de empacotamento bloqueados. Reproduziu as respostas de ambas as cabeças. O PyTorch ainda compilou uma pequena sondagem de capacidade AVX de CPU no cache inicialmente vazio; nenhum grafo de modelo ou kernel CUDA foi compilado. Uma afirmação geral de «nenhuma atividade do compilador ao carregar» seria, portanto, imprecisa nesta runtime.

As mesmas entradas guardadas contêm sete linhas de decisão em três lotes, com texto real tokenizado de faturação/reembolso, perguntas choice/score/noul, padding e contagens de marcadores válidos desiguais. Cada latência é a mediana de cinco grupos de 30 forwards após dez aquecimentos, com medição de tempo de relógio de parede sincronizada. torch.compile(dynamic=True) usa o modo padrão. Nenhum grafo CUDA explícito, tokenização, exportação ou compilação está incluído nestes números de latência.

Aqui AMP significa parâmetros fp32 com autocast bf16. bf16 explícito significa que os pesos e o fluxo residual são bf16, sem autocast; usa essas colunas para a comparação de execução mais próxima do pacote. bf16 explícito eager/compiled e AOTI estavam indisponíveis antes desta correção.

Tanto o argmax de decisão como o de ação concordam em 7/7 linhas para ambas as comparações de AOTI. As probabilidades são saídas softmax brutas, sem calibração. A distribuição de ação está saturada nestas entradas, por isso o seu delta de probabilidade zero não implica logits subjacentes idênticos em relação ao AMP. O forward compilado em bf16 explícito também difere do bf16 explícito eager (delta máximo de logit de decisão 0.5, delta de probabilidade 0.00659859). A compilação em precisão reduzida não é exata bit a bit.

A GPU foi partilhada com aplicações de ambiente de trabalho e outras tarefas de Python; a compilação de CPU também foi partilhada, e os clocks não estavam fixados. Estes são tempos observados, não uma reivindicação de aceleração isolada. Snapshots do nvidia-smi que enquadram as execuções:

Limites restantes

• O artefacto é específico para este checkpoint, esta precisão, esta pilha de PyTorch/runtime e este alvo de GPU; a portabilidade para outro hardware ou versões do PyTorch não foi testada.
• Esta verificação exporta as linhas 1–8, tokens 32–512 em múltiplos de 16 e slots de marcador 2–8. Exercita três formas, incluindo formas diferentes do exemplo de exportação, em vez de cada ponto dessas faixas. Uma opção válida pode usar slots de marcador com padding; um tensor real de um único slot usa o ramo separado de opção única e precisa de uma exportação separada. Comprimentos de token arbitrários e uma camada de produção de bucketing/dispatch estão fora desta mudança.
• Sete linhas verificam a regressão de empacotamento, não a exatidão ampla do checkpoint nem a paridade de confiança calibrada. Converter uma cópia de exportação para bf16 difere de manter os pesos fp32 sob autocast; o próprio caminho eager existente permanece inalterado.
• O script exige um build do PyTorch compatível com CUDA e uma cadeia de ferramentas de compilador local para criar o artefacto. O .pt2 embute o modelo e os kernels CUDA; nenhum serviço alojado está envolvido.