Introduction

Physical intelligence requires not only more capable models, but also faster and more efficient systems.

Beyond scaling foundation models, when these models are deployed into robots, industrial systems, and real-world environments, the inference system often becomes a critical factor determining overall performance and user experience. A single action decision may involve multiple stages, including visual understanding, semantic reasoning, state modeling, and action generation. Any millisecond-level latency introduced during these processes can be amplified into noticeable response delays in physical interactions.

Therefore, we believe that the competition in Physical AI is not only a competition of model capabilities, but also a competition of AI Infrastructure.

In this work, we conduct a systematic inference optimization study based on a widely adopted Vision-Language-Action (VLA) model. From framework restructuring and CUDA Graph optimization to Triton Kernel acceleration and Kernel Fusion, we develop a cross-stack co-optimization approach spanning Framework, Runtime, and Kernel layers. This approach reduces inference latency from 130 ms to 32 ms, achieving approximately 4× performance improvement.

More than an optimization effort for a single model, this work represents our exploration of the next-generation Physical AI Inference Stack.

Background and Optimization Results

Robot Vision-Language-Action (VLA) models integrate visual perception, language understanding, and action generation into a unified neural network. In such systems, inference latency directly impacts real-time robotic control. Excessive inference time can lead to delayed action commands, execution stuttering, and even task failures in highly dynamic environments.

The original StarVLA training framework exhibits three major efficiency bottlenecks in its inference pipeline.
First, framework-level scheduling overhead: the general-purpose interface of HuggingFace Transformers introduces unnecessary conditional checks and configuration parsing during inference.
Second, excessive kernel launch overhead: a single inference step requires approximately 10,000 CUDA kernel launches, with scheduling overhead accounting for more than 40% of the total latency.
Third, general-purpose operators are not optimized for inference-specific workloads: native PyTorch RMSNorm, MLP, and other operators fail to exploit the characteristics of VLA inference, such as batch size = 1 and sequence lengths ranging from 100 to 400.

Considering that model architectures remain relatively stable in practical deployment scenarios, while achieving ultra-low inference latency is a primary objective, we did not adopt general-purpose inference frameworks such as vLLM or TensorRT-LLM. Instead, we developed a lightweight inference framework from the ground up and implemented customized operators with Triton. This design choice enables fine-grained control over every stage of the inference pipeline and provides a foundation for further optimizations, including CUDA Graph and Kernel Fusion.

Based on this framework, we completed a three-stage inference optimization process on NVIDIA RTX 4090, covering pure language models, vision-language models, and full VLA models. We implemented optimized versions of all 10 operators involved in the inference pipeline. Ultimately, with a single-camera 224×224 input and maintaining accuracy degradation within 1%, we reduced the average pure-model inference latency from 130 ms to 32 ms, achieving approximately 4× acceleration.

Performance Optimization Approach

Overall Approach

The StarVLA model consists of three sequential modules: the Vision Encoder (24-layer ViT, approximately 300 MBparameters), the LLM Decoder (36 layers, including projection layers, with approximately 2.75B parameters), and the Action Head (a 4-step × 36-layer DiT, with approximately 1.1B parameters).

The LLM adopts a GQA (Grouped Query Attention) mechanism with 32/8 heads and a SwiGLU MLP with an intermediate dimension of 9,728. It operates in a pure prefill mode, without an autoregressive decoding stage. The Action Head employs Flow-Matching Euler denoising, starting from random noise and progressively generating meaningful action trajectories through four denoising steps.

The complete inference pipeline contains approximately 240 layers, while involving only 10 operator types: RMSNorm, LayerNorm, RoPE, GELU, SiLU, MatMul (QKV/MLP), SDPA, Sin/Cos, Embedding, and Element-wise operations. This characteristic makes each operator a valuable target for manual optimization and fine-grained acceleration.

Based on these characteristics, we developed a comprehensive optimization approach consisting of two key components: a custom-built inference framework and operator-level optimization. At the system level, we leverage CUDA Graphto reduce runtime overhead and improve execution efficiency. At the kernel level, we develop and fuse customized operators with Triton, further exploiting hardware capabilities to accelerate inference.

The optimized model architecture and inference pipeline are illustrated in the figure above.

Custom-Built Inference Framework

Problem

The original inference pipeline was built on top of HuggingFace Transformers and PyTorch nn.Module abstractions. Designed primarily for training flexibility, the framework contains substantial logic that is unnecessary during inference. The Python-level forward calls of each nn.Module introduce additional scheduling overhead, which becomes significant when accumulated across 36–144 layers. In addition, the from_pretrained weight-loading interface requires model structure and configuration parsing, resulting in limited transparency and control over the execution path.

Approach

We rebuilt the entire inference pipeline in-house, from weight loading to operator scheduling, without relying on HuggingFace Transformers. The optimization consists of three key efforts:

First, direct weight loading. We directly load weight tensors from PyTorch checkpoints or safetensors, bypassing the configuration parsing stage.

Second, modular code organization. Each model component is implemented as an independent module (vision_encoder.py, language_model.py, dit_model.py, and action_head.py), with a unified dispatch layerresponsible for operator scheduling and execution control.

Third, three-stage progressive optimization. We progressively adapt the inference pipeline from the pure language model to the vision-language model, and finally to the complete VLA model. Lower-level implementations are reused by higher-level components, enabling incremental validation and optimization.

Benefits

The custom-built inference framework eliminates framework-level scheduling overhead and provides full transparency over the execution flow. This optimization contributes approximately 15% acceleration, reducing inference latency from 130 ms to 115 ms.

More importantly, this optimization demonstrates strong general applicability: any model deployed in production using HuggingFace Transformers can potentially achieve comparable gains by reducing unnecessary framework-level overhead, with the required engineering effort largely independent of model complexity.

CUDA Graph

Problem

In GPU inference, the CPU needs to invoke a separate kernel launch through the driver for each GPU kernel execution. A single kernel launch typically incurs approximately 4 μs of overhead. Among the more than 10,000 kernel launches required for a single inference pass, approximately 5,900 are small kernels, including element-wise operations, reshaping, and memory copy operations. For these kernels, the launch overhead can become comparable to the actual computation time.

As a result, a significant portion of inference latency is spent on CPU-to-GPU command dispatch, rather than actual computation.

Approach

CUDA Graph records a sequence of kernel launches into a single execution graph. During subsequent inference runs, the entire sequence can be executed through a single graph replay, eliminating the overhead associated with individual kernel launches.

We adopt the standard warmup → capture → replay workflow: first performing a complete inference pass for warmup, then re-executing the pipeline within the torch.cuda.graph() context to capture the kernel execution sequence. Subsequent inference runs directly replay the captured graph.

We apply CUDA Graph optimization to three core components:

• Vision Graph: covers the complete pipeline from PatchEmbed, through 24 ViT blocks, to Merger.
• LLM Graph: covers the 36-layer Transformer stack through the final RMSNorm.
• Action Graph: covers 4-step denoising × 36-layer DiT execution.

Two implementation details require particular attention.

First, memory pre-allocation. Input and output buffers are pre-allocated before graph capture. During replay, input data is updated through copy_(), avoiding dynamic memory allocation inside the graph execution.

Second, sequence length adaptation. When the input sequence length changes, the LLM Graph automatically falls back to eager execution and triggers graph re-capture. Cached sequence lengths are used for compatibility checks to determine whether an existing graph can be reused.

Benefits

CUDA Graph reduces the number of CPU-side kernel launch calls from 10,605 to 3 (one replay() call per graph), effectively eliminating the CPU launch bottleneck.

According to Nsight Systems (nsys) profiling, the number of GPU-side kernel executions is reduced from 10,605 to 5,139 (52% reduction). Among them, small kernels are reduced from 5,913 to 499 (92% reduction) through graph-level optimization and kernel elimination/fusion. The remaining large computational kernels, such as matrix multiplication and attention, are captured into execution graphs and scheduled directly by the GPU.

Overall, CUDA Graph contributes approximately 60% acceleration, reducing latency from 57 ms to 115 ms, and represents the single largest performance gain among the four optimization techniques.

CUDA Graph is broadly effective for models with relatively static inference computation graphs. The more stable the execution shape, the greater the potential benefit. This model is particularly suitable because it operates in a pure prefillscenario without autoregressive decoding, where sequence lengths remain constant throughout inference, enabling significant gains from graph-based execution.

Custom Operator Development

Problem

PyTorch’s general-purpose operator implementations are designed to balance performance across a wide range of input shapes, and therefore do not fully exploit the specific characteristics of VLA inference workloads, such as batch size = 1, sequence lengths of 100–400, and fixed hidden dimensions of 1024 or 2560.

Some operators are invoked extremely frequently. For example, RMSNorm is executed 144 times in total across LLM pre-normalization (pre-norm) and QK normalization (QK norm). Therefore, even minor optimizations at the operator level can accumulate into significant overall performance gains.

Approach

We developed custom kernels using the Triton language specifically optimized for VLA inference workloads. Compared with CUDA development, Triton provides higher development efficiency while maintaining performance close to hand-optimized kernels. Some kernels leverage the @triton.autotune mechanism to automatically search for optimal block size and warp count configurations during compilation.

Below, we use RMSNorm and RoPE as examples to illustrate the core principles of operator optimization.

RMSNorm

RMSNorm (Root Mean Square Normalization) is executed 72 times in LLM pre-normalization and another 72 times in QK normalization, resulting in 144 invocations in total.

The native PyTorch implementation consists of a sequence of separate operations:

x.float() → x*x → mean(-1) → torch.rsqrt() → * x → * w.float()

Each operation launches an independent kernel, and intermediate results require approximately 3–4 rounds of HBM read/write operations.

We implement a single-pass kernel that loads an entire row into SRAM, completes the computation of sum of squares, reciprocal standard deviation (rstd), and final scaling entirely in registers, and then writes the result back to HBM. The optimized kernel performs only one HBM read and one HBM write.

Under the LLM prefill configuration (M=182, N=2560), this optimization achieves 1.81× acceleration.

RoPE

RoPE (Rotary Position Embedding) performs half-dimension rotation transformations on Q/K tensors across all 36 Transformer layers, with every token requiring computation.

The native PyTorch implementation creates new q_rot / k_rot tensors and involves multiple kernel operations, including split, element-wise computation, concatenation (cat), and dtype conversion.

We implement an in-place Triton kernel that directly loads the first and second halves of the original Q/K tensors into registers, performs the rotation transformation, and writes the results back to the same memory locations.

Under the configuration B=1, S=182, H=32, D=128, this optimization achieves 1.52× acceleration.

The two operators above represent typical examples with positive performance gains. Overall, most operators—including RMSNorm, RoPE, GELU MLP, and SwiGLU MLP—achieve 1.3×–7× acceleration with Triton optimization.

However, not all operators are suitable for Triton implementation. For example, SDPA (Scaled Dot-Product Attention) is deeply optimized in the cuDNN backend, and our Triton implementation performs worse. Therefore, the attention operator continues to rely on the cuDNN backend in the final inference pipeline.

Benefits

Custom operator development contributes approximately 10% acceleration, reducing latency from 57 ms to 47 ms.

It is important to note that this acceleration is measured after CUDA Graph has eliminated kernel launch overhead, representing the true improvement from operator-level computation efficiency.

More importantly, these optimized operators are fundamental components shared across Transformer architectures. Models such as LLaMA, Qwen, Gemma, and GPT all adopt similar normalization and positional encoding mechanisms. Therefore, the Triton kernels developed in this work can be directly transferred to inference optimization for other models without additional operator redevelopment.

Kernel Fusion

Problem

The NVIDIA RTX 4090 provides approximately 1 TB/s HBM bandwidth, while the on-chip SRAM bandwidthreaches around 20 TB/s, creating an order-of-magnitude gap between the two memory systems. In PyTorch, each intermediate result in a sequential operator chain is written back to HBM, making memory bandwidth a significant bottleneck.

Taking GELU MLP as an example: the computation first performs up = x @ W_up and writes the result to HBM. It then computes gelu = GELU(up), requiring another HBM read/write operation. Finally, out = gelu @ W_down reads gelu back from HBM for the final matrix multiplication.

This process involves four HBM memory accesses:

• read x
• write up
• read up
• write out

Among these operations, the intermediate write and read of up can be eliminated by retaining intermediate results in SRAM.

Approach

We retain intermediate results in SRAM during the tiled computation loop of the Triton kernel.

The optimized execution flow is as follows:

1. Load a tile of x from HBM into SRAM.
2. Compute up = x @ W_up while keeping intermediate results in SRAM registers.
3. Compute gelu = GELU(up) directly in SRAM.
4. Accumulate partial results through out += gelu @ W_down.
5. Write the final out tensor back to HBM.

Through this approach, HBM memory accesses are reduced from four operations to two, significantly reducing memory bandwidth pressure.

Benefits

Based on standalone operator benchmarks (PyTorch eager mode vs. Triton fused implementation, without CUDA Graph), the acceleration of Fused GELU MLP scales positively with matrix size.

In the Merger module (M=49, C=4096), the fused implementation achieves 7.33× acceleration. In the VisionMLPmodule (M=512, C=1024), it achieves 3.82× acceleration.

The same principle is further applied to Fused SwiGLU MLP in the LLM component (1.27× acceleration) and Fused QKV, where three independent matrix multiplications are combined into a single fused operation.

Overall, kernel fusion contributes approximately 15% acceleration, reducing inference latency from 47 ms to 32 ms.

Future Work

Among the potential optimization directions, quantization is expected to provide the most significant performance gains.

Different quantization strategies, including FP8 and INT4, have different levels of applicability across model components. The LLM Decoder, dominated by matrix multiplication operations, is expected to benefit the most from quantization, with the most mature industry practices available. The Action Head involves an iterative denoising process that is sensitive to numerical precision, requiring careful evaluation of error accumulation introduced by quantization. The Vision Encoder contributes a relatively small portion of the overall inference latency, making it a lower-priority target for quantization.

Overall, quantization techniques are expected to further reduce end-to-end inference latency, potentially lowering it from 32 ms to below 20 ms.

Citation

If you use this work in research or engineering projects, please cite as follows:

@article{rimbot2026aivlaqwen3,
  author = {RimBot Research Team},
  title = {AI Infra Optimization for VLA: A Technical Breakdown of Model Inference Performance Improvements},
  year = {2026},
  howpublished = {Technical Report},
  note = {An engineering report on optimizing Qwen3-VLA inference with NVIDIA RTX 4090 through a custom inference framework, CUDA Graph, Triton kernels, and operator fusion, achieving 4× acceleration.}
}