diff --git a/mstar/engine/resources/attn/ragged/base.py b/mstar/engine/resources/attn/ragged/base.py index f4f0a35131..ad6b3b48bd 100644 --- a/mstar/engine/resources/attn/ragged/base.py +++ b/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, ) diff --git a/mstar/engine/resources/attn/ragged/config.py b/mstar/engine/resources/attn/ragged/config.py index 6e958aa63e..07d15bbe65 100644 --- a/mstar/engine/resources/attn/ragged/config.py +++ b/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 diff --git a/test/modular/test_ragged_attention_config.py b/test/modular/test_ragged_attention_config.py new file mode 100644 index 0000000000..102b1e4679 --- /dev/null +++ b/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))