Add experimental inference runtime API envelope by jarcherNV · Pull Request #403 · NVIDIA/flashdreams

jarcherNV · 2026-07-31T07:46:27Z

Define the initial flashdreams.runtime package with minimal T1 boundaries for runtime config, user/model inputs, schemas, input mapping, model adapters, runtime/session protocols, output targets, and metrics. Add focused CPU tests for the new API surface without migrating existing runners.

greptile-apps · 2026-07-31T07:48:28Z

Greptile Summary

Adds the initial experimental flashdreams.runtime API envelope.

Defines runtime configuration, user/model input schemas, input mapping, adapter/runtime/session protocols, output targets, and metrics.
Adds explicit required-user-snapshot validation and resets NullOutputTarget state when starting each run.
Adds focused CPU tests and records the completed T1 design decisions.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; both previously reported issues are addressed by explicit snapshot-validation helpers and per-run output-state reset behavior.

Important Files Changed

Filename	Overview
flashdreams/flashdreams/runtime/inputs.py	Defines validated user/model input envelopes and adds explicit required-snapshot validation helpers.
flashdreams/flashdreams/runtime/output.py	Defines the output-target boundary and correctly resets headless target accounting for each run.
flashdreams/flashdreams/runtime/interfaces.py	Defines the experimental adapter, reusable runtime, and isolated session protocols.
flashdreams/flashdreams/runtime/mapping.py	Defines the user-to-model input mapping protocol and identity implementation.
flashdreams/tests/test_inference_runtime_api.py	Exercises envelope validation, protocol composition, per-run output state, input mapping, and metrics behavior.

Class Diagram

%%{init: {'theme': 'neutral'}}%%
classDiagram
  ModelAdapter --> InferenceRuntime : creates
  InferenceRuntime --> InferenceSession : starts
  InputMapping --> UserInputs : consumes
  InputMapping --> ModelInputs : produces
  InferenceSession --> StepRequest : requests
  InferenceSession --> StepResult : produces
  OutputTarget --> StepResult : consumes
  MetricsRecorder --> RuntimeMetricSample : records

_{Reviews (9): Last reviewed commit: "docs: clarify input timing ownership" | Re-trigger Greptile}

Address Greptile feedback by adding required user snapshot validation and resetting NullOutputTarget state on each open(). Add focused tests for both behaviors and apply the CI pre-commit formatting/import fixes.

Add explicit session completion, step output metadata, and lower-level step data carriers to reduce coupling between runtime modules. Make runtime mapping payloads shallow-read-only, require sorted user events, enforce null output open/write lifecycle, and validate timing metrics use seconds. Expand tests with a fake adapter/runtime/session flow that composes the T1 API end to end.

Validate input mapping before runtime creation in the reference test loop and ensure created runtimes close when session startup fails. Make NullOutputTarget use internal open state, document why ModelAdapter is not runtime-checkable, and expand focused validation coverage for the runtime envelopes.

ArielG-NV

We should likely stub __post_init__ validation of types so that agents don't forget to include the validation.

I commented rather-heavily-- feel free to ignore a suggestion if it seems rather-nonsense

ArielG-NV · 2026-07-31T16:14:37Z

+    """Model/runtime execution settings.
+
+    Prompts, user controls, browser settings, output paths, and benchmark
+    directories intentionally live outside this object. The typed optimization
+    fields cover common cross-backend knobs; adapter-specific choices belong in
+    :attr:`runtime_options`.
+    """


Suggested change

"""Model/runtime execution settings.

Prompts, user controls, browser settings, output paths, and benchmark

directories intentionally live outside this object. The typed optimization

fields cover common cross-backend knobs; adapter-specific choices belong in

:attr:`runtime_options`.

"""

"""Model/runtime inference config.

Controls all knobs that functionally affect the inference of a model: Technique to host inference compute, `model_id` we are running inference on, `checkpoint` of model to load, whether we run CUDA via a CUDA Graph or not, etc...

"""

Yeah this could probably be clearer, but maybe we should avoid ‘all knobs’ because prompts/output/browser settings are intentionally outside InferenceConfig. What about ‘Runtime settings that affect model execution.’?

ArielG-NV · 2026-07-31T16:16:40Z

+from flashdreams.runtime._utils import freeze_mapping
+
+ExecutionBackend = Literal["local", "local-distributed", "external", "hosted"]
+"""Execution backend families the v0 envelope leaves room for."""


We should define what each "ENUM" literal means.

Suggested change

"""Execution backend families the v0 envelope leaves room for."""

"""Techniques to host inference compute."""

What about ‘Where and how inference compute is run’ rather than ‘host’, since local backends are not hosted services.

ArielG-NV · 2026-07-31T16:28:06Z

+    """Whether CUDA graph capture is requested, disabled, or left to the adapter."""
+
+    attention_backend: str | None = None
+    """Optional attention implementation selector."""


Suggested change

"""Optional attention implementation selector."""

"""Optional - attention implementation selector. `None` means unused."""

What're available attention_backend exposed here? Is this field going to be used to choose

sdpa_backend = { "math": torch.nn.attention.SDPBackend.MATH, "efficient": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, "cudnn": torch.nn.attention.SDPBackend.CUDNN_ATTENTION, "flash": torch.nn.attention.SDPBackend.FLASH_ATTENTION, }[self.backend]

One of the suggestions I have is we remove all these implementation details selector out of InferenceConfig, or we create a ModelConfig class to hold all the implementation details in a model. For now, we can store one ModelConfig per InferenceConfig. Later when we start building auto tune system we can store a list of candidate ModelConfig for the auto tuner to pick.

See Triton autotune system for reference: https://triton-lang.org/main/python-api/generated/triton.autotune.html

Yeah that probably makes sense. The fields like attention_backend and cache_policy are not really universal inference config and if we keep adding stuff like that then InferenceConfig could probably become bloated pretty quickly. I'll see if I can split this stuff out into a ModelConfig thing.

Regarding the suggested wording change, what about "Optional attention implementation selector; None leaves the choice to the adapter."

ArielG-NV · 2026-07-31T16:28:14Z

+    """Optional attention implementation selector."""
+
+    cache_policy: str | None = None
+    """Optional cache policy selector."""


Suggested change

"""Optional cache policy selector."""

"""Optional - cache policy selector. `None` means unused."""

What are some available options here? We should make a Enum instead of str.

Yeah a free string is probably ambiguous if this is a shared field. I'd prefer not to define the canonical cache policies now though, it's probably better to leave that for the T9 task.

ArielG-NV · 2026-07-31T17:08:08Z

+    resource_hints: Mapping[str, Any] = field(default_factory=dict)
+    """Cheap resource hints for launchers, schedulers, or hosted backends."""
+
+    def __post_init__(self) -> None:


We likely need more validation on parameter values and checks if types assigned were correct & if elements like runtime_options are valid

maybe we stub these checks?

Yeah that sounds reasonable, I'll see what checks I can add.

ArielG-NV · 2026-07-31T17:15:19Z

+class ModelInputSchema:
+    """Minimal metadata for model-facing initial and per-step inputs."""
+
+    initial_fields: tuple[InputField, ...] = ()


Suggested change

initial_fields: tuple[InputField, ...] = ()

initial_fields: tuple[InputField, ...] = ()

"""Inputs needed to run first generation."""

What about "These are model inputs required before starting the initial generation/session."

ArielG-NV · 2026-07-31T17:15:48Z

+    """Minimal metadata for model-facing initial and per-step inputs."""
+
+    initial_fields: tuple[InputField, ...] = ()
+    step_fields: tuple[InputField, ...] = ()


Suggested change

step_fields: tuple[InputField, ...] = ()

step_fields: tuple[InputField, ...] = ()

"""Inputs needed to run generation at an arbitrary step after the initial generation."""

What about phrasing it as "per-step model inputs required after the session starts."

ArielG-NV · 2026-07-31T17:25:02Z

+    kind: str
+    payload: Mapping[str, Any] = field(default_factory=dict)
+    source: str | None = None
+    event_id: str | None = None


why string and not int?

also- how do we have no id?

do we mean event_type here? otherwise sounds hash should just be the id

hmm, but we already have event.kind

Yeah, I think part of the confusion is naming. kind is meant to be the event type, e.g. move, look, reset, etc. event_id is not meant to be the primary identity/hash of the event; it is only an optional source/correlation ID from a transport, replay trace, or external system. Many runtime-stamped events will not need one.

Maybe we just rename event_id to something like source_event_id or correlation_id and maybe rename kind to event_type.

What do you guys think?

ArielG-NV · 2026-07-31T17:35:52Z

@@ -0,0 +1,189 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.


We should at some point doc-comment these members in the file to better explain purpose (and add a requirement to CI so that our API does not fall apart over time)

Yeah we should do that at some point before we merge all this stuff to main. I guess once the API is more settled we can throw that in.

gtong-nv · 2026-07-31T21:41:54Z

+    value: float | int
+    unit: str = "s"
+    step_index: int | None = None
+    category: str = "runtime"


category should be enum?

Yeah it probably should be once we know the canonical categories. For the T1 task I’d keep it open and we can revisit it in the future T9 task when benchmark/profiling categories are defined.

fangjunzhou-nv · 2026-07-31T21:29:40Z

+    """Whether CUDA graph capture is requested, disabled, or left to the adapter."""
+
+    attention_backend: str | None = None
+    """Optional attention implementation selector."""


What're available attention_backend exposed here? Is this field going to be used to choose

sdpa_backend = { "math": torch.nn.attention.SDPBackend.MATH, "efficient": torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION, "cudnn": torch.nn.attention.SDPBackend.CUDNN_ATTENTION, "flash": torch.nn.attention.SDPBackend.FLASH_ATTENTION, }[self.backend]

One of the suggestions I have is we remove all these implementation details selector out of InferenceConfig, or we create a ModelConfig class to hold all the implementation details in a model. For now, we can store one ModelConfig per InferenceConfig. Later when we start building auto tune system we can store a list of candidate ModelConfig for the auto tuner to pick.

See Triton autotune system for reference: https://triton-lang.org/main/python-api/generated/triton.autotune.html

fangjunzhou-nv · 2026-07-31T21:30:38Z

+    """Optional attention implementation selector."""
+
+    cache_policy: str | None = None
+    """Optional cache policy selector."""


What are some available options here? We should make a Enum instead of str.

fangjunzhou-nv · 2026-07-31T21:32:44Z

+ExecutionBackend = Literal["local", "local-distributed", "external", "hosted"]
+"""Execution backend families the v0 envelope leaves room for."""
+
+Precision = Literal["auto", "fp32", "fp16", "bf16"]


Suggested change

Precision = Literal["auto", "fp32", "fp16", "bf16"]

class Precision(Enum):

AUTO = "auto"

FP32 = "fp32"

FP16 = "fp16"

BF16 = "bf16"

We should use enum.Enum for all enums, this should make LSP happier.

Yeah there may be an LSP benefit. I’d keep literal for the T1 task and revisit enum when we actually implement this part of the API.

fangjunzhou-nv · 2026-07-31T21:35:12Z

+    cache_policy: str | None = None
+    """Optional cache policy selector."""
+
+    runtime_options: Mapping[str, Any] = field(default_factory=dict)


Should we use TypedDict over dict for all dictionary? Again, this should make LSP happier.

For runtime_options/resource_hints, a shared TypedDict would be too generic to help much because the keys are adapter/backend-specific. It's probably better to use adapter-owned typed configs once those shapes are known when we implement them.

fangjunzhou-nv · 2026-07-31T21:36:33Z

Should we split user input and model input into two source files?

Yeah that might make sense. We could probably do that in the T2/T3 tasks, I'll leave that up to whoever works on those tasks to decide.

fangjunzhou-nv · 2026-08-01T00:00:27Z

I'm not sure if the current "pipeline" way of interface definition is a good idea or not. This requires users to compose the application in a sequential way. They need to choose the session, an adapter, then a model.

I think a better to structure the application is providing a model interface and a session interface. The model interface specifies model input schema. The session interface specifies user input schemas. The user will implement their Session declaratively. The input mapping happens inside a Session implementation.

Yeah that might be a useful design direction, we should probably compare it before committing T2-T4. My intent for this PR is the lower-level T1 envelope rather than the final user-facing authoring API. I think the current primitives can still sit underneath a declarative model/session API: a model can expose ModelInputSchema, a session can expose UserInputSchema, and the session can own or wrap an InputMapping. I’d prefer not to do a full pivot in this PR until we have the alternative proposal and can check it against LingBot, OmniDreams, the-internal-demo-that-shall-not-be-named and Reactor requirements.

fangjunzhou-nv · 2026-08-01T00:04:58Z

+
+
+@dataclass(frozen=True, kw_only=True, slots=True)
+class ModelInputs:


As commented in flashdreams/flashdreams/runtime/interfaces.py, I'd like a declarative way of defining model inputs. The user should be able to build a Model that expose the input scheme directly. Instead of using a dedicated ModelInput object.

Yeah a declarative Model API could expose its schema directly. I still think we need a transport-neutral ModelInputs payload underneath for mapping, replay, tests, and benchmarks. The later declarative layer can wrap or produce ModelInputs rather than removing that boundary.

fangjunzhou-nv · 2026-08-01T00:06:48Z

Similar to the comments on ModelInput, if we design the pipeline with declarative primitives the user should be able to design the mapping stage inside Session

A declarative session can own or wrap the mapping, but I’d keep the UserInputs -> InputMapping -> ModelInputs boundary explicit. Reactor specifically liked that separation, and it helps replay/benchmark/runtime overrides without changing the model step.

fangjunzhou-nv · 2026-08-01T00:08:03Z

Is there any chance we can have Python decorator style metrics recorder like nvtx? We need more ergonomic profiling/benchmark tooling.

Yes, I like that direction. For T1 I’d keep MetricsRecorder as the low-level sink, then add decorator/context-manager helpers in the metrics/profiling work once benchmark needs are clearer.

fangjunzhou-nv · 2026-08-01T00:32:26Z

I'm going to make an alternative API proposal over the weekend.

jarcherNV · 2026-08-01T10:25:46Z

I'm going to make an alternative API proposal over the weekend.

Sounds good. I’ll keep this PR scoped as the minimal T1 envelope for now, and we can compare it against your proposal. But we should not deviate from what reactor specifically liked about the design unless we have a good reason to, this will be their API too after all and ideally should be a joint effort with them.

Co-authored-by: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com> Signed-off-by: jarcherNV <jarcher@nvidia.com>

greptile-apps Bot reviewed Jul 31, 2026

View reviewed changes

Comment thread flashdreams/flashdreams/runtime/output.py Outdated

Comment thread flashdreams/flashdreams/runtime/inputs.py

jarcherNV self-assigned this Jul 31, 2026

jarcherNV requested review from ArielG-NV and gtong-nv July 31, 2026 07:52

jarcherNV added 4 commits July 31, 2026 08:00

Fix runtime API validation and formatting

da683d8

Address Greptile feedback by adding required user snapshot validation and resetting NullOutputTarget state on each open(). Add focused tests for both behaviors and apply the CI pre-commit formatting/import fixes.

docs: record completed runtime API milestones

63f8b38

ArielG-NV reviewed Jul 31, 2026

View reviewed changes

gtong-nv reviewed Jul 31, 2026

View reviewed changes

fangjunzhou-nv reviewed Aug 1, 2026

View reviewed changes

jarcherNV and others added 5 commits August 1, 2026 03:50

Update resource hints wording

2ce669d

Co-authored-by: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com> Signed-off-by: jarcherNV <jarcher@nvidia.com>

Update compile wording

390a523

Co-authored-by: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com> Signed-off-by: jarcherNV <jarcher@nvidia.com>

Update cuda_graph wording

b891eaf

Co-authored-by: ArielG-NV <159081215+ArielG-NV@users.noreply.github.com> Signed-off-by: jarcherNV <jarcher@nvidia.com>

docs: clarify Reactor-aligned runtime API boundaries

c9da5d3

docs: clarify input timing ownership

5f83769

	"""Execution backend families the v0 envelope leaves room for."""
	"""Techniques to host inference compute."""

	"""Optional attention implementation selector."""
	"""Optional - attention implementation selector. `None` means unused."""

	"""Optional cache policy selector."""
	"""Optional - cache policy selector. `None` means unused."""

	initial_fields: tuple[InputField, ...] = ()
	initial_fields: tuple[InputField, ...] = ()
	"""Inputs needed to run first generation."""

		@@ -0,0 +1,189 @@
		# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.

-Precision = Literal["auto", "fp32", "fp16", "bf16"]
+class Precision(Enum):
+  AUTO = "auto"
+  FP32 = "fp32"
+  FP16 = "fp16"
+  BF16 = "bf16"



		@dataclass(frozen=True, kw_only=True, slots=True)
		class ModelInputs:

Uh oh!

Conversation

jarcherNV commented Jul 31, 2026

Uh oh!

greptile-apps Bot commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Greptile Summary

Confidence Score: 5/5

Important Files Changed

Class Diagram

Uh oh!

Uh oh!

Uh oh!

ArielG-NV left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ArielG-NV Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

greptile-apps Bot commented Jul 31, 2026 •

edited

Loading

ArielG-NV Jul 31, 2026 •

edited

Loading