Cosmos

[micwill755]

Below is a ready-to-paste PR description. I left test checkboxes unchecked because only pre-commit has been verified in this session.

## What does this PR do?

**Type of change:** New feature

**Overview:**

This PR adds Vision-Language-Action (VLA) and World Foundation Model (WFM) inference support to TensorRT Edge-LLM.

### Vision-Language-Action support

- Adds `VlaInferenceRuntime`, which orchestrates language, vision, optional audio, and action engines.
- Adds `LLMEngineRunner` and `LinearKVCache` for VLA-compatible prefill, decoding, system-prompt KV-cache reuse, and hybrid recurrent/conv state handling.
- Adds a model-agnostic fixed-shape `VitRunner` and registers the new `vit` multimodal model type.
- Adds `ActionRunner` supporting:
  - Prefix-KV context handoff with flow-matching rollout, used by Alpamayo-style models.
  - Context-tensor handoff with velocity integration, used by GR00T and PI0.5-style action heads.
  - Configurable action noise seeds and denoising schedules.
  - Robot-state and embodiment inputs.
  - PI0.5 suffix position IDs and attention masks.
- Adds `ActionContextRunner` for projecting language hidden states into action-context embeddings.
- Extends `LLMGenerationRequest` with robot state, embodiment ID, and action batch size.
- Extends `LLMGenerationResponse` with generated robot actions.
- Adds compact VLA prefix-padding support through tokenizer configuration.
- Adds `vla_inference`, a JSON-driven VLA example supporting text, image, trajectory, robot-state, LoRA, batching, and action output.

### Cosmos3-Edge WFM support

- Adds `WFMInferenceRuntime`, which orchestrates the Cosmos encode → denoise → decode pipeline.
- Adds dedicated TensorRT runners for:
  - `VaeEncodeRunner`: video pixels → vision latents.
  - `VaeDecodeRunner`: vision latents → generated video.
  - `AudioEncodeRunner`: optional waveform → sound latents.
  - `AudioDecodeRunner`: sound latents → generated waveform.
  - `EmbedRunner`: noisy latents and timestep → generation tokens.
  - `MotBackboneRunner`: text, generation tokens, and rotary embeddings → hidden states.
  - `DenoiseHeadRunner`: hidden states → predicted vision or sound latents.
  - `CosmosDenoiseRunner`: orchestrates embedding, MoT backbone, denoise heads, and scheduler updates.
- Adds Cosmos prompt packing, multimodal RoPE preparation, latent noise seeding, and a UniPC flow scheduler.
- Supports vision-only Cosmos3-Edge inference and optional Cosmos3-Omni sound generation when audio engines are present.
- Reuses shared TensorRT execution-context memory across sequential engines.
- Adds the `wfm_inference` example with:
  - JSON batch requests.
  - Per-request and global seeds/inference-step settings.
  - Optional FP16 pixel and waveform inputs.
  - Warmup and profiling support.
  - Raw FP16 video and waveform outputs.
- Documents Cosmos3-Edge engine export, build, and inference in the README.

## Usage

### VLA inference

Build the example:

```bash
cmake --build build --target vla_inference -j$(nproc)

Example input:

{
  "batch_size": 1,
  "max_generate_length": 128,
  "action_batch_size": 1,
  "requests": [
    {
      "messages": [
        {
          "role": "user",
          "content": [
            {
              "type": "image",
              "image": "/path/to/observation.png"
            },
            {
              "type": "text",
              "text": "Pick up the red cube."
            }
          ]
        }
      ],
      "robot_state": [0.0, 0.0, 0.0],
      "embodiment_id": 0
    }
  ]
}

Run inference:

export EDGELLM_PLUGIN_PATH=/path/to/libNvInfer_edgellm_plugin.so

./build/examples/llm/vla_inference \
  --engineDir=/path/to/language_engines \
  --multimodalEngineDir=/path/to/vla_engines \
  --inputFile=/path/to/vla_input.json \
  --outputFile=/tmp/vla_output.json \
  --actionNoiseSeed=5

The output JSON contains generated text and, when an action engine is loaded, either actions or trajectory for each request.

Cosmos3-Edge WFM inference

Build the example:

cmake --build build --target wfm_inference -j$(nproc)

Example input:

{
  "num_inference_steps": 2,
  "seed": 42,
  "requests": [
    {
      "prompt": "A robot arm picks up a red cube.",
      "generate_sound": false,
      "output_video_file": "/tmp/wfm_output_video.fp16"
    }
  ]
}

Run inference:

export EDGELLM_PLUGIN_PATH=/path/to/libNvInfer_edgellm_plugin.so

./build/examples/wfm/wfm_inference \
  --engineDir=/path/to/cosmos_edge_engines \
  --inputFile=examples/wfm/wfm_input_example.json \
  --outputFile=/tmp/wfm_output.json

The video and waveform files are raw FP16 tensor dumps rather than MP4 or WAV containers.

Pull Request Checklist

Thank you for contributing to TensorRT Edge-LLM! Before we review your pull request, please make sure the following items are complete.
Please also refer to Contributor guidelines for general guidelines.

Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit.
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

Tests

• Tests have been added or updated as needed.
• All tests are passing.

Pre-commit validation:

isort................................Passed
CRLF end-lines remover...............Passed
Insert license in comments...........Passed
yapf.................................Passed
clang-format.........................Passed
cmake-format.........................Passed
codespell............................Passed
autoflake............................Passed

Documentation

• Updated necessary documentation.

The README now documents Cosmos3-Edge engine export, build, engine layout, request format, and inference usage.

Compatibility

• The change is backward compatible.

The new runtimes, request fields, model types, and examples are additive. Existing LLM and multimodal inference paths remain available.

Additional Information

Cosmos3-Edge engine export currently uses the companion Test/wfm/export_wfm_cosmos_edge.py workflow. The prompt used during inference must match the export prompt so that packing_static.json remains compatible with the packed text sequence.

Before submitting, the branch still needs its build/runtime tests recorded. Also, the successful pre-commit run produced uncommitted formatting and license changes across multiple files; those should be reviewed before committing.

[micwill755]

cc: @narendasan