Embed DeepGEMM source (not submodule) for SM100 raw CUDA GEMM primitives

third_party/DeepGEMM/scripts/generate_pyi.py

@@ -0,0 +1,890 @@
+import re
+from pathlib import Path
+
+
+def build_cpp_function_index(root_path):
+    func_index = {}
+    extensions = {'.cpp', '.cc', '.cxx', '.c', '.hpp', '.h'}
+
+    pattern = re.compile(
+        r'([\w:\s*<&>,\[\]\(\)]+?)'
+        r'\s+'
+        r'([a-zA-Z_][a-zA-Z0-9_:]*)'
+        r'\s*\(',
+    )
+
+    for file_path in Path(root_path).rglob('*'):
+        if file_path.suffix.lower() not in extensions:
+            continue
+        if not file_path.is_file():
+            continue
+
+        try:
+            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                content = f.read()
+        except Exception as e:
+            print(f'Failed to read file {file_path}: {e}')
+            continue
+
+        # Remove the compile directives and comments
+        lines = content.split('\n')
+        clean_lines = [line for line in lines if not line.strip().startswith(('#', '//'))]
+        content = '\n'.join(clean_lines)
+
+        for match in pattern.finditer(content):
+            return_type_part = match.group(1).strip()
+            full_func_name = match.group(2).strip()
+
+            if not return_type_part or not re.match(r'^[a-zA-Z_]', return_type_part):
+                continue
+
+            first_token = return_type_part.split()[0]
+            if first_token in {'return', 'if', 'else', 'for', 'while', 'switch', 'case', 'throw', 'catch', 'auto'}:
+                continue
+
+            # Extract base name
+            if '::' in full_func_name:
+                base_name = full_func_name.split('::')[-1]
+            else:
+                base_name = full_func_name
+
+            if not re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', base_name):
+                continue
+
+            # Find matching ')'
+            paren_start = match.end() - 1
+            paren_count = 0
+            pos = paren_start
+            while pos < len(content):
+                ch = content[pos]
+                if ch == '(':
+                    paren_count += 1
+                elif ch == ')':
+                    paren_count -= 1
+                    if paren_count == 0:
+                        break
+                    elif paren_count < 0:
+                        pos = -1
+                        break
+                pos += 1
+            else:
+                continue
+
+            if pos == -1:
+                continue
+
+            # Check context before match: should be at statement boundary
+            match_start = match.start()
+            context_before = content[max(0, match_start - 50):match_start]
+            if context_before and re.search(r'[a-zA-Z0-9_]$', context_before.rstrip()):
+                continue
+
+            # Check for definition or header declaration
+            is_header = file_path.suffix.lower() in {'.h', '.hpp', '.cuh'}
+            after_paren = content[pos+1:pos+500]
+            has_brace = '{' in after_paren
+            has_semicolon = ';' in after_paren.split('{')[0]
+
+            if has_brace or (is_header and has_semicolon):
+                sig_start = match.start(1)
+                full_signature = content[sig_start:pos+1].strip()
+                if base_name not in func_index:
+                    func_index[base_name] = full_signature
+
+    return func_index
+
+
+class BracketTracker:
+    """
+    Tracks nesting levels of various brackets in C++ code:
+      - () → paren
+      - [] → bracket
+      - {} → brace
+      - <> → angle (treated as template brackets only at top level)
+    Provides is_top_level() to check if currently outside all brackets.
+    """
+    def __init__(self):
+        self.paren = 0      # ()
+        self.bracket = 0    # []
+        self.brace = 0      # {}
+        self.angle = 0      # <>
+
+    def update(self, char: str):
+        """
+        Update internal counters based on the given character.
+        """
+        if char == '(':
+            self.paren += 1
+        elif char == ')':
+            self.paren -= 1
+        elif char == '[':
+            self.bracket += 1
+        elif char == ']':
+            self.bracket -= 1
+        elif char == '{':
+            self.brace += 1
+        elif char == '}':
+            self.brace -= 1
+        # Angle brackets < > are only treated as template delimiters
+        # when not inside (), [], or {}
+        elif char == '<' and self._in_top_level_of_other_brackets():
+            self.angle += 1
+        elif char == '>' and self.angle > 0 and self._in_top_level_of_other_brackets():
+            self.angle -= 1
+
+    def _in_top_level_of_other_brackets(self):
+        """
+        Check if not inside parentheses, square brackets, or braces (for correct template bracket recognition).
+        """
+        return self.paren == 0 and self.bracket == 0 and self.brace == 0
+
+    def is_top_level(self):
+        """
+        Check if completely at top level (all bracket counters are zero).
+        """
+        return (self.paren == 0 and
+                self.bracket == 0 and
+                self.brace == 0 and
+                self.angle == 0)
+
+
+def extract_m_def_statements(root_path):
+    """
+    Scan all c files under root_path and extract all m.def(...) statements.
+    """
+    results = []
+    extensions = {'.hpp', '.cpp', '.h', '.cc'}
+
+    # Regex: match m.def( ... ), supports multi-line
+    pattern = re.compile(r'm\.def\s*\(')
+
+    for file_path in Path(root_path).rglob('*'):
+        if file_path.suffix.lower() not in extensions:
+            continue
+        if not file_path.is_file():
+            continue
+
+        try:
+            with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
+                content = f.read()
+        except Exception as e:
+            print(f'Failed to read file {file_path}: {e}')
+            continue
+
+        m_def_list = []
+        lines = content.splitlines(keepends=True)
+        i = 0
+        while i < len(lines):
+            line = lines[i]
+            if 'm.def(' in line:
+                # Found a potential starting line
+                start_i = i
+                # Check if it's a comment
+                stripped = line.lstrip()
+                if stripped.startswith('//') or stripped.startswith('/*'):
+                    i += 1
+                    continue
+
+                # Try to match the complete m.def(...) call
+                paren_count = 0
+                j = i
+                found_start = False
+                while j < len(lines):
+                    current_line = lines[j]
+                    for k, char in enumerate(current_line):
+                        if char == '(':
+                            if not found_start and re.search(r'm\.def\s*\(', current_line[:k+1]):
+                                found_start = True
+                            if found_start:
+                                paren_count += 1
+                        elif char == ')':
+                            if found_start:
+                                paren_count -= 1
+                                if paren_count == 0:
+                                    # Found complete statement
+                                    full_stmt = ''.join(lines[i:j+1]).rstrip()
+                                    m_def_list.append(full_stmt)
+                                    i = j
+                                    break
+                    if paren_count <= 0 and found_start:
+                        break
+                    j += 1
+                else:
+                    pass
+            i += 1
+
+        if m_def_list:
+            results.append({
+                'file': str(file_path),
+                'm_def_statements': m_def_list
+            })
+
+    return results
+
+
+def parse_m_def_statement(m_def_str):
+    result = {
+        'python_function_name': None,
+        'num_args': 0,
+        'default_args': {},
+        'is_lambda': False,
+    }
+
+    # Extract top-level arguments
+    start = m_def_str.find('m.def(')
+    if start == -1:
+        raise ValueError(f'[{m_def_str}] Could not find m.def start position')
+
+    paren_count = 0
+    content_start = start + len('m.def(')
+    content_end = -1
+    for i in range(content_start, len(m_def_str)):
+        ch = m_def_str[i]
+        if ch == '(':
+            paren_count += 1
+        elif ch == ')':
+            if paren_count == 0:
+                content_end = i
+                break
+            else:
+                paren_count -= 1
+    if content_end == -1:
+        raise ValueError(f'[{m_def_str}] m.def parentheses not closed')
+
+    args_content = m_def_str[content_start:content_end]
+
+    # Split arguments using BracketTracker
+    args_list = []
+    current = []
+    tracker = BracketTracker()
+
+    for ch in args_content:
+        if ch in '()[]{}<>':
+            tracker.update(ch)
+        if ch == ',' and tracker.is_top_level():
+            args_list.append(''.join(current).strip())
+            current = []
+        else:
+            current.append(ch)
+
+    if current:
+        args_list.append(''.join(current).strip())
+
+    if len(args_list) < 2:
+        raise ValueError(f'[{m_def_str}] m.def has insufficient arguments')
+
+    # Extract Python function name
+    first = args_list[0].strip()
+    str_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"', first)
+    if str_match:
+        result['python_function_name'] = str_match.group(1)
+    else:
+        raise ValueError(f'[{m_def_str}] m.def first argument should be a string literal')
+
+    cpp_func_part = args_list[1].strip()
+    if cpp_func_part.startswith('&'):
+        cpp_func_part = cpp_func_part[1:].strip()
+
+    if cpp_func_part.startswith('['):
+        result['is_lambda'] = True
+        result['cpp_function_name'] = None
+    else:
+        if '::' in cpp_func_part:
+            cpp_func_name = cpp_func_part.split('::')[-1]
+        else:
+            cpp_func_name = cpp_func_part
+
+        match = re.match(r'^([a-zA-Z_][a-zA-Z0-9_]*)', cpp_func_name)
+        if match:
+            result['cpp_function_name'] = match.group(1)
+        else:
+            result['cpp_function_name'] = cpp_func_name
+
+    # Parse py::arg arguments
+    py_args = args_list[2:]
+    result['num_args'] = len(py_args)
+
+    for idx, arg_expr in enumerate(py_args):
+        expr = arg_expr.strip()
+        # Find top-level '='
+        eq_pos = -1
+        p_depth = b_depth = br_depth = angle_depth = 0
+        i = 0
+        while i < len(expr):
+            ch = expr[i]
+            if ch == '(':
+                p_depth += 1
+            elif ch == ')':
+                p_depth -= 1
+            elif ch == '[':
+                b_depth += 1
+            elif ch == ']':
+                b_depth -= 1
+            elif ch == '{':
+                br_depth += 1
+            elif ch == '}':
+                br_depth -= 1
+            elif ch == '<' and p_depth == 0 and b_depth == 0 and br_depth == 0:
+                angle_depth += 1
+            elif ch == '>' and angle_depth > 0 and p_depth == 0 and b_depth == 0 and br_depth == 0:
+                angle_depth -= 1
+            elif ch == '=' and all(d == 0 for d in [p_depth, b_depth, br_depth, angle_depth]):
+                eq_pos = i
+                break
+            i += 1
+
+        if eq_pos != -1:
+            default_val = expr[eq_pos + 1:].strip()
+            if not default_val:
+                raise ValueError(f'[{expr}] Default value is empty (arg {idx})')
+            result['default_args'][idx] = default_val
+
+    return result
+
+
+def extract_cpp_signature_from_content(cpp_func_name, content):
+    """
+    Search for the C++ function signature of cpp_func_name in the given file content.
+    """
+    if not cpp_func_name:
+        return None
+
+    # Build regex: match function starting with cpp_func_name (after word boundary)
+    # Note: function name may be preceded by return type (with templates, namespaces, etc.), followed by '('
+    pattern = re.compile(
+        r'^\s*'                                        # leading whitespace
+        r'([\w:\s*<&>,\[\]\(\)]+?)'                    # return type (non-greedy, allows templates, pointers, etc.)
+        r'\s+'                                         # at least one space
+        r'\b' + re.escape(cpp_func_name) + r'\b'       # function name (word boundary)
+                                           r'\s*\(',   # optional whitespace + start of param list
+        re.MULTILINE
+    )
+
+    for match in pattern.finditer(content):
+        # Find '(' position after function name
+        paren_start = match.end() - 1
+        if content[paren_start] != '(':
+            paren_start = content.find('(', match.end(0) - 1)
+            if paren_start == -1:
+                continue
+
+        # From '(', match to corresponding ')'
+        paren_count = 0
+        pos = paren_start
+        while pos < len(content):
+            ch = content[pos]
+            if ch == '(':
+                paren_count += 1
+            elif ch == ')':
+                paren_count -= 1
+                if paren_count == 0:
+                    start_sig = match.start(1)
+                    full_signature = content[start_sig:pos+1].strip()
+                    return full_signature
+            pos += 1
+
+    return None
+
+
+def parse_mdef_and_attach_cpp_signatures(item, func_index):
+    """
+    Enhance item by parsing m.def and extracting C++ function signature from global index
+    """
+    statements_with_parsed_signatures = []
+
+    for stmt in item['m_def_statements']:
+        parsed = parse_m_def_statement(stmt,)
+        cpp_func_name = parsed.get('cpp_function_name')
+
+        cpp_sig = None
+        if cpp_func_name and cpp_func_name in func_index:
+            cpp_sig = func_index[cpp_func_name]
+        else:
+            if not parsed['is_lambda']:
+                print(f'Warning: C++ function "{cpp_func_name}" not found in any .cpp file')
+
+        parsed['cpp_signature'] = cpp_sig
+        statements_with_parsed_signatures.append({
+            'raw': stmt,
+            'parsed': parsed
+        })
+
+    return {
+        'm_def_statements': statements_with_parsed_signatures
+    }
+
+
+def parse_cpp_signature(cpp_sig):
+    """
+    Parse a C++ function signature and extract return type, parameter types, and names.
+    """
+    if not cpp_sig or not cpp_sig.strip():
+        return None
+
+    # Find function name: last identifier before '('
+    paren_pos = cpp_sig.find('(')
+    if paren_pos == -1:
+        return None
+
+    before_paren = cpp_sig[:paren_pos].strip()
+    if not before_paren:
+        return None
+
+    # Function name is the last word in before_paren (may include templates like func<int>)
+    tokens = before_paren.split()
+    if len(tokens) < 2:
+        return None
+
+    # Heuristic: function name is usually the last token (may include <>)
+    func_name_part = tokens[-1]
+    return_type = ' '.join(tokens[:-1]).strip()
+
+    # Now extract parameter list content
+    param_list_str = cpp_sig[paren_pos+1:cpp_sig.rfind(')')].strip()
+    parameters = []
+
+    if param_list_str and param_list_str != 'void':  # 'void' means no parameters
+        # Split parameters (handle commas not inside templates/brackets)
+        param_decls = split_cpp_parameters(param_list_str)
+        for decl in param_decls:
+            decl = decl.strip()
+            if not decl:
+                continue
+            # Try to split type and name from right to left
+            param_info = parse_parameter_declaration(decl)
+            if param_info:
+                parameters.append(param_info)
+
+    return {
+        'return_type': return_type,
+        'parameters': parameters,
+        'num_parameters': len(parameters)
+    }
+
+
+def split_cpp_parameters(param_str: str):
+    """
+    Split a C++ parameter list string by top-level commas,
+    e.g., 'int a, std::vector<float> b' → ['int a', 'std::vector<float> b']
+    """
+    if not param_str.strip() or param_str == 'void':
+        return []
+    params = []
+    current = []
+    tracker = BracketTracker()
+
+    for ch in param_str:
+        if ch in '()[]{}<>':
+            tracker.update(ch)
+        if ch == ',' and tracker.is_top_level():
+            param = ''.join(current).strip()
+            if param:  # Only add non-empty parameters
+                params.append(param)
+            current = []
+        else:
+            current.append(ch)
+
+    if current:
+        final_param = ''.join(current).strip()
+        if final_param:  # Only add non-empty parameters
+            params.append(final_param)
+    return params
+
+
+def parse_parameter_declaration(decl: str):
+    """
+    Parse a single parameter declaration, e.g., 'const std::string& name' → {'type': 'const std::string&', 'name': 'name'}
+    Improved version that better handles template types.
+    """
+    decl = decl.strip()
+    if not decl:
+        return None
+
+    # Remove possible default value (starting from top-level '=')
+    tracker = BracketTracker()
+    eq_pos = -1
+    for i, ch in enumerate(decl):
+        if ch in '()[]{}<>':
+            tracker.update(ch)
+        elif ch == '=' and tracker.is_top_level():
+            eq_pos = i
+            break
+
+    if eq_pos != -1:
+        decl = decl[:eq_pos].strip()
+
+    # Now decl is 'type name' or just 'type'
+    # Instead of simple splitting, we'll use a more robust approach
+    # to find the parameter name
+
+    # First, let's handle the case where there's no explicit parameter name
+    # (this sometimes happens in function declarations)
+    if not re.search(r'[a-zA-Z_][a-zA-Z0-9_]*$', decl):
+        # No parameter name found, just return the type
+        return {
+            'type': decl,
+            'name': None
+        }
+
+    # Use bracket tracking to find where the type ends and name begins
+    tracker = BracketTracker()
+    name_start = -1
+
+    # Scan from the end to find the start of the parameter name
+    # We look for the first identifier that's outside all brackets
+    i = len(decl) - 1
+    while i >= 0:
+        ch = decl[i]
+
+        if ch in '()[]{}<>':
+            tracker.update(ch)
+
+        # If we're at top level and find an identifier character
+        if tracker.is_top_level() and re.match(r'[a-zA-Z0-9_]', ch):
+            # Track back to find the start of this identifier
+            name_start = i
+            while name_start > 0 and re.match(r'[a-zA-Z0-9_]', decl[name_start - 1]):
+                name_start -= 1
+
+            # Check if this might be part of a type keyword (like 'int', 'bool', etc.)
+            potential_name = decl[name_start:i+1]
+            type_keywords = {'int', 'long', 'short', 'char', 'bool', 'float', 'double',
+                             'void', 'auto', 'const', 'static', 'volatile', 'mutable',
+                             'unsigned', 'signed'}
+
+            # If it's not a type keyword and looks like a parameter name, use it
+            if (potential_name not in type_keywords and
+                    re.match(r'^[a-zA-Z_][a-zA-Z0-9_]*$', potential_name)):
+                break
+
+        i -= 1
+
+    if name_start != -1 and i >= 0:
+        param_name = decl[name_start:i+1]
+        param_type = decl[:name_start].strip()
+
+        # Clean up the type - remove trailing &, * and whitespace
+        param_type = param_type.rstrip('&* \t')
+
+        return {
+            'type': param_type,
+            'name': param_name
+        }
+
+    # Fallback: if we can't find a clear parameter name, just return the type
+    return {
+        'type': decl,
+        'name': None
+    }
+
+
+def extract_cpp_signature_details(item):
+    """
+    For each m.def entry in item, parse cpp_signature to extract return type and parameter details.
+    """
+    statements_with_parsed_signatures = []
+    for stmt_info in item['m_def_statements']:
+        parsed = stmt_info['parsed']
+        cpp_sig = parsed.get('cpp_signature')
+
+        cpp_params_info = None
+        if cpp_sig:
+            try:
+                cpp_params_info = parse_cpp_signature(cpp_sig)
+            except Exception as e:
+                print(f'Failed to parse C++ signature: {e}')
+
+        parsed['cpp_parsed_signature'] = cpp_params_info
+        statements_with_parsed_signatures.append({
+            'raw': stmt_info['raw'],
+            'parsed': parsed
+        })
+
+    return {
+        'm_def_statements': statements_with_parsed_signatures
+    }
+
+
+def cpp_type_to_python_type(cpp_type: str) -> str:
+    if not cpp_type:
+        return 'Any'
+
+    original = cpp_type.strip()
+    if not original:
+        return 'Any'
+
+    # Remove C++ specifiers that don't affect Python type
+    cleaned = re.sub(r'\b(static|inline|constexpr|thread_local|extern|mutable|const|volatile|endif)\b', '', original)
+    cleaned = cleaned.replace('&', '').replace('*', '').strip()
+    cleaned = re.sub(r'\s+', ' ', cleaned).strip()
+
+    # Handle void
+    if cleaned == 'void':
+        return 'None'
+
+    # Handle template types — ORDER MATTERS! Must come before internal type checks.
+
+    # std::pair<T1, T2>
+    if cleaned.startswith('std::pair<'):
+        inner = cleaned[10:-1].strip()  # len('std::pair<') == 10
+        args = split_template_args(inner)
+        if len(args) == 2:
+            t1 = cpp_type_to_python_type(args[0])
+            t2 = cpp_type_to_python_type(args[1])
+            return f'tuple[{t1}, {t2}]'
+        else:
+            print(f'Warning: std::pair with unexpected number of args: {cleaned}')
+            return 'Any'
+
+    # std::tuple<T1, T2, ...>
+    if cleaned.startswith('std::tuple<'):
+        inner = cleaned[11:-1].strip()  # len('std::tuple<') == 11
+        args = split_template_args(inner)
+        py_types = [cpp_type_to_python_type(arg) for arg in args]
+        return f"tuple[{', '.join(py_types)}]"
+
+    # std::vector<T>
+    if cleaned.startswith('std::vector<'):
+        inner = cleaned[12:-1].strip()  # len('std::vector<') == 12
+        args = split_template_args(inner)
+        if len(args) == 1:
+            inner_py = cpp_type_to_python_type(args[0])
+            return f'list[{inner_py}]'
+        else:
+            print(f'Warning: std::vector with unexpected args: {cleaned}')
+            return 'Any'
+
+    # std::optional<T>
+    if cleaned.startswith('std::optional<'):
+        inner = cleaned[14:-1].strip()  # len('std::optional<') == 14
+        args = split_template_args(inner)
+        if len(args) == 1:
+            inner_py = cpp_type_to_python_type(args[0])
+            return f'Optional[{inner_py}]'
+        else:
+            print(f'Warning: std::optional with unexpected args: {cleaned}')
+            return 'Any'
+
+    # std::string
+    if re.search(r'\bstd::string\b', original):
+        return 'str'
+
+    # C-style strings: char*, const char*, char[], etc.
+    if re.search(r'\b(?:const\s+)?char\s*[\*\[]', original):
+        return 'str'
+
+    # Boolean
+    if re.search(r'\bbool\b', cleaned):
+        return 'bool'
+
+    # Integer types (including fixed-width and common aliases)
+    if re.search(r'\b(int|long|short|size_t|ssize_t|ptrdiff_t|'
+                 r'int8_t|int16_t|int32_t|int64_t|'
+                 r'uint8_t|uint16_t|uint32_t|uint64_t)\b', cleaned):
+        return 'int'
+
+    # Floating-point
+    if re.search(r'\b(float|double|long\s+double)\b', cleaned):
+        return 'float'
+
+    # torch::Tensor
+    if re.search(r'\btorch::Tensor\b', original):
+        return 'torch.Tensor'
+
+    # Unrecognized type
+    print(f'Warning: Unrecognized C++ type: {original}')
+    return 'Any'
+
+
+def split_template_args(template_args: str):
+    """
+    Split template arguments, e.g., 'int, std::vector<float>' → ['int', 'std::vector<float>']
+    """
+    if not template_args.strip():
+        return []
+    args = []
+    current = []
+    tracker = BracketTracker()
+
+    for ch in template_args:
+        if ch in '()[]{}<>':
+            tracker.update(ch)
+        if ch == ',' and tracker.is_top_level():
+            args.append(''.join(current).strip())
+            current = []
+        else:
+            current.append(ch)
+
+    if current:
+        args.append(''.join(current).strip())
+    return args
+
+
+def cpp_default_to_python_default(cpp_default: str):
+    """
+    Convert C++ default value string to valid Python expression string.
+    """
+    if not cpp_default:
+        return 'None'
+
+    s = cpp_default.strip()
+
+    # Handle string literals: 'bf16' → 'bf16'
+    # Match: starts and ends with unescaped double quotes
+    string_match = re.match(r'^"([^"\\]*(?:\\.[^"\\]*)*)"$', s)
+    if string_match:
+        return s
+
+    # Handle boolean literals
+    if s == 'false':
+        return 'False'
+    if s == 'true':
+        return 'True'
+
+    # Handle null-like values: nullptr, nullopt, NULL, etc.
+    if s in ('nullptr', 'NULL') or 'nullopt' in s:
+        return 'None'
+
+    # Handle std::tuple<int, int>({128, 128}) → (128, 128)
+    tuple_match = re.match(r'std::tuple\s*<[^>]*>\s*\(\s*({.*?})\s*\)', s)
+    if tuple_match:
+        inner = tuple_match.group(1)  # {128, 128}
+        inner_py = inner.replace('{', '(').replace('}', ')')
+        return inner_py
+
+    # Handle std::make_tuple(1, 2, 3) → (1, 2, 3)
+    make_tuple_match = re.match(r'std::make_tuple\s*\(\s*(.*?)\s*\)', s)
+    if make_tuple_match:
+        inner = make_tuple_match.group(1)
+        # Ensure it's a valid tuple even with one element: add comma if needed?
+        # But in C++ default args, it's usually multi-element, so we assume valid.
+        return f'({inner})'
+
+    # Handle std::vector<int>({1,2,3}) → [1, 2, 3]
+    vector_match = re.match(r'std::vector\s*<[^>]*>\s*\(\s*({.*?})\s*\)', s)
+    if vector_match:
+        inner = vector_match.group(1)
+        inner_py = inner.replace('{', '[').replace('}', ']')
+        return inner_py
+
+    # Handle numeric literals: integers and floats
+    if re.match(r'^[+-]?\d+$', s):  # integer
+        return s
+    if re.match(r'^[+-]?\d*\.\d+([eE][+-]?\d+)?$', s):  # float
+        return s
+
+    # Fallback: unrecognized → warn and return None
+    print(f'Warning: Unrecognized default value: {s}')
+    return 'None'
+
+
+def generate_pyi_function(item_entry):
+    parsed = item_entry['parsed']
+    py_name = parsed['python_function_name']
+
+    if parsed.get('is_lambda'):
+        return f'def {py_name}(*args, **kwargs) -> Any: ...'
+
+    sig_info = parsed.get('cpp_parsed_signature')
+    default_args = parsed.get('default_args', {})
+
+    if not sig_info:
+        return f'def {py_name}(*args, **kwargs) -> Any: ...'
+
+    return_type = cpp_type_to_python_type(sig_info['return_type'])
+    params = sig_info['parameters']
+    num_params = len(params)
+
+    # Build parameter list
+    param_lines = []
+    for i in range(num_params):
+        param_info = params[i] if i < len(params) else {'type': 'Any', 'name': f'arg{i}'}
+        param_type = cpp_type_to_python_type(param_info['type'])
+        param_name = param_info['name'] or f'arg{i}'
+
+        # Replace invalid Python identifiers (e.g., keywords)
+        if param_name in {'def', 'class', 'from', 'import', 'None', 'True', 'False'}:
+            param_name = f'{param_name}_'
+
+        # Check for default value
+        if i in default_args:
+            cpp_default = default_args[i]
+            py_default = cpp_default_to_python_default(cpp_default)
+            param_str = f'    {param_name}: {param_type} = {py_default}'
+        else:
+            param_str = f'    {param_name}: {param_type}'
+
+        param_lines.append(param_str)
+
+    if param_lines:
+        params_block = ',\n'.join(param_lines)
+        func_def = f'def {py_name}(\n{params_block}\n) -> {return_type}: ...'
+    else:
+        func_def = f'def {py_name}() -> {return_type}: ...'
+
+    return func_def
+
+
+def generate_pyi_file_content(enhanced_results, module_name: str = 'my_module'):
+    function_decls = []
+    has_optional = False
+    has_torch = False
+    has_numpy = False
+
+    for item in enhanced_results:
+        for stmt in item['m_def_statements']:
+            try:
+                decl = generate_pyi_function(stmt)
+                function_decls.append(decl)
+
+                if 'Optional[' in decl:
+                    has_optional = True
+                if 'torch.Tensor' in decl:
+                    has_torch = True
+                if 'numpy.ndarray' in decl or 'py::array' in str(stmt):
+                    has_numpy = True
+            except Exception as e:
+                func_name = stmt['parsed'].get('python_function_name', 'unknown')
+                function_decls.append(f'# ERROR: failed to generate stub for {func_name}: {e}')
+
+    imports = ['from typing import Any']
+    if has_optional:
+        imports[0] += ', Optional'
+
+    if has_torch:
+        imports.append('import torch')
+    if has_numpy:
+        imports.append('import numpy')
+
+    lines = [f'# Stubs for module: {module_name}', '']
+    lines.extend(imports)
+    lines.append('')
+    lines.append('')
+
+    for decl in function_decls:
+        lines.append(decl)
+        lines.append('')
+        lines.append('')
+
+    return '\n'.join(lines)
+
+
+def generate_pyi_file(name, root, output_dir='.'):
+    func_index = build_cpp_function_index(root)
+    results = extract_m_def_statements(root)
+
+    cpp_results = []
+    for item in results:
+        enhanced_item = parse_mdef_and_attach_cpp_signatures(item, func_index)
+        cpp_item = extract_cpp_signature_details(enhanced_item)
+        cpp_results.append(cpp_item)
+
+    pyi_content = generate_pyi_file_content(cpp_results, module_name=name)
+
+    output_path = Path(output_dir) / f'{name}.pyi'
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    with open(output_path, 'w', encoding='utf-8') as f:
+        f.write(pyi_content)
+
+    print(f'.pyi file generated: {output_path}')

third_party/DeepGEMM/scripts/quick_plot_pm.py

@@ -0,0 +1,448 @@
+#!/usr/bin/env python3
+"""Plot a curated set of NCU PM metrics from an .ncu-rep report.
+
+Usage:
+    python scripts/quick_plot_pm.py [report.ncu-rep]
+
+By default the script saves a PNG next to the report.
+With --interactive, it opens a Qt window instead.
+"""
+
+import argparse
+import csv
+import io
+import subprocess
+from dataclasses import dataclass
+
+import matplotlib
+import numpy as np
+
+
+@dataclass(frozen=True)
+class MetricSpec:
+    name: str
+    metric: str
+    kind: str
+    category: str
+    aliases: tuple[str, ...] = ()
+
+
+@dataclass(frozen=True)
+class ResolvedMetricSpec:
+    name: str
+    metric: str
+    kind: str
+    category: str
+
+
+@dataclass(frozen=True)
+class MetricSeries:
+    name: str
+    metric: str
+    category: str
+    unit: str
+    values: tuple[float, ...]
+
+
+CATEGORY_ORDER = [
+    "Overview",
+    "SM",
+    "L1",
+    "L2",
+    "DRAM",
+    "Interconnect",
+]
+
+
+KIND_SUFFIXES = {
+    "pct_peak": [".avg.pct_of_peak_sustained_elapsed"],
+    "pct": [".pct", ".avg.pct_of_peak_sustained_elapsed"],
+    "avg": [".avg"],
+    "sum": [".sum"],
+    "avg_per_second": [".avg.per_second"],
+    "sum_per_second": [".sum.per_second"],
+    "avg_per_cycle_active": [".avg.per_cycle_active"],
+    "avg_per_cycle_elapsed": [".avg.per_cycle_elapsed"],
+    "sum_per_cycle_elapsed": [".sum.per_cycle_elapsed"],
+}
+
+
+# Curated from scripts/ncu-metrics.txt, with a few corrections against
+# `ncu --query-metrics --chip gb100`:
+# - Blocks launched uses `gr__ctas_launched_realtime`
+# - SM active cycles uses `sm__cycles_active`
+# - L2 throughput for GCC requests uses `lts__t_sector_throughput_srcunit_gcc`
+# - C2C throughput uses `ctc__throughput`
+# - NVLink RX metrics use the `NVLRX` domain
+CURATED_METRICS = [
+    MetricSpec("Blocks Launched", "FE_B.TriageCompute.gr__ctas_launched_realtime", "sum_per_cycle_elapsed", "Overview"),
+    MetricSpec("Average Blocks Active", "TPC.TriageCompute.tpc__ctas_active_realtime", "avg_per_cycle_elapsed", "Overview"),
+    MetricSpec("Total Blocks Active", "TPC.TriageCompute.tpc__ctas_active_realtime", "sum_per_cycle_elapsed", "Overview"),
+    MetricSpec("Average CGAs Active", "GPC_B.TriageCompute.gpc__cgas_active_realtime", "avg_per_cycle_elapsed", "Overview"),
+    MetricSpec("Total CGAs Active", "GPC_B.TriageCompute.gpc__cgas_active_realtime", "sum_per_cycle_elapsed", "Overview"),
+    MetricSpec("SM Active Cycles", "TPC.TriageCompute.sm__cycles_active", "avg", "SM"),
+    MetricSpec("Executed IPC Active", "TPC.TriageCompute.sm__inst_executed_realtime", "avg_per_cycle_active", "SM"),
+    MetricSpec("Executed IPC Elapsed", "TPC.TriageCompute.sm__inst_executed_realtime", "avg_per_cycle_elapsed", "SM"),
+    MetricSpec("SM Throughput", "TPC.TriageCompute.sm__inst_executed_realtime", "pct_peak", "SM"),
+    MetricSpec("SM ALU Pipe Throughput", "TPC.TriageCompute.sm__inst_executed_pipe_alu_realtime", "pct_peak", "SM"),
+    MetricSpec("SM FMA Pipe Throughput", "TPC.TriageCompute.sm__pipe_fma_cycles_active_realtime", "pct_peak", "SM"),
+    MetricSpec("SM FMA Heavy Pipe Throughput", "TPC.TriageCompute.sm__pipe_fmaheavy_cycles_active_realtime", "pct_peak", "SM"),
+    MetricSpec("SM FMA Light Pipe Throughput", "TPC.TriageCompute.sm__pipe_fmalite_cycles_active_realtime", "pct_peak", "SM"),
+    MetricSpec("SM Tensor Pipe Throughput", "TPC.TriageCompute.sm__pipe_tensor_cycles_active_realtime", "pct_peak", "SM"),
+    MetricSpec("SM TMEM Pipe Throughput", "SM_A.TriageCompute.sm__mem_tensor_cycles_active_realtime", "pct_peak", "SM"),
+    MetricSpec("SM Uniform Pipe Throughput", "SM_A.TriageCompute.sm__inst_executed_pipe_uniform_realtime", "pct_peak", "SM"),
+    MetricSpec("SM XU Pipe Throughput", "SM_A.TriageCompute.sm__inst_executed_pipe_xu_realtime", "pct_peak", "SM"),
+    MetricSpec("L1 Throughput", "SM_A.TriageCompute.l1tex__throughput", "pct_peak", "L1"),
+    MetricSpec("L1 Sectors", "SM_B.TriageCompute.l1tex__t_sectors", "sum", "L1"),
+    MetricSpec("L1 Hit Rate", "SM_B.TriageCompute.l1tex__t_sector_hit_rate", "pct", "L1"),
+    MetricSpec("L1 Lookup Hit", "SM_B.TriageCompute.l1tex__t_sectors_lookup_hit", "sum", "L1"),
+    MetricSpec("L1 Lookup Miss", "SM_B.TriageCompute.l1tex__t_sectors_lookup_miss", "sum", "L1"),
+    MetricSpec("L1 Wavefronts (Data)", "SM_A.TriageCompute.l1tex__data_pipe_lsu_wavefronts", "avg", "L1"),
+    MetricSpec("L1 Wavefronts (LGDS)", "SM_A.TriageCompute.l1tex__data_pipe_lsu_wavefronts_mem_lgds", "avg", "L1"),
+    MetricSpec("L1 Wavefronts (Shared)", "SM_A.TriageCompute.l1tex__data_pipe_lsu_wavefronts_mem_shared", "avg", "L1"),
+    MetricSpec("L2 Throughput", "LTS.TriageCompute.lts__throughput", "pct_peak", "L2"),
+    MetricSpec("L2 Throughput for L1 Requests", "LTS.TriageCompute.lts__t_sector_throughput_srcunit_tex", "pct_peak", "L2"),
+    MetricSpec("L2 Throughput for GCC Requests", "LTS.TriageCompute.lts__t_sector_throughput_srcunit_gcc", "pct_peak", "L2"),
+    MetricSpec("L2 Throughput to DRAM", "LTS.TriageCompute.lts__t_sector_throughput_srcnode_fbp", "pct_peak", "L2"),
+    MetricSpec("SysL2 Throughput to Peer Memory", "SYSLTS.TriageCompute.syslts__t_sector_throughput_aperture_peer", "pct_peak", "L2"),
+    MetricSpec("SysL2 Throughput to System Memory", "SYSLTS.TriageCompute.syslts__t_sector_throughput_aperture_sysmem", "pct_peak", "L2"),
+    MetricSpec("L2 Hit Rate", "LTS.TriageCompute.lts__average_t_sector_hit_rate_realtime", "pct", "L2"),
+    MetricSpec("L2 Hit Rate From L1", "LTS.TriageCompute.lts__average_t_sector_hit_rate_srcunit_tex_realtime", "pct", "L2"),
+    MetricSpec("DRAM Frequency", "FBSP.TriageCompute.dram__cycles_elapsed", "avg_per_second", "DRAM"),
+    MetricSpec("DRAM Throughput", "FBSP.TriageCompute.dram__throughput", "pct_peak", "DRAM"),
+    MetricSpec("DRAM Read Throughput", "FBSP.TriageCompute.dram__read_throughput", "pct_peak", "DRAM"),
+    MetricSpec("DRAM Write Throughput", "FBSP.TriageCompute.dram__write_throughput", "pct_peak", "DRAM"),
+    MetricSpec("C2C Throughput", "TriageCompute.ctc__throughput", "pct_peak", "Interconnect", aliases=("TriageCompute.ctx__throughput",)),
+    MetricSpec("NVLink Transmitted Throughput", "NVLTX.TriageCompute.nvltx__bytes", "pct_peak", "Interconnect"),
+    MetricSpec("NVLink Received Throughput", "NVLRX.TriageCompute.nvlrx__bytes", "pct_peak", "Interconnect"),
+    MetricSpec("NVLink Transmitted Bandwidth", "NVLTX.TriageCompute.nvltx__bytes", "sum_per_second", "Interconnect"),
+    MetricSpec("NVLink Received Bandwidth", "NVLRX.TriageCompute.nvlrx__bytes", "sum_per_second", "Interconnect"),
+    MetricSpec("PCIe Throughput", "PCI.TriageCompute.pcie__throughput", "pct_peak", "Interconnect"),
+    MetricSpec("PCIe Read Bandwidth", "PCI.TriageCompute.pcie__read_bytes", "sum_per_second", "Interconnect"),
+    MetricSpec("PCIe Write Bandwidth", "PCI.TriageCompute.pcie__write_bytes", "sum_per_second", "Interconnect"),
+]
+
+
+def _run_csv_command(command, timeout):
+    result = subprocess.run(command, capture_output=True, text=True, timeout=timeout)
+    if result.returncode != 0 and not result.stdout:
+        return None
+    reader = csv.reader(io.StringIO(result.stdout))
+    return list(reader)
+
+
+def _query_available_metrics(chip):
+    result = subprocess.run(
+        ["ncu", "--query-metrics", "--chip", chip],
+        capture_output=True,
+        text=True,
+        timeout=30,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(result.stderr.strip() or f"failed to query metrics for chip {chip}")
+
+    metrics = set()
+    for line in result.stdout.splitlines():
+        parts = line.split()
+        if not parts:
+            continue
+        token = parts[0]
+        if "__" not in token:
+            continue
+        metrics.add(token)
+    return metrics
+
+
+def _metric_candidates(metric):
+    candidates = [metric]
+    marker = ".TriageCompute."
+    if marker in metric:
+        candidates.append(metric.split(marker, 1)[1])
+    return candidates
+
+
+def resolve_metric_specs(chip):
+    available = _query_available_metrics(chip)
+    resolved = []
+    missing = []
+
+    for spec in CURATED_METRICS:
+        candidates = []
+        for metric in (spec.metric, *spec.aliases):
+            candidates.extend(_metric_candidates(metric))
+
+        actual_metric = next((metric for metric in candidates if metric in available), None)
+        if actual_metric is None:
+            missing.append(spec)
+            continue
+
+        resolved.append(ResolvedMetricSpec(spec.name, actual_metric, spec.kind, spec.category))
+
+    return resolved, missing
+
+
+def _parse_metric_values(raw):
+    if not raw or raw == "no data":
+        return ()
+
+    try:
+        if raw.startswith("(") and raw.endswith(")"):
+            rest = raw[1:-1]
+            return tuple(float(v.strip().replace(",", "")) for v in rest.split(";") if v.strip())
+        if " (" in raw:
+            _agg, rest = raw.split(" (", 1)
+            rest = rest.rstrip(")")
+            return tuple(float(v.strip().replace(",", "")) for v in rest.split(";") if v.strip())
+        return (float(raw.replace(",", "")),)
+    except ValueError:
+        return ()
+
+
+def _probe_metric_series(report, metric_name):
+    rows = _run_csv_command(
+        [
+            "ncu",
+            "--import",
+            report,
+            "--page",
+            "raw",
+            "--csv",
+            "--metrics",
+            metric_name,
+            "--print-metric-instances",
+            "values",
+        ],
+        timeout=60,
+    )
+    if not rows or len(rows) < 3 or len(rows[0]) <= 11:
+        return None
+
+    header, units, row = rows[0], rows[1], rows[2]
+    unit = units[11] if len(units) > 11 else ""
+    raw = row[11] if len(row) > 11 else ""
+    values = _parse_metric_values(raw)
+    return header[11], unit, values
+
+
+def collect_metric_series(report, resolved_specs):
+    collected = []
+    skipped = []
+
+    for spec in resolved_specs:
+        series = None
+        for suffix in KIND_SUFFIXES[spec.kind]:
+            probe = _probe_metric_series(report, f"{spec.metric}{suffix}")
+            if probe is None:
+                continue
+            full_metric, unit, values = probe
+            if len(values) > 1:
+                series = MetricSeries(spec.name, full_metric, spec.category, unit, values)
+                break
+
+        if series is None:
+            skipped.append(spec)
+            continue
+
+        collected.append(series)
+
+    return collected, skipped
+
+
+def _format_value(value):
+    if value == 0:
+        return "0"
+    abs_value = abs(value)
+    if abs_value >= 1e12:
+        return f"{value / 1e12:.2f} T"
+    if abs_value >= 1e9:
+        return f"{value / 1e9:.2f} G"
+    if abs_value >= 1e6:
+        return f"{value / 1e6:.2f} M"
+    if abs_value >= 1e3:
+        return f"{value / 1e3:.2f} K"
+    if abs_value >= 1:
+        return f"{value:.1f}"
+    return f"{value:.2f}"
+
+
+def _format_with_unit(value, unit):
+    if not unit:
+        return _format_value(value)
+    return f"{_format_value(value)} {unit}"
+
+
+def plot_pm(report, metrics, save=False):
+    """Plot curated PM metrics as shared-x subplots in a light theme."""
+    import matplotlib.pyplot as plt
+    from matplotlib.gridspec import GridSpec
+
+    if not metrics:
+        print("No curated metrics had time-series data in the report.")
+        return
+
+    bg_fig = "#ffffff"
+    bg_row = "#f6f8fb"
+    text_primary = "#1f2937"
+    text_secondary = "#6b7280"
+    text_header = "#111827"
+    grid_color = "#d7deea"
+    border = "#c7d0dd"
+
+    wave_colors = {
+        "Overview": "#7c8aa5",
+        "SM": "#4f87c2",
+        "L1": "#2f9d8f",
+        "L2": "#dd8452",
+        "DRAM": "#c95d63",
+        "Interconnect": "#8c6bb1",
+    }
+
+    category_rank = {category: index for index, category in enumerate(CATEGORY_ORDER)}
+    metrics = sorted(metrics, key=lambda item: (category_rank.get(item.category, 99), item.name))
+
+    row_h = 0.55
+    label_w = 3.6
+    plot_w = 14.0
+    fig_w = label_w + plot_w
+    fig_h = row_h * len(metrics) + 0.6
+
+    fig = plt.figure(figsize=(fig_w, fig_h), facecolor=bg_fig)
+    gs = GridSpec(
+        len(metrics),
+        1,
+        figure=fig,
+        left=label_w / fig_w,
+        right=0.97,
+        top=1 - 0.45 / fig_h,
+        bottom=0.35 / fig_h,
+        hspace=0.18,
+    )
+    axes = [fig.add_subplot(gs[i, 0]) for i in range(len(metrics))]
+
+    prev_category = None
+    for idx, metric in enumerate(metrics):
+        ax = axes[idx]
+        values = np.array(metric.values)
+        x = np.arange(len(values))
+        wave_color = wave_colors.get(metric.category, "#5b9bd5")
+
+        ax.set_facecolor(bg_row)
+        ax.fill_between(x, values, alpha=0.35, color=wave_color, linewidth=0)
+        ax.plot(x, values, linewidth=0.8, color=wave_color)
+
+        ax.set_xlim(0, len(values) - 1)
+        if metric.unit == "%":
+            ax.set_ylim(0, 100)
+        else:
+            ymax = np.max(values)
+            ax.set_ylim(0, ymax * 1.15 if ymax > 0 else 1)
+
+        ax.grid(True, axis="both", color=grid_color, linewidth=0.5, alpha=0.85)
+        ax.tick_params(axis="both", colors=text_secondary, labelsize=6, length=0)
+
+        if idx < len(metrics) - 1:
+            ax.tick_params(axis="x", labelbottom=False)
+        else:
+            ax.set_xlabel("Sample Index", fontsize=8, color=text_secondary)
+
+        ymin_v, ymax_v = ax.get_ylim()
+        ax.set_yticks([ymin_v, ymax_v])
+        ax.set_yticklabels([_format_value(ymin_v), _format_value(ymax_v)], fontsize=6, color=text_secondary)
+
+        peak = np.max(values)
+        ax.text(
+            1.005,
+            0.5,
+            _format_with_unit(peak, metric.unit),
+            transform=ax.transAxes,
+            fontsize=7,
+            color=text_secondary,
+            va="center",
+            ha="left",
+            family="monospace",
+        )
+
+        for spine in ax.spines.values():
+            spine.set_color(border)
+            spine.set_linewidth(0.5)
+
+        if metric.category != prev_category:
+            cat_y = ax.get_position().y1 + 0.008
+            fig.text(
+                0.005,
+                cat_y,
+                f"  {metric.category}",
+                fontsize=8.5,
+                fontweight="bold",
+                color=text_header,
+                va="bottom",
+                family="sans-serif",
+                transform=fig.transFigure,
+                bbox=dict(boxstyle="square,pad=0.15", facecolor="#e9eef5", edgecolor="none"),
+            )
+        prev_category = metric.category
+
+        label_y = (ax.get_position().y0 + ax.get_position().y1) / 2
+        fig.text(
+            label_w / fig_w - 0.012,
+            label_y,
+            metric.name,
+            fontsize=7.5,
+            color=text_primary,
+            va="center",
+            ha="right",
+            family="sans-serif",
+            transform=fig.transFigure,
+        )
+
+    fig.text(
+        0.5,
+        1 - 0.15 / fig_h,
+        f"PM Sampling - {report}",
+        fontsize=11,
+        fontweight="bold",
+        color=text_header,
+        ha="center",
+        va="top",
+        family="sans-serif",
+        transform=fig.transFigure,
+    )
+
+    if save:
+        out_path = report.replace(".ncu-rep", ".pm_sampling.png")
+        fig.savefig(out_path, dpi=150, facecolor=bg_fig, bbox_inches="tight", pad_inches=0.2)
+        print(f"Saved: {out_path}")
+        plt.close(fig)
+    else:
+        plt.show()
+
+
+def main():
+    parser = argparse.ArgumentParser(description="NCU PM Sampling plotter")
+    parser.add_argument("report", nargs="?", default="mega-moe-kk.3.ncu-rep", help="Path to .ncu-rep file")
+    parser.add_argument("--chip", default="gb100", help="Chip name used for `ncu --query-metrics`")
+    parser.add_argument("--interactive", action="store_true", help="Open an interactive Qt window instead of saving a PNG")
+    args = parser.parse_args()
+
+    if args.interactive:
+        matplotlib.use("QtAgg")
+    else:
+        matplotlib.use("Agg")
+
+    resolved_specs, missing_specs = resolve_metric_specs(args.chip)
+    if missing_specs:
+        print(f"Skipped {len(missing_specs)} curated metrics not available on {args.chip}.")
+        for spec in missing_specs:
+            print(f"  missing: {spec.name} -> {spec.metric}")
+
+    metric_series, skipped_specs = collect_metric_series(args.report, resolved_specs)
+    if skipped_specs:
+        print(f"Skipped {len(skipped_specs)} curated metrics with no time-series data in {args.report}.")
+        for spec in skipped_specs:
+            print(f"  no series: {spec.name} -> {spec.metric}")
+
+    plot_pm(args.report, metric_series, save=not args.interactive)
+
+
+if __name__ == "__main__":
+    main()

third_party/DeepGEMM/scripts/run_ncu_mega_moe.sh

@@ -0,0 +1,89 @@
+#!/bin/bash
+
+set -e
+
+# parse num-processes, output_dir and separate python args
+num_processes=8
+output_dir=work
+python_args=()
+for ((arg_idx = 1; arg_idx <= $#; ++arg_idx)); do
+    arg="${!arg_idx}"
+    case "$arg" in
+        --num-processes)
+            python_args+=("$arg")
+            if ((arg_idx < $#)); then
+                ((arg_idx++))
+                num_processes="${!arg_idx}"
+                python_args+=("$num_processes")
+            fi
+            ;;
+        -h|--help)
+            echo "Usage: $0 [--num-processes N] [--output DIR] [python args...]"
+            exit 0
+            ;;
+        --num-processes=*)
+            num_processes="${arg#*=}"
+            python_args+=("$arg")
+            ;;
+        -o|--output)
+            if ((arg_idx < $#)); then
+                ((arg_idx++))
+                output_dir="${!arg_idx}"
+            fi
+            ;;
+        --output=*)
+            output_dir="${arg#*=}"
+            ;;
+        *)
+            python_args+=("$arg")
+            ;;
+    esac
+done
+
+echo "Python Args: ${python_args[*]}"
+echo "Num Processes: $num_processes"
+echo "Output Dir: $output_dir"
+mkdir -p $output_dir
+
+export DG_JIT_WITH_LINEINFO=1 # for source counters
+
+echo "Warm up JIT cache"
+python tests/test_mega_moe.py --ncu-profile-only "${python_args[@]}"
+
+sleep 2
+
+ncu_args=(
+    --config-file off
+    --force-overwrite
+    --kernel-name sm100_fp8_fp4_mega_moe_impl
+    --import-source yes
+    --replay-mode application
+    --section PmSampling
+    --section SourceCounters
+    --rule LocalMemoryUsage
+    --launch-skip 0
+    --launch-count 1
+    --lockstep-kernel-launch
+    --communicator tcp
+    --clock-control none
+    --pm-sampling-interval 1000
+    --pm-sampling-max-passes 1
+    --disable-pm-warp-sampling
+    --communicator-tcp-num-peers "$num_processes"
+    --kill yes
+    --app-replay-buffer memory
+)
+
+echo "Run Job"
+
+for ((i = 0; i < num_processes; ++i)); do
+    ncu ${ncu_args[@]} -o "${output_dir%/}/mega-moe.$i" \
+        python tests/test_mega_moe.py \
+            --local-rank-idx=$i \
+            --ncu-profile-only \
+            "${python_args[@]}" &
+done
+
+echo "Waiting"
+wait
+echo "Done"