[None][feat] Default GPT-OSS to the Python KV-cache transceiver

tensorrt_llm/_torch/models/modeling_gpt_oss.py

@@ -1,4 +1,4 @@
-from typing import Dict, Optional
+from typing import Any, Dict, Literal, Optional
 
 import torch
 from torch import nn
@@ -552,6 +552,13 @@ def forward(
 @register_auto_model("GptOssForCausalLM")
 class GptOssForCausalLM(SpecDecOneEngineForCausalLM[Transformer, GptOssConfig]):
 
+    @classmethod
+    def get_preferred_transceiver_runtime(
+        cls,
+        pretrained_config: Any = None,
+    ) -> Optional[Literal["CPP", "PYTHON"]]:
+        return "PYTHON"
+
     params_map = {
         # TRTLLM module name : GptOss module name
         "qkv_proj": "qkv",

tests/unittest/_torch/modeling/test_modeling_gpt_oss.py

@@ -47,6 +47,10 @@
 """
 
 
+def test_gpt_oss_prefers_python_transceiver() -> None:
+    assert GptOssForCausalLM.get_preferred_transceiver_runtime() == "PYTHON"
+
+
 def dump_config_json(dst_dir):
     if os.path.exists(dst_dir):
         shutil.rmtree(dst_dir)