diff --git a/nemo_rl/algorithms/grpo.py b/nemo_rl/algorithms/grpo.py index dcf14a3d940..c667b29b479 100644 --- a/nemo_rl/algorithms/grpo.py +++ b/nemo_rl/algorithms/grpo.py @@ -93,7 +93,10 @@ run_async_nemo_gym_rollout, run_multi_turn_rollout, ) -from nemo_rl.models.generation.interfaces import GenerationInterface +from nemo_rl.models.generation.interfaces import ( + GenerationInterface, + resolve_routed_experts_dtype_name_for_model, +) from nemo_rl.models.generation.megatron import MegatronGeneration from nemo_rl.models.generation.sglang.config import SGLangConfig from nemo_rl.models.generation.sglang.sglang_generation import SGLangGeneration @@ -588,6 +591,11 @@ def _spinup_nemo_gym(base_urls, model_name): invalid_tool_call_patterns=invalid_tool_call_patterns, thinking_tags=thinking_tags, require_routed_experts=router_replay_enabled(policy_config), + routed_experts_dtype=( + resolve_routed_experts_dtype_name_for_model(model_name) + if router_replay_enabled(policy_config) + else "int16" + ), initial_global_config_dict=nemo_gym_dict, ) nemo_gym_opts = {} diff --git a/nemo_rl/environments/nemo_gym.py b/nemo_rl/environments/nemo_gym.py index b3f8dcbfbd1..a5243087cdb 100644 --- a/nemo_rl/environments/nemo_gym.py +++ b/nemo_rl/environments/nemo_gym.py @@ -29,6 +29,15 @@ from nemo_rl.environments.interfaces import EnvironmentInterface from nemo_rl.utils.timer import Timer +# Kept local (not imported from models.generation) so the gym actor stays free of +# generation-module imports. Must cover every name resolve_routed_experts_dtype +# can produce. +_ROUTED_EXPERTS_DTYPES = { + "int8": torch.int8, + "int16": torch.int16, + "int32": torch.int32, +} + DEFAULT_INVALID_TOOL_CALL_PATTERNS = [ "", "", @@ -80,6 +89,9 @@ class NemoGymConfig(TypedDict): require_routed_experts: NotRequired[ bool ] # Require Gym output items to carry R3 routed_experts + routed_experts_dtype: NotRequired[ + str + ] # Carry dtype name for routed_experts tensors ("int8"/"int16"/"int32"), resolved from the model's expert count def _detect_invalid_tool_call_and_malformed_thinking( @@ -351,7 +363,12 @@ def _postprocess_nemo_gym_to_nemo_rl_result( routed_experts = None if routed_experts_raw is not None: - routed_experts = torch.as_tensor(routed_experts_raw, dtype=torch.int32) + routed_experts_dtype = _ROUTED_EXPERTS_DTYPES[ + self.cfg.get("routed_experts_dtype", "int16") + ] + routed_experts = torch.as_tensor( + routed_experts_raw, dtype=routed_experts_dtype + ) if routed_experts.dim() != 3: raise ValueError( "NeMo Gym returned routed_experts with invalid shape. " diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 025707c7001..4584ee0b1e5 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/nemo_rl/models/generation/interfaces.py @@ -12,13 +12,75 @@ # See the License for the specific language governing permissions and # limitations under the License. from abc import ABC, abstractmethod -from typing import Any, NotRequired, TypedDict, Union +from typing import Any, NotRequired, Optional, TypedDict, Union import ray import torch from nemo_rl.distributed.batched_data_dict import BatchedDataDict +# Routed-expert index tensors ([seq, layers, topk]) are carried in the narrowest +# signed dtype that fits ids 0..num_experts-1 plus the -1 missing-route sentinel: +# int8 for <=128 experts (e.g. Qwen3-MoE), int16 for <=32768 (e.g. DeepSeek-V3), +# int32 beyond. This shrinks message logs, transports, replay buffers, and +# checkpoints 2-4x vs int32. The Megatron replay install converts to int64 at the +# gather site, so training math is unaffected. When the expert count cannot be +# determined, fall back to int16. +ROUTED_EXPERTS_FALLBACK_DTYPE = torch.int16 + +_ROUTED_EXPERTS_DTYPE_NAMES = { + torch.int8: "int8", + torch.int16: "int16", + torch.int32: "int32", +} + + +def get_num_routed_experts(hf_config: Any) -> Optional[int]: + """Best-effort read of the routed-expert count from a HF model config. + + Checks the attribute names used by the common MoE architectures (Qwen-MoE, + DeepSeek, Mixtral), including nested ``text_config`` for VLMs. Returns None + for dense models or unrecognized configs. + """ + for owner in (hf_config, getattr(hf_config, "text_config", None)): + if owner is None: + continue + for attr in ("num_experts", "n_routed_experts", "num_local_experts"): + value = getattr(owner, attr, None) + if isinstance(value, int) and value > 0: + return value + return None + + +def resolve_routed_experts_dtype(num_experts: Optional[int]) -> torch.dtype: + """Return the narrowest signed dtype that fits expert ids and the -1 sentinel.""" + if num_experts is None: + return ROUTED_EXPERTS_FALLBACK_DTYPE + if num_experts - 1 <= torch.iinfo(torch.int8).max: + return torch.int8 + if num_experts - 1 <= torch.iinfo(torch.int16).max: + return torch.int16 + return torch.int32 + + +def resolve_routed_experts_dtype_name_for_model(model_name: str) -> str: + """Resolve the routed-experts carry dtype name ("int8"/"int16"/"int32") for a model. + + Used where only the model name is available (e.g. building the NeMo-Gym env + config on the driver). Falls back to the default dtype name if the config + cannot be loaded. + """ + # Deferred import: transformers config loading is only needed for this sizing. + from transformers import AutoConfig + + try: + hf_config = AutoConfig.from_pretrained(model_name, trust_remote_code=True) + except (OSError, ValueError): + return _ROUTED_EXPERTS_DTYPE_NAMES[ROUTED_EXPERTS_FALLBACK_DTYPE] + return _ROUTED_EXPERTS_DTYPE_NAMES[ + resolve_routed_experts_dtype(get_num_routed_experts(hf_config)) + ] + def verify_right_padding( data: Union[ diff --git a/nemo_rl/models/generation/vllm/utils.py b/nemo_rl/models/generation/vllm/utils.py index 2125de3cff4..4d389b3aedb 100644 --- a/nemo_rl/models/generation/vllm/utils.py +++ b/nemo_rl/models/generation/vllm/utils.py @@ -18,10 +18,42 @@ import torch from nemo_rl.distributed.batched_data_dict import BatchedDataDict -from nemo_rl.models.generation.interfaces import GenerationDatumSpec +from nemo_rl.models.generation.interfaces import ( + ROUTED_EXPERTS_FALLBACK_DTYPE, + GenerationDatumSpec, +) R3_MISSING_ROUTE_SENTINEL = -1 +# The expert-id range vs carry dtype is model-constant, so it is verified on the +# first non-empty routed-experts tensor per process and skipped afterwards. +G_ROUTED_EXPERTS_RANGE_CHECKED = False + + +def _as_routed_experts_tensor( + value: Any, *, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Convert backend routed-expert ids to the resolved carry dtype. + + Guards against expert ids overflowing ``dtype`` before the narrowing cast, + which would otherwise wrap silently (e.g. if the expert count was + mis-detected when resolving the dtype). + """ + global G_ROUTED_EXPERTS_RANGE_CHECKED + tensor = torch.as_tensor(value, device=device) + if not G_ROUTED_EXPERTS_RANGE_CHECKED and tensor.numel() > 0: + max_id = int(tensor.max()) + limit = torch.iinfo(dtype).max + if max_id > limit: + raise ValueError( + f"routed expert id {max_id} exceeds the resolved carry dtype " + f"{dtype} (max {limit}); the model's expert count was likely " + "mis-detected (see resolve_routed_experts_dtype in " + "nemo_rl.models.generation.interfaces)." + ) + G_ROUTED_EXPERTS_RANGE_CHECKED = True + return tensor.to(dtype=dtype) + def format_prompt_for_vllm_generation( data: BatchedDataDict[GenerationDatumSpec], sample_idx: Optional[int] = None @@ -114,15 +146,20 @@ def pad_and_align_routed_expert_indices( require_complete_routed_experts: bool = False, allow_missing_routed_experts_fallback: bool = True, return_stats: bool = False, + routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE, ) -> Optional[torch.Tensor] | tuple[Optional[torch.Tensor], dict[str, int]]: - """Return full-sequence-aligned routed experts as ``[S, L, topk]`` int32.""" + """Return full-sequence-aligned routed experts as ``[S, L, topk]`` in ``routed_experts_dtype``.""" routed = getattr(completion_output, "routed_experts", None) prompt_routed = getattr(request_output, "prompt_routed_experts", None) if prompt_routed is not None: - prompt_routed = torch.as_tensor(prompt_routed, dtype=torch.int32, device=device) + prompt_routed = _as_routed_experts_tensor( + prompt_routed, device=device, dtype=routed_experts_dtype + ) if routed is not None: - routed = torch.as_tensor(routed, dtype=torch.int32, device=device) + routed = _as_routed_experts_tensor( + routed, device=device, dtype=routed_experts_dtype + ) if prompt_routed is not None and routed is not None: routed = torch.cat((prompt_routed, routed), dim=0) @@ -180,7 +217,7 @@ def pad_and_align_routed_expert_indices( default_route = torch.arange( routed.shape[2], - dtype=torch.int32, + dtype=routed_experts_dtype, device=device, ) full = ( @@ -188,7 +225,6 @@ def pad_and_align_routed_expert_indices( .expand(padded_length, routed.shape[1], routed.shape[2]) .clone() ) - full = full.to(dtype=torch.int32) routes_to_copy = min(expected_routes, routed.shape[0]) if routes_to_copy > 0: full[:routes_to_copy] = routed[:routes_to_copy].to(device=device) @@ -203,6 +239,7 @@ def attach_routed_experts_to_chat_response_choices( *, device: torch.device, logger: Any = None, + routed_experts_dtype: torch.dtype = ROUTED_EXPERTS_FALLBACK_DTYPE, ) -> Any: """Attach aligned routed experts to OpenAI chat response choices.""" outputs_by_index = { @@ -229,6 +266,7 @@ def attach_routed_experts_to_chat_response_choices( device=device, require_complete_routed_experts=True, return_stats=True, + routed_experts_dtype=routed_experts_dtype, ) if not isinstance(routed_result, tuple): raise RuntimeError( @@ -254,7 +292,9 @@ def attach_routed_experts_to_chat_response_choices( r3_stats["actual_routes"], r3_stats["expected_routes"], ) - choice.message.routed_experts = routed_experts.to(dtype=torch.int32).tolist() + choice.message.routed_experts = routed_experts.to( + dtype=routed_experts_dtype + ).tolist() if len(attached_choice_indices) != len(choices): missing_choice_indices = sorted( diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index 0d6a68b9f30..eed09f8c642 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -30,8 +30,11 @@ ) from nemo_rl.distributed.worker_group_utils import get_nsight_config_if_pattern_matches from nemo_rl.models.generation.interfaces import ( + ROUTED_EXPERTS_FALLBACK_DTYPE, GenerationDatumSpec, GenerationOutputSpec, + get_num_routed_experts, + resolve_routed_experts_dtype, verify_right_padding, ) from nemo_rl.models.generation.vllm.config import VllmConfig @@ -203,6 +206,8 @@ def _init_config( """Lightweight config setup. No model loading, no heavy imports.""" self.cfg = config self.model_name = self.cfg["model_name"] + # Refined from the model's expert count in _load_model. + self.routed_experts_dtype = ROUTED_EXPERTS_FALLBACK_DTYPE self.tensor_parallel_size = self.cfg["vllm_cfg"]["tensor_parallel_size"] self.pipeline_parallel_size = self.cfg["vllm_cfg"]["pipeline_parallel_size"] self.expert_parallel_size = self.cfg["vllm_cfg"]["expert_parallel_size"] @@ -349,6 +354,9 @@ def _load_model(self, bundle_indices, seed): # Override HF config for gpt-oss models to ensure compatibility with megatron # The megatron --> hf export is done in bf16, so we disable quantization hf_config = AutoConfig.from_pretrained(self.model_name, trust_remote_code=True) + self.routed_experts_dtype = resolve_routed_experts_dtype( + get_num_routed_experts(hf_config) + ) if "GptOssForCausalLM" in getattr(hf_config, "architectures", []): if "quantization_config" in hf_config: assert load_format == "dummy", ( @@ -734,6 +742,7 @@ def generate( device=input_ids.device, require_complete_routed_experts=return_routed_experts, return_stats=True, + routed_experts_dtype=self.routed_experts_dtype, ) if return_routed_experts and full_routed_experts is None: raise RuntimeError( diff --git a/nemo_rl/models/generation/vllm/vllm_worker_async.py b/nemo_rl/models/generation/vllm/vllm_worker_async.py index fab2e1330e8..f30abb5f774 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker_async.py +++ b/nemo_rl/models/generation/vllm/vllm_worker_async.py @@ -717,6 +717,7 @@ async def capture_result_generator(): final_res, device=torch.device("cpu"), logger=LOGGER, + routed_experts_dtype=worker_self.routed_experts_dtype, ) class NeMoRLOpenAIServingChat(NeMoRLOpenAIServingChatMixin, OpenAIServingChat): @@ -1212,6 +1213,7 @@ async def process_single_sample(sample_idx): device=original_input_ids_single_row.device, require_complete_routed_experts=return_routed_experts, return_stats=True, + routed_experts_dtype=self.routed_experts_dtype, ) if return_routed_experts and routed_experts is None: raise RuntimeError( diff --git a/tests/unit/environments/test_nemo_gym_router_replay.py b/tests/unit/environments/test_nemo_gym_router_replay.py index fdc7a021f7e..d1171856213 100644 --- a/tests/unit/environments/test_nemo_gym_router_replay.py +++ b/tests/unit/environments/test_nemo_gym_router_replay.py @@ -88,3 +88,34 @@ class _MockSelf: NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( _MockSelf(), nemo_gym_result, _Tokenizer() ) + + +def test_nemo_gym_postprocess_casts_routed_experts_to_configured_dtype(): + import torch + + nemo_gym_result = { + "response": { + "output": [ + { + "prompt_token_ids": [1, 2], + "generation_token_ids": [3], + "generation_log_probs": [-0.1], + "routed_experts": _routes(3), + }, + ] + }, + "responses_create_params": {"input": []}, + } + + class _MockSelf: + cfg = {"require_routed_experts": True, "routed_experts_dtype": "int8"} + + result = ( + NemoGym.__ray_metadata__.modified_class._postprocess_nemo_gym_to_nemo_rl_result( + _MockSelf(), nemo_gym_result, _Tokenizer() + ) + ) + + for message in result["message_log"]: + if "routed_experts" in message: + assert message["routed_experts"].dtype == torch.int8 diff --git a/tests/unit/models/generation/test_vllm_utils.py b/tests/unit/models/generation/test_vllm_utils.py index 9216718c649..e4e033e0a99 100644 --- a/tests/unit/models/generation/test_vllm_utils.py +++ b/tests/unit/models/generation/test_vllm_utils.py @@ -20,6 +20,12 @@ import torch from nemo_rl.distributed.batched_data_dict import BatchedDataDict +from nemo_rl.models.generation.interfaces import ( + ROUTED_EXPERTS_FALLBACK_DTYPE, + get_num_routed_experts, + resolve_routed_experts_dtype, +) +from nemo_rl.models.generation.vllm import utils as vllm_utils from nemo_rl.models.generation.vllm.utils import ( R3_MISSING_ROUTE_SENTINEL, aggregate_spec_decode_counters, @@ -223,11 +229,14 @@ class Output: ) assert routed_experts.shape == (8, 3, 2) - assert routed_experts.dtype == torch.int32 + assert routed_experts.dtype == ROUTED_EXPERTS_FALLBACK_DTYPE assert torch.equal( - routed_experts[:5], completion_output.routed_experts.to(torch.int32) + routed_experts[:5], + completion_output.routed_experts.to(ROUTED_EXPERTS_FALLBACK_DTYPE), ) - expected_default_route = torch.tensor([0, 1], dtype=torch.int32).view(1, 1, 2) + expected_default_route = torch.tensor( + [0, 1], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ).view(1, 1, 2) assert torch.equal(routed_experts[5:], expected_default_route.expand(3, 3, 2)) @@ -237,8 +246,12 @@ class Output: request_output = Output() completion_output = Output() - request_output.prompt_routed_experts = torch.ones(2, 1, 2, dtype=torch.int32) - completion_output.routed_experts = 2 * torch.ones(3, 1, 2, dtype=torch.int32) + request_output.prompt_routed_experts = torch.ones( + 2, 1, 2, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ) + completion_output.routed_experts = 2 * torch.ones( + 3, 1, 2, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ) routed_experts = pad_and_align_routed_expert_indices( request_output, @@ -248,7 +261,9 @@ class Output: device=torch.device("cpu"), ) - expected_default_route = torch.tensor([0, 1], dtype=torch.int32).view(1, 1, 2) + expected_default_route = torch.tensor( + [0, 1], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ).view(1, 1, 2) assert torch.equal(routed_experts[:2], request_output.prompt_routed_experts) assert torch.equal(routed_experts[2:4], completion_output.routed_experts[:2]) assert torch.equal(routed_experts[4:], expected_default_route.expand(1, 1, 2)) @@ -265,7 +280,7 @@ class Output: [[4, 5, 6], [7, 8, 9]], [[1, 2, 3], [10, 11, 12]], ], - dtype=torch.int32, + dtype=ROUTED_EXPERTS_FALLBACK_DTYPE, ) routed_experts = pad_and_align_routed_expert_indices( @@ -276,7 +291,9 @@ class Output: device=torch.device("cpu"), ) - expected_default_route = torch.tensor([0, 1, 2], dtype=torch.int32).view(1, 1, 3) + expected_default_route = torch.tensor( + [0, 1, 2], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ).view(1, 1, 3) assert torch.equal(routed_experts[:2], completion_output.routed_experts) assert torch.equal(routed_experts[2:], expected_default_route.expand(3, 2, 3)) @@ -293,7 +310,7 @@ class Output: [[1, 2, 3], [10, 11, 12]], [[0, 0, 0], [0, 0, 0]], ], - dtype=torch.int32, + dtype=ROUTED_EXPERTS_FALLBACK_DTYPE, ) routed_experts = pad_and_align_routed_expert_indices( @@ -304,7 +321,9 @@ class Output: device=torch.device("cpu"), ) - expected_default_route = torch.tensor([0, 1, 2], dtype=torch.int32).view(1, 1, 3) + expected_default_route = torch.tensor( + [0, 1, 2], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ).view(1, 1, 3) assert torch.equal(routed_experts[:2], completion_output.routed_experts[:2]) assert torch.equal(routed_experts[2:], expected_default_route.expand(1, 2, 3)) @@ -316,7 +335,9 @@ class Output: request_output = Output() request_output.num_cached_tokens = 4 completion_output = Output() - completion_output.routed_experts = torch.ones(2, 1, 2, dtype=torch.int32) + completion_output.routed_experts = torch.ones( + 2, 1, 2, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ) routed_experts, stats = pad_and_align_routed_expert_indices( request_output, @@ -337,9 +358,13 @@ class Output: assert torch.equal(routed_experts[:2], completion_output.routed_experts) assert torch.equal( routed_experts[2:5], - torch.full((3, 1, 2), R3_MISSING_ROUTE_SENTINEL, dtype=torch.int32), + torch.full( + (3, 1, 2), R3_MISSING_ROUTE_SENTINEL, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), ) - expected_default_route = torch.tensor([0, 1], dtype=torch.int32).view(1, 1, 2) + expected_default_route = torch.tensor( + [0, 1], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ).view(1, 1, 2) assert torch.equal(routed_experts[5:], expected_default_route) @@ -350,7 +375,9 @@ class Output: request_output = Output() request_output.num_cached_tokens = 4 completion_output = Output() - completion_output.routed_experts = torch.ones(2, 1, 2, dtype=torch.int32) + completion_output.routed_experts = torch.ones( + 2, 1, 2, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ) with pytest.raises(ValueError, match="incomplete routed_experts"): pad_and_align_routed_expert_indices( @@ -371,7 +398,9 @@ class Output: request_output = Output() request_output.num_cached_tokens = 0 completion_output = Output() - completion_output.routed_experts = torch.ones(4, 1, 2, dtype=torch.int32) + completion_output.routed_experts = torch.ones( + 4, 1, 2, dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ) with pytest.raises(ValueError, match="too many routed_experts routes"): pad_and_align_routed_expert_indices( @@ -387,17 +416,23 @@ class Output: def test_attach_routed_experts_to_chat_response_choices_reassociates_by_choice_index(): final_res = SimpleNamespace( prompt_token_ids=[101, 102, 103], - prompt_routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + prompt_routed_experts=torch.tensor( + [[[10]], [[11]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), outputs=[ SimpleNamespace( index=1, token_ids=[201, 202], - routed_experts=torch.tensor([[[31]], [[32]]], dtype=torch.int32), + routed_experts=torch.tensor( + [[[31]], [[32]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), ), SimpleNamespace( index=0, token_ids=[200], - routed_experts=torch.tensor([[[30]]], dtype=torch.int32), + routed_experts=torch.tensor( + [[[30]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), ), ], ) @@ -453,7 +488,9 @@ def test_attach_routed_experts_to_chat_response_choices_warns_on_missing_routes( SimpleNamespace( index=0, token_ids=[200, 201], - routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + routed_experts=torch.tensor( + [[[10]], [[11]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), ) ], ) @@ -496,7 +533,9 @@ def test_attach_routed_experts_to_chat_response_choices_raises_for_unmatched_cho SimpleNamespace( index=1, token_ids=[200], - routed_experts=torch.tensor([[[10]], [[11]]], dtype=torch.int32), + routed_experts=torch.tensor( + [[[10]], [[11]]], dtype=ROUTED_EXPERTS_FALLBACK_DTYPE + ), ) ], ) @@ -602,3 +641,71 @@ def test_compute_spec_decode_metrics(): assert math.isclose(metrics["vllm/spec_acceptance_length"], 3.4, rel_tol=1e-6) # acceptance_rate = accepted / draft_tokens = 240 / 300 = 0.8 assert math.isclose(metrics["vllm/spec_acceptance_rate"], 0.8, rel_tol=1e-6) + + +def test_resolve_routed_experts_dtype_boundaries(): + assert resolve_routed_experts_dtype(None) == ROUTED_EXPERTS_FALLBACK_DTYPE + assert resolve_routed_experts_dtype(8) == torch.int8 + assert resolve_routed_experts_dtype(128) == torch.int8 + assert resolve_routed_experts_dtype(129) == torch.int16 + assert resolve_routed_experts_dtype(256) == torch.int16 + assert resolve_routed_experts_dtype(32768) == torch.int16 + assert resolve_routed_experts_dtype(32769) == torch.int32 + + +def test_get_num_routed_experts_across_config_conventions(): + qwen_moe = SimpleNamespace(num_experts=128) + deepseek = SimpleNamespace(n_routed_experts=256) + mixtral = SimpleNamespace(num_local_experts=8) + dense = SimpleNamespace() + vlm = SimpleNamespace(text_config=SimpleNamespace(num_experts=128)) + + assert get_num_routed_experts(qwen_moe) == 128 + assert get_num_routed_experts(deepseek) == 256 + assert get_num_routed_experts(mixtral) == 8 + assert get_num_routed_experts(dense) is None + assert get_num_routed_experts(vlm) == 128 + + +def test_pad_and_align_uses_resolved_dtype(): + class Output: + pass + + request_output = Output() + completion_output = Output() + completion_output.routed_experts = torch.arange(5 * 3 * 2).reshape(5, 3, 2) % 128 + + routed_experts = pad_and_align_routed_expert_indices( + request_output, + completion_output, + valid_length=6, + padded_length=8, + device=torch.device("cpu"), + routed_experts_dtype=torch.int8, + ) + + assert routed_experts.dtype == torch.int8 + assert torch.equal( + routed_experts[:5], completion_output.routed_experts.to(torch.int8) + ) + + +def test_pad_and_align_rejects_expert_ids_overflowing_dtype(monkeypatch): + monkeypatch.setattr(vllm_utils, "G_ROUTED_EXPERTS_RANGE_CHECKED", False) + + class Output: + pass + + request_output = Output() + completion_output = Output() + completion_output.routed_experts = torch.full((2, 1, 2), 200) + + with pytest.raises(ValueError, match="exceeds the resolved carry dtype"): + pad_and_align_routed_expert_indices( + request_output, + completion_output, + valid_length=3, + padded_length=3, + device=torch.device("cpu"), + routed_experts_dtype=torch.int8, + )