Revert "[Bugfix] Restrict TRTLLM attention to SM100, fixing GB300 (SM103) hang (#38730)"

This reverts commit c284a6671c.
This commit is contained in:
khluu
2026-04-01 15:14:58 -07:00
parent c284a6671c
commit cfad6a509c
3 changed files with 10 additions and 14 deletions

View File

@@ -235,11 +235,10 @@ def _resolve_import_to_file(
def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
"""Find a compute capability from is_device_capability*() calls in a function.
"""Find a compute capability from is_device_capability_family() calls in a function.
Handles two patterns:
- is_device_capability_family(N): "M.x" (e.g. 100 -> "10.x")
- is_device_capability(N): "M.m" (e.g. 100 -> "10.0")
Looks for the pattern: current_platform.is_device_capability_family(N)
and converts N (e.g. 100) to a CC string (e.g. "10.x").
"""
for node in ast.walk(tree):
if not isinstance(node, ast.FunctionDef) or node.name != func_name:
@@ -248,15 +247,12 @@ def _find_cc_in_function(tree: ast.AST, func_name: str) -> str | None:
if (
isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and n.func.attr == "is_device_capability_family"
and n.args
and isinstance(n.args[0], ast.Constant)
and isinstance(n.args[0].value, int)
):
val = n.args[0].value
if n.func.attr == "is_device_capability_family":
return f"{val // 10}.x"
elif n.func.attr == "is_device_capability":
return f"{val // 10}.{val % 10}"
return f"{n.args[0].value // 10}.x"
return None