engine: let a ragged attention node declare its activation dtype

mstar/engine/resources/attn/ragged/base.py

@@ -26,8 +26,15 @@ def build(cls, spec: RaggedAttentionSpec, info: EngineResourceInfo):
         config = spec.config
         if info.joint_comm_group is not None:
             config.shard(info.joint_comm_group.world_size)
+        dtype = config.dtype if config.dtype is not None else info.kv_dtype
+        if dtype is None:
+            raise ValueError(
+                f"ragged attention {spec.resource_key!r}: no activation dtype. Set "
+                "RaggedAttentionConfig.dtype on a node without a KV-backed attention "
+                "(the engine has no KV dtype to fall back on)."
+            )
         return FlashInferRaggedManager(
             device=info.device,
-            dtype=info.kv_dtype,
+            dtype=dtype,
             config=config,
         )

mstar/engine/resources/attn/ragged/config.py

@@ -7,6 +7,8 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING
 
+import torch
+
 from mstar.engine.resources.spec import NodeResourceSpec
 
 if TYPE_CHECKING:
@@ -41,6 +43,12 @@ class RaggedAttentionConfig:
 
     flashinfer_backend: str = "auto"
 
+    # Activation dtype of the q / k / v the node hands the kernel (and of its
+    # output). A cacheless attention doesn't need to adhere to the dtype of its
+    # KV cache, so it is free to declare a separate dtype here; when None, the
+    # engine's KV dtype (which is actually a misnamed autocast dtype) is used.
+    dtype: torch.dtype | None = None
+
     def __post_init__(self):
         if self.sm_scale is None:
             self.sm_scale = self.head_dim ** -0.5

test/modular/test_ragged_attention_config.py

@@ -0,0 +1,48 @@
+"""CPU tests for the ragged attention resource's activation dtype.
+
+A cacheless attention has no KV cache to inherit a dtype from, so a node without a
+KV-backed attention (a DiT) declares it on the config; the engine's KV dtype stays the
+fallback for nodes that have one, and neither being set is an error, not an fp32 kernel.
+"""
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+import torch
+
+sys.path.insert(0, ".")
+
+from mstar.engine.resources import RaggedAttentionConfig, RaggedAttentionSpec  # noqa: E402
+from mstar.engine.resources.attn.ragged.base import RaggedAttnManager  # noqa: E402
+from mstar.engine.resources.base import EngineResourceInfo  # noqa: E402
+
+CPU = torch.device("cpu")
+
+
+def _spec(dtype=None) -> RaggedAttentionSpec:
+    return RaggedAttentionSpec(
+        resource_key="dit_attn", nodes={"dit"},
+        config=RaggedAttentionConfig(num_qo_heads=2, num_kv_heads=2, head_dim=128, dtype=dtype),
+    )
+
+
+def test_config_dtype_reaches_the_kernel_without_a_kv_dtype():
+    manager = RaggedAttnManager.build(_spec(torch.bfloat16), EngineResourceInfo(device=CPU, kv_dtype=None))
+    assert manager._kwargs["q_data_type"] is torch.bfloat16
+
+
+def test_config_dtype_wins_over_the_engine_kv_dtype():
+    manager = RaggedAttnManager.build(_spec(torch.float16), EngineResourceInfo(device=CPU, kv_dtype=torch.bfloat16))
+    assert manager._kwargs["q_data_type"] is torch.float16
+
+
+def test_engine_kv_dtype_is_the_fallback():
+    manager = RaggedAttnManager.build(_spec(None), EngineResourceInfo(device=CPU, kv_dtype=torch.bfloat16))
+    assert manager._kwargs["q_data_type"] is torch.bfloat16
+
+
+def test_no_dtype_anywhere_is_an_error():
+    with pytest.raises(ValueError, match="no activation dtype"):
+        RaggedAttnManager.build(_spec(None), EngineResourceInfo(device=CPU, kv_dtype=None))