Cosmos3-Edge on Blackwell: dense attention on FlashInfer's ragged prefill, deterministic compile, fast whole-clip decode, compiled batched step

docs/adding_models.rst

@@ -521,7 +521,14 @@ The spec types are:
        ``AttnBackend.DENSE``, ``AttnBackend.FLEX`` or ``AttnBackend.FLASHINFER_MLA``.
        Leave it at ``AUTO`` unless the model needs a specific kernel. FlashInfer, dense
        and MLA attention require a ``PagedKVConfig``; FlexAttention requires a
-       ``RingKVConfig``.
+       ``RingKVConfig``. ``DENSE`` runs one varlen pass over each stream's frozen
+       prefix (gathered from the pages once) and the step's fresh K/V, eager only:
+       FA3 (AITER on ROCm), or FlashInfer's ragged prefill where FA3 has no kernel for
+       the GPU, the same choice the ragged resources make (``dense_or_flashinfer``), with
+       ``flashinfer_backend`` picking its kernel. The engine plans and runs a tiny layout on
+       that kernel when it builds the resource, so a backend the GPU or FlashInfer build
+       cannot run stops the server with a ``ValueError`` naming it. Its wrappers share one
+       workspace, taken at build.
        ``flashinfer_backend`` selects a kernel generation: ``"auto"``, ``"fa2"`` or
        ``"fa3"`` when the FlashInfer backend is selected. ``decode_tensor_cores``
        (default ``True``) runs one-token decode rows on FlashInfer's tensor-core
@@ -1436,9 +1443,15 @@ Both types share the base ``CudaGraphConfig`` fields:
 - ``compile_mode`` is that ``torch.compile``'s mode. The default, ``None``, takes the
   deployment's compile level (``mstar serve -O``, ``MSTAR_COMPILE_LEVEL``); naming a
   mode takes the config out of that level. ``default`` keeps Inductor's fusion but takes cuBLAS for
-  every GEMM and skips autotuning. Building the config refuses an unknown mode, and one
-  that turns on Inductor's own CUDA graphs (``reduce-overhead``, ``max-autotune``), which
-  fail inside the runner's capture.
+  every GEMM and skips autotuning. ``deterministic`` is the engine's own: Inductor's
+  defaults in its deterministic mode, which picks kernel configs by rule rather than timing
+  them on the device, so every process compiles the same kernels (a torch without the mode
+  logs a warning and compiles with the defaults). A compile outside the runner takes the same
+  options from ``accelerator_graph_runner.compile_kwargs("deterministic")``, or
+  ``deterministic_inductor_options()`` where it builds its own; the Wan VAE decode, the DiT
+  scaffold's ``compile_blocks`` and Cosmos3's denoise step do. Building the config refuses
+  an unknown mode, and one that turns on Inductor's own CUDA graphs (``reduce-overhead``,
+  ``max-autotune``), which fail inside the runner's capture.
 - ``required`` makes every bucket in the config mandatory. If capture fails locally or
   on another participating rank, warmup raises after rank-wide agreement instead of
   dropping the bucket and falling back to eager execution. The default is ``False``.

docs/models.rst

@@ -929,11 +929,21 @@ Cosmos3 environment requirements
   (AITER's varlen kernel on ROCm).
   That wheel is ABI-tied to the installed torch/CUDA build (Hopper builds
   exist for at least torch 2.9 + cu12.8 and torch 2.11 + cu13.0); install the
-  one matching your environment. When it is not importable, the engine logs a
-  warning at startup and automatically falls back to the paged ``flashinfer``
-  backend — serving still works, only the bs=1 dense fast path is lost.
-  ``model_kwargs.attention_backend: flashinfer`` in the config YAML selects
-  the paged backend explicitly.
+  one matching your environment. It ships kernels for sm_80 and sm_90a only.
+  Where it is not importable or has no kernel for the GPU (Blackwell), the
+  same dense pass runs on FlashInfer's ragged prefill, whose ``auto`` kernel on
+  a B300 is CUTLASS's FMHA: 1.5 ms per layer for the 832x480x121 guidance pair
+  against 6.3 ms for the paged prefill. ``resources: attn_gen:
+  flashinfer_backend:`` picks that kernel (``cudnn``, ``cutlass``, ``fa2``); the
+  server checks it at startup with a tiny plan and run, and a backend this GPU or
+  FlashInfer build cannot run (a typo, ``fa3`` on Blackwell, ``cutlass`` on
+  Hopper) stops it with a ``ValueError`` naming the backend.
+  Without a usable FlashInfer either, the engine logs a warning and switches to
+  the paged backend, but that is FlashInfer too (the default backend on every
+  platform for now), so it cannot run there either: the dense pass needs
+  fa3-fwd (AITER on ROCm) or FlashInfer.
+  ``model_kwargs.attention_backend: flashinfer`` in the config YAML selects the
+  paged backend explicitly.
 - Video-input requests (video-to-video, action inverse-dynamics) decode the
   conditioning clip with ``torchcodec``; environments without it reject those
   requests at preprocessing (other modes are unaffected).
@@ -996,6 +1006,32 @@ so an image or clip produced alongside other requests differs from its solo
 result at the kernel-drift level (~30 dB PSNR at guidance 6). Serve with one
 request at a time when outputs must be bitwise repeatable.
 
+The denoise step is compiled (``compile_denoise``) in Inductor's deterministic
+mode by default (``denoise_compile_mode: deterministic``, the engine's compile
+mode of that name): kernel configs are picked by rule, so a seed gives the same
+video from every server process. ``default`` is torch.compile's own defaults,
+and ``level`` takes the ``mstar serve -O`` level's options (``-O2``:
+max-autotune).
+
+Video clips decode through the shared Wan VAE decoder (``vae_decode_fast``):
+one decoder call per latent, channels-last, in the decode dtype (bf16 with
+cuDNN >= 9.16, otherwise fp32; see above), on the decoder-only copy of the VAE
+the channels-last decoders on a GPU share. The first two calls (the clip-start latent and the one after it)
+run eager; every later call is compiled in Inductor's deterministic mode. An
+image (one latent) keeps the compiled stock ``vae.decode`` on that copy, which
+is faster for a single frame. False decodes everything with the compiled stock
+``vae.decode`` on the shared instance.
+
+Two or more requests at the same denoise step run one cross-request forward
+(``max_gen_batch_size``, default 8) for video, image and sound; the action step
+runs eager. With ``compile_denoise_batched`` (default on) its layer stack is
+compiled with the same mode, once, while the server starts, with the packed
+length dynamic: no batch size or mix of resolutions compiles mid-traffic. The
+per-request patchify, rotary tables and unpatchify around it run eager. At
+832x480x121 one request's guidance pair is 24k tokens, which already fills the
+GPU's GEMM and attention throughput, so a batch of N costs about N single
+steps: batching keeps requests in step rather than adding throughput.
+
 The action policy (``cosmos3_edge_droid``, or ``cosmos3_edge`` with an action
 ``domain_name``) predicts a chunk of robot actions from the current observation:
 ``output_modalities=action`` with ``model_kwargs``

mstar/engine/accelerator_graph_config.py

@@ -8,11 +8,18 @@
 from mstar.engine.resources import SubmoduleStep
 from mstar.model.submodule_base import ModelInputsFromEngine, NodeInputs
 
+# A compile mode of the engine's own rather than torch.compile's: Inductor's
+# defaults in its deterministic mode, which picks each kernel's config by rule
+# instead of timing candidates on the device, so every process compiles the
+# same kernels and one input gives one output from any server. See
+# ``accelerator_graph_runner.compile_kwargs``.
+DETERMINISTIC_COMPILE_MODE = "deterministic"
+
 
 def check_compile_mode(mode: str | None) -> str | None:
     """A capture config's ``compile_mode``, refused at build rather than at capture,
     where the per-capture fallback would serve the bucket eager."""
-    if mode is None or mode == "default":
+    if mode is None or mode in ("default", DETERMINISTIC_COMPILE_MODE):
         return mode
     if not mode:
         # torch.compile takes a falsy mode for "default"

mstar/engine/accelerator_graph_runner.py

@@ -17,6 +17,7 @@
     CapturedGraph,
 )
 from mstar.engine.accelerator_graph_config import (
+    DETERMINISTIC_COMPILE_MODE,
     AcceleratorGraphConfig,
     AcceleratorGraphConfigType,
     PiecewiseAcceleratorGraphConfig,
@@ -69,13 +70,37 @@ def compile_level() -> int:
     return level
 
 
+@cache
+def _has_deterministic_mode() -> bool:
+    import torch._inductor.config as inductor_config
+
+    if hasattr(inductor_config, "deterministic"):
+        return True
+    logger.warning(
+        "torch %s has no Inductor deterministic mode; compiles that ask for it time "
+        "their kernels instead, so their output can differ between server processes",
+        torch.__version__,
+    )
+    return False
+
+
+def deterministic_inductor_options() -> dict[str, Any]:
+    """Inductor options for its deterministic mode (see
+    ``DETERMINISTIC_COMPILE_MODE``), for a compile that sets its own options;
+    empty, with a warning, on a torch without the mode."""
+    return {"deterministic": True} if _has_deterministic_mode() else {}
+
+
 def compile_kwargs(config_mode: str | None) -> dict[str, Any]:
     """``torch.compile``'s arguments for a capture config whose ``compile_mode``
     is ``config_mode``, which the config checked when it was built. None takes
     the compile level's arguments."""
     kwargs: dict[str, Any] = {"fullgraph": False, "dynamic": False}
     if config_mode is None:
         return {**kwargs, **COMPILE_LEVELS[compile_level()]}
+    if config_mode == DETERMINISTIC_COMPILE_MODE:
+        options = deterministic_inductor_options()
+        return {**kwargs, "options": options} if options else kwargs
     # "default" is torch.compile's own: Inductor fusion, cuBLAS GEMMs, no autotuning
     if config_mode != "default":
         kwargs["mode"] = config_mode

mstar/engine/engine.py

@@ -732,6 +732,9 @@ def warmup(self) -> None:
         # torch.compile applied after CUDA graph capture because the cuda
         # graph runner compiles internally
         self._compile_submodules()
+        for submodule_mgmt in self._submodules.values():
+            with fp32_precision_scope(submodule_mgmt.fp32_precision):
+                submodule_mgmt.submodule.warmup_compiles()
 
         for resource in self._resources.values():
             resource.post_warmup_validate()

mstar/engine/resources/attn/base.py

@@ -51,8 +51,8 @@ class AttentionManager(AttentionResource):
     # Label / layer cursors come from `AttentionResource`; `run` resolves them.
 
     def fingerprint(self) -> bytes:
-        # the class, not the declared backend: `build` degrades dense to paged
-        # where the FA3 kernel is unusable, and the two do not agree bit for bit
+        # the class, not the declared backend: `build` moves dense to another
+        # kernel where the FA3 one is unusable, and they do not agree bit for bit
         return fingerprint(type(self).__name__)
 
     @classmethod
@@ -86,26 +86,34 @@ def build(cls, spec: AttentionSpec, info: EngineResourceInfo):
         if info.joint_comm_group is not None:
             kv_config.shard(info.joint_comm_group.world_size)
         if backend == AttnBackend.DENSE:
-            # A dense backend needs the FlashAttention-3 kernel; where the
-            # wheel does not match the installed torch/CUDA build or has no
-            # kernel for the GPU, degrade to the paged backend AUTO picks
-            # rather than failing to serve.
-            # The two are drop-in for each other: the step declaration is
+            # A dense backend runs on FlashAttention-3 (AITER on ROCm). Where
+            # that has no kernel for the GPU or does not import, the same pass
+            # runs on FlashInfer's ragged prefill, and where FlashInfer cannot
+            # run either it degrades to the paged backend AUTO picks rather
+            # than failing to serve.
+            # All three are drop-in for each other: the step declaration is
             # identical (see DenseAttentionManager) and `requires_kv_write`
             # tells the layer which one it is talking to.
             from mstar.engine.resources.attn.dense import (
                 DenseAttentionManager,
-                _fa3_unavailable_reason,
+                _dense_unavailable_reason,
                 _warn_dense_fallback,
             )
+            from mstar.engine.resources.attn.ragged.base import dense_or_flashinfer
+            from mstar.engine.resources.attn.ragged.dense_wrappers import DensePrefixedPrefillWrapper
+            from mstar.engine.resources.attn.ragged.wrappers import FlashInferPrefixedPrefillWrapper
 
-            reason = _fa3_unavailable_reason(info.device)
+            reason = _dense_unavailable_reason(info.device)
             if reason is None:
                 return DenseAttentionManager(
                     kv_cache=spec.config.kv_cache,
                     device=info.device,
                     dtype=info.kv_dtype,
                     kv_config=kv_config,
+                    wrapper_class=dense_or_flashinfer(
+                        info.device, DensePrefixedPrefillWrapper, FlashInferPrefixedPrefillWrapper, "Attention",
+                    ),
+                    flashinfer_backend=spec.config.flashinfer_backend,
                 )
             backend = resolve_attn_backend(AttnBackend.AUTO, info.device)
             _warn_dense_fallback(reason, backend)
@@ -221,6 +229,11 @@ def captured_page_capacity(num_rows: int, num_tokens: int, kv_config: KVConfig)
     )
 
 
+def workspace_bytes() -> int:
+    """Size of one FlashInfer workspace buffer (``MSTAR_WORKSPACE_BUFFER_MB``)."""
+    return int(os.environ.get("MSTAR_WORKSPACE_BUFFER_MB", "512")) * 1024 * 1024
+
+
 class WorkspacePool:
     """FlashInfer workspace buffers, one per (plan label, cg slot).
 
@@ -230,9 +243,7 @@ class WorkspacePool:
 
     def __init__(self, device: torch.device):
         self._device = device
-        self._size = int(
-            os.environ.get("MSTAR_WORKSPACE_BUFFER_MB", "512")
-        ) * 1024 * 1024
+        self._size = workspace_bytes()
         self._buffers: dict[str, torch.Tensor] = {}
 
     def get(self, label: str, cg_slot: int | None = None) -> torch.Tensor: