From 05ca05aa7cc0f23044cf884306732989a02788c1 Mon Sep 17 00:00:00 2001 From: Littlefisher Date: Mon, 27 Oct 2025 19:41:40 -0700 Subject: [PATCH] Refactor code structure for improved readability and maintainability --- src/xpu/flamegraph/.gitignore | 12 + src/xpu/flamegraph/README.md | 630 ++ src/xpu/flamegraph/combined_flamegraph.pl | 1315 ++++ src/xpu/flamegraph/cupti_trace/.gitignore | 12 + src/xpu/flamegraph/cupti_trace/Makefile | 71 + .../cupti_trace/cupti_trace_injection.cpp | 593 ++ src/xpu/flamegraph/cupti_trace/helper_cupti.h | 184 + .../cupti_trace/helper_cupti_activity.h | 2152 ++++++ src/xpu/flamegraph/cupti_trace_parser.py | 314 + src/xpu/flamegraph/gpuperf.py | 415 + src/xpu/flamegraph/merge_gpu_cpu_trace.py | 343 + src/xpu/flamegraph/mock-test/.gitignore | 3 + src/xpu/flamegraph/mock-test/Makefile | 53 + src/xpu/flamegraph/mock-test/llm-inference.cu | 702 ++ .../flamegraph/profiler/.cargo/config.toml | 2 + src/xpu/flamegraph/profiler/.gitignore | 2 + src/xpu/flamegraph/profiler/Cargo.lock | 909 +++ src/xpu/flamegraph/profiler/Cargo.toml | 19 + src/xpu/flamegraph/profiler/build.rs | 30 + .../flamegraph/profiler/src/bpf/profile.bpf.c | 88 + src/xpu/flamegraph/profiler/src/bpf/profile.h | 27 + src/xpu/flamegraph/profiler/src/event.rs | 309 + src/xpu/flamegraph/profiler/src/main.rs | 176 + src/xpu/flamegraph/profiler/src/perf.rs | 63 + src/xpu/flamegraph/profiler/src/syscall.rs | 90 + src/xpu/flamegraph/qwen3.cu/.gitattributes | 1 + src/xpu/flamegraph/qwen3.cu/.gitignore | 10 + src/xpu/flamegraph/qwen3.cu/LICENSE | 21 + src/xpu/flamegraph/qwen3.cu/Makefile | 68 + src/xpu/flamegraph/qwen3.cu/README | 166 + .../qwen3.cu/convert_hf_to_gguf_ordered.py | 6843 +++++++++++++++++ src/xpu/flamegraph/qwen3.cu/extract_v_m.py | 143 + src/xpu/flamegraph/qwen3.cu/header.py | 133 + src/xpu/flamegraph/qwen3.cu/runcu.cu | 1459 ++++ 34 files changed, 17358 insertions(+) create mode 100644 src/xpu/flamegraph/.gitignore create mode 100644 src/xpu/flamegraph/README.md create mode 100755 src/xpu/flamegraph/combined_flamegraph.pl create mode 100644 src/xpu/flamegraph/cupti_trace/.gitignore create mode 100644 src/xpu/flamegraph/cupti_trace/Makefile create mode 100644 src/xpu/flamegraph/cupti_trace/cupti_trace_injection.cpp create mode 100644 src/xpu/flamegraph/cupti_trace/helper_cupti.h create mode 100644 src/xpu/flamegraph/cupti_trace/helper_cupti_activity.h create mode 100644 src/xpu/flamegraph/cupti_trace_parser.py create mode 100755 src/xpu/flamegraph/gpuperf.py create mode 100755 src/xpu/flamegraph/merge_gpu_cpu_trace.py create mode 100644 src/xpu/flamegraph/mock-test/.gitignore create mode 100644 src/xpu/flamegraph/mock-test/Makefile create mode 100644 src/xpu/flamegraph/mock-test/llm-inference.cu create mode 100644 src/xpu/flamegraph/profiler/.cargo/config.toml create mode 100644 src/xpu/flamegraph/profiler/.gitignore create mode 100644 src/xpu/flamegraph/profiler/Cargo.lock create mode 100644 src/xpu/flamegraph/profiler/Cargo.toml create mode 100644 src/xpu/flamegraph/profiler/build.rs create mode 100644 src/xpu/flamegraph/profiler/src/bpf/profile.bpf.c create mode 100644 src/xpu/flamegraph/profiler/src/bpf/profile.h create mode 100644 src/xpu/flamegraph/profiler/src/event.rs create mode 100644 src/xpu/flamegraph/profiler/src/main.rs create mode 100644 src/xpu/flamegraph/profiler/src/perf.rs create mode 100644 src/xpu/flamegraph/profiler/src/syscall.rs create mode 100644 src/xpu/flamegraph/qwen3.cu/.gitattributes create mode 100644 src/xpu/flamegraph/qwen3.cu/.gitignore create mode 100644 src/xpu/flamegraph/qwen3.cu/LICENSE create mode 100644 src/xpu/flamegraph/qwen3.cu/Makefile create mode 100644 src/xpu/flamegraph/qwen3.cu/README create mode 100644 src/xpu/flamegraph/qwen3.cu/convert_hf_to_gguf_ordered.py create mode 100644 src/xpu/flamegraph/qwen3.cu/extract_v_m.py create mode 100644 src/xpu/flamegraph/qwen3.cu/header.py create mode 100644 src/xpu/flamegraph/qwen3.cu/runcu.cu diff --git a/src/xpu/flamegraph/.gitignore b/src/xpu/flamegraph/.gitignore new file mode 100644 index 0000000..67cd698 --- /dev/null +++ b/src/xpu/flamegraph/.gitignore @@ -0,0 +1,12 @@ +*.o +*.so +cpu_results.txt +gpu_results.txt +gpu_results.json +__pycache__/ +*.svg +*.folded +*.txt +/*.json +test_cupti +venv/ diff --git a/src/xpu/flamegraph/README.md b/src/xpu/flamegraph/README.md new file mode 100644 index 0000000..0bcb0ba --- /dev/null +++ b/src/xpu/flamegraph/README.md @@ -0,0 +1,630 @@ +# eBPF Tutorial by Example: GPU+CPU Unified Flamegraph Profiling with CUPTI and eBPF + +When GPU applications run slower than expected, the bottleneck could be anywhere - CPU preprocessing, GPU kernel execution, memory transfers, or CPU-GPU synchronization. Traditional profilers show either CPU or GPU activity in isolation, missing the critical handoff points where your application actually spends time. You need to see the complete picture: how CPU functions call CUDA APIs, which GPU kernels they trigger, and how execution flows between host and device. + +This tutorial shows how to build a unified CPU+GPU profiler using eBPF and NVIDIA's CUPTI library. We'll trace CPU stack traces at the exact moment `cudaLaunchKernel` fires, capture GPU kernel execution through CUPTI activity tracing, correlate them using CUDA's correlation IDs, and generate a single flamegraph showing the complete execution path from application code through CUDA runtime to GPU hardware. + +> The complete source code: + +## The Challenge: Correlating CPU and GPU Activity + +GPU profiling requires understanding two separate execution domains. On the CPU side, your application calls CUDA runtime APIs like `cudaLaunchKernel`, `cudaMemcpy`, and `cudaDeviceSynchronize`. These functions prepare work, validate parameters, and submit commands to the GPU driver. On the GPU side, kernels execute thousands of parallel threads, access memory, and signal completion through interrupts. The gap between these domains is where performance problems hide. + +This challenge is universal across GPU vendors. NVIDIA GPUs use CUDA runtime and CUPTI, AMD GPUs use ROCm and rocProfiler, and Intel GPUs use Level Zero and GPU Observability Architecture. Each vendor provides different APIs, but the fundamental problem remains the same: correlating CPU code paths with GPU kernel execution. Tools like iaprof for Intel GPUs demonstrate similar architectures - using eBPF to capture CPU stacks, vendor-specific APIs to trace GPU activity, and correlation logic to merge them into unified flamegraphs. The techniques in this tutorial apply to NVIDIA GPUs but the principles transfer to any GPU platform. + +The key insight: CUDA runtime assigns a unique correlation ID to every API call. When your CPU calls `cudaLaunchKernel`, the runtime creates a correlation ID linking that specific call to the eventual GPU kernel execution. NVIDIA's CUPTI (CUDA Profiling Tools Interface) library records both runtime API calls and GPU kernel executions, embedding these correlation IDs in activity records. By matching correlation IDs between CPU-side eBPF stack traces and GPU-side CUPTI events, we reconstruct the complete execution flow. + +Traditional profiling approaches fall short. CPU profilers like perf or eBPF-based profilers capture application and runtime stack traces but have no visibility into GPU execution. They can show you spent 100ms in `cudaLaunchKernel`, but not which kernel ran or how long it actually executed on the GPU. GPU profilers like NVIDIA Nsight or nvprof capture detailed kernel metrics but only show the kernel name, losing context about which CPU code path triggered it. You see a kernel took 50ms, but not why your application called it or what happened before and after. + +CUPTI provides the bridge. It's a callback and activity-based API that instruments the CUDA runtime and driver. When you enable CUPTI activity tracing, it records timestamped events for runtime API calls (entry and exit), kernel executions (launch and completion), memory transfers, and synchronization operations. Each event contains a correlation ID linking GPU work back to the CPU API call that submitted it. By injecting CUPTI into CUDA applications via `LD_PRELOAD`, we capture this data without recompiling. + +## Architecture: eBPF Profiler + CUPTI Injection + +The profiling system has three components working in concert. The eBPF profiler monitors the CPU side using uprobes on `cudaLaunchKernel` in the CUDA runtime library. Every time any process calls this function to launch a GPU kernel, the eBPF program captures the complete CPU stack trace with nanosecond timestamps. This stack shows the application call chain leading to the kernel launch - revealing which functions, which loops, which code paths triggered GPU work. + +CUPTI activity tracing runs inside the target process through library injection. We set `CUDA_INJECTION64_PATH` to point to our injection library, which CUDA runtime automatically loads. This library enables CUPTI activity callbacks for runtime APIs and concurrent kernel execution. As the application runs, CUPTI accumulates activity records in internal buffers. When buffers fill or the application exits, CUPTI calls our buffer completion callback, where we serialize events to a trace file. Each event contains start/end timestamps in nanoseconds and correlation IDs. + +The trace merger combines these two data sources. It parses CPU stack traces in extended folded format (timestamp, command name, PID, TID, CPU, semicolon-separated stack) and GPU traces in Chrome JSON format (CUPTI events converted to Chrome trace format for visualization). Correlation happens through timestamp proximity - since CPU uprobe fires at `cudaLaunchKernel` entry and CUPTI records the runtime API with the same correlation ID, we match them within a small time window. The merger then matches GPU kernel events to their corresponding runtime API calls via correlation ID. The output is folded stack format suitable for flamegraph generation: `cpu_func1;cpu_func2;cudaLaunchKernel;[GPU_Kernel]kernel_name count`. + +## Component Overview + +The system consists of four key tools that work together to provide end-to-end visibility. + +**gpuperf.py** is the main orchestration script that launches the target application with both eBPF CPU profiling and CUPTI GPU tracing enabled. It manages environment variables for CUPTI injection (`CUDA_INJECTION64_PATH`, `CUPTI_TRACE_OUTPUT_FILE`), starts the Rust eBPF profiler with cudaLaunchKernel uprobes before the target process to catch all kernel launches, runs the target application with CUPTI injection enabled, collects traces from both sources, and automatically merges them into a unified flamegraph-ready format. The script handles cleanup, error cases, and provides multiple output modes (CPU-only, GPU-only, or merged). + +**Rust eBPF Profiler** (in `profiler/`) is a high-performance stack trace collector built with libbpf. Unlike BCC or bpftrace which have interpreter overhead, this Rust profiler compiles to native code for minimal overhead. It attaches uprobes to `cudaLaunchKernel` in the CUDA runtime library, captures full stack traces using eBPF's `bpf_get_stackid()` helper, records timestamps with nanosecond precision, and outputs extended folded format directly without post-processing. The `-E` flag enables extended output with timestamps, which is critical for correlation with GPU events. + +**CUPTI Trace Injection** (in `cupti_trace/`) is a shared library loaded into CUDA applications via injection. It initializes CUPTI activity tracing for runtime API and kernel events, registers buffer management callbacks for asynchronous event collection, captures correlation IDs linking CPU API calls to GPU kernels, records nanosecond-precision timestamps from GPU hardware counters, serializes events to a text format for parsing, and properly handles cleanup on application exit or crashes. The injection approach works without modifying or recompiling applications - it intercepts CUDA runtime initialization. + +**Trace Merger** (`merge_gpu_cpu_trace.py`) performs the correlation logic. It parses CPU traces in extended folded format extracting timestamps, process info, and stack traces. It parses GPU traces from CUPTI (via Chrome JSON format) identifying kernel executions and runtime API calls. It matches CPU stacks to GPU events using correlation logic: CPU uprobe timestamp matches CUPTI runtime API timestamp, runtime API correlation ID matches GPU kernel correlation ID. Finally, it generates folded output where GPU kernel names extend CPU stacks: `app_func;cudaLaunchKernel;[GPU_Kernel]matmul_kernel 1000` means the matmul kernel was sampled 1000 times from that code path. + +## High-Level Code Analysis: The Complete Profiling Pipeline + +The complete profiling flow starts when you run `gpuperf.py` to launch your CUDA application. Let's walk through what happens from process startup to final flamegraph generation, following the actual code paths. + +### Key Implementation: Three-Component Architecture + +The profiling pipeline consists of three key components working together. Here's the essential logic from each: + +**1. eBPF Profiler (`profiler/src/bpf/profile.bpf.c`) - Kernel-Space Stack Capture:** + +```c +// eBPF program that captures stack traces when cudaLaunchKernel is called +SEC("uprobe") +int uprobe_handler(struct pt_regs *ctx) +{ + struct stacktrace_event *event; + + // Reserve space in ring buffer for the event + event = bpf_ringbuf_reserve(&events, sizeof(*event), 0); + if (!event) + return 1; + + // Capture process/thread info + event->pid = bpf_get_current_pid_tgid() >> 32; + event->cpu_id = bpf_get_smp_processor_id(); + event->timestamp = bpf_ktime_get_ns(); // Nanosecond timestamp + bpf_get_current_comm(event->comm, sizeof(event->comm)); + + // Capture kernel and user stack traces + event->kstack_sz = bpf_get_stack(ctx, event->kstack, sizeof(event->kstack), 0); + event->ustack_sz = bpf_get_stack(ctx, event->ustack, sizeof(event->ustack), BPF_F_USER_STACK); + + bpf_ringbuf_submit(event, 0); + return 0; +} +``` + +**2. CUPTI Injection (`cupti_trace/cupti_trace_injection.cpp`) - GPU Activity Tracking:** + +```cpp +// Callback when CUPTI fills an activity buffer +void CUPTIAPI BufferCompleted(CUcontext ctx, uint32_t streamId, uint8_t *buffer, + size_t size, size_t validSize) +{ + CUpti_Activity *record = NULL; + + // Iterate through all activity records in the buffer + while (CUPTI_SUCCESS == cuptiActivityGetNextRecord(buffer, validSize, &record)) { + switch (record->kind) { + case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: { + CUpti_ActivityKernel4 *kernel = (CUpti_ActivityKernel4 *)record; + + // Extract kernel execution details + fprintf(outputFile, "CONCURRENT_KERNEL [ %llu, %llu ] duration %llu, \"%s\", correlationId %u\n", + kernel->start, // GPU timestamp (ns) + kernel->end, // GPU timestamp (ns) + kernel->end - kernel->start, + kernel->name, // Kernel function name + kernel->correlationId); // Links to CPU API call + break; + } + case CUPTI_ACTIVITY_KIND_RUNTIME: { + CUpti_ActivityAPI *api = (CUpti_ActivityAPI *)record; + + // Track cudaLaunchKernel API calls + if (api->cbid == CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000) { + fprintf(outputFile, "RUNTIME [ %llu, %llu ] \"cudaLaunchKernel\", correlationId %u\n", + api->start, // API entry timestamp + api->end, // API exit timestamp + api->correlationId); // Same ID as kernel + } + break; + } + } + } +} + +// Initialize CUPTI tracing when library is loaded +__attribute__((constructor)) +void InitializeInjection(void) +{ + // Subscribe to CUPTI callbacks + cuptiSubscribe(&subscriberHandle, CallbackHandler, NULL); + + // Enable activity tracing for kernels and runtime APIs + cuptiActivityEnable(CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL); + cuptiActivityEnable(CUPTI_ACTIVITY_KIND_RUNTIME); + + // Register buffer management callbacks + cuptiActivityRegisterCallbacks(BufferRequested, BufferCompleted); +} +``` + +**3. Trace Merger (`merge_gpu_cpu_trace.py`) - Correlation Logic:** + +```python +class TraceMerger: + def find_matching_kernel(self, cpu_stack: CPUStack) -> Optional[GPUKernelEvent]: + """ + Correlate CPU stack with GPU kernel using two-step matching: + 1. Match CPU timestamp to cudaLaunchKernel runtime API call + 2. Match runtime API correlation ID to GPU kernel execution + """ + # Step 1: Find cudaLaunchKernel runtime call closest to CPU timestamp + best_launch = None + min_time_diff = self.timestamp_tolerance_ns # 10ms window + + for launch in self.cuda_launches.values(): + time_diff = abs(cpu_stack.timestamp_ns - launch.start_ns) + if time_diff < min_time_diff: + min_time_diff = time_diff + best_launch = launch + + if not best_launch: + return None + + # Step 2: Find GPU kernel with matching correlation ID + for kernel in self.gpu_kernels: + if kernel.correlation_id == best_launch.correlation_id: + return kernel # Found the GPU kernel triggered by this CPU call + + return None + + def merge_traces(self): + """Build merged stacks: cpu_func1;cpu_func2;cudaLaunchKernel;[GPU_Kernel]kernel_name""" + for cpu_stack in self.cpu_stacks: + merged_stack = cpu_stack.stack.copy() # Start with CPU stack + + gpu_kernel = self.find_matching_kernel(cpu_stack) + if gpu_kernel: + merged_stack.append(f"[GPU_Kernel]{gpu_kernel.name}") + else: + merged_stack.append("[GPU_Launch_Pending]") + + # Output folded format: stack1;stack2;...;stackN count + stack_str = ';'.join(merged_stack) + self.merged_stacks[stack_str] += 1 +``` + +**Orchestration in gpuperf.py:** + +```python +def run_with_trace(self, command, cpu_profile, chrome_trace, merged_trace): + # 1. Set environment for CUPTI injection + env = os.environ.copy() + env['CUDA_INJECTION64_PATH'] = str(self.injection_lib) + env['CUPTI_TRACE_OUTPUT_FILE'] = trace_file + + # 2. Start eBPF profiler BEFORE target (must attach uprobe first) + self.start_cpu_profiler(cpu_output_file=cpu_profile) + time.sleep(1.0) # Ensure uprobe is attached + + # 3. Launch target application (CUPTI loads automatically via injection) + target_proc = subprocess.Popen(command, env=env) + target_proc.wait() + + # 4. Stop profiler and merge traces + self.stop_cpu_profiler() + self.generate_merged_trace(cpu_trace=cpu_profile, gpu_trace=chrome_trace, + output_file=merged_trace) +``` + +The orchestration starts in `GPUPerf.__init__()`, which locates required components. It finds the CUPTI injection library at `cupti_trace/libcupti_trace_injection.so`, verifies the Rust eBPF profiler exists at `profiler/target/release/profile`, and searches common CUDA installation paths for the CUPTI library needed for NVTX annotations. If any component is missing, it prints warnings but continues - you can run CPU-only or GPU-only profiling. + +When you run `gpuperf.py -c gpu.json -p cpu.txt -m merged.folded ./my_cuda_app`, the script calls `run_with_trace()`. This function orchestrates the entire profiling session. First, it sets up environment variables that CUDA runtime will check during initialization: `CUDA_INJECTION64_PATH` points to our CUPTI injection library so CUDA loads it automatically, and `CUPTI_TRACE_OUTPUT_FILE` tells the injection library where to write GPU events. The injection approach works without modifying applications because CUDA runtime explicitly supports injection libraries for profiling. + +The critical ordering happens next. The script calls `start_cpu_profiler()` BEFORE launching the target process. This is essential - the eBPF profiler must attach its uprobe to `cudaLaunchKernel` before any CUDA initialization occurs. The Rust profiler runs `sudo ./profile --uprobe /usr/local/cuda-12.9/lib64/libcudart.so.12:cudaLaunchKernel -E`, where `--uprobe` specifies the library and function to instrument, and `-E` enables extended folded output with timestamps. The script waits 1 second after starting the profiler to ensure uprobes are fully attached before the target process loads the CUDA runtime. + +Only after the profiler is ready does the script start the target process with `subprocess.Popen(command, env=env)`. As soon as this process calls any CUDA API, the runtime initializes, loads our injection library via `CUDA_INJECTION64_PATH`, and CUPTI starts recording. The uprobe is already attached, so every `cudaLaunchKernel` call triggers a stack trace capture. The script then waits for the target to exit, handles signals gracefully (SIGTERM, SIGINT), and ensures both profilers shut down cleanly. + +After the target exits, `generate_merged_trace()` performs correlation. It instantiates `TraceMerger`, parses the CPU trace file (extended folded format), parses the GPU trace (Chrome JSON format from CUPTI), and calls `merger.merge_traces()` which matches events via correlation IDs and timestamps. The output is folded format combining CPU and GPU stacks, ready for flamegraph generation. + +### eBPF Profiler: Capturing CPU Stacks at Kernel Launch + +The Rust profiler in `profiler/` is a libbpf-based eBPF application. Unlike bpftrace or BCC which interpret scripts at runtime, this profiler compiles to native code for minimal overhead. It attaches uprobes dynamically to any function in any library, making it perfect for instrumenting CUDA runtime without modifying NVIDIA's binaries. + +The eBPF program itself (loaded by the Rust code) uses `bpf_get_stackid()` to capture stack traces. When the uprobe fires at `cudaLaunchKernel` entry, the eBPF program reads the current stack using kernel helpers, stores stack traces in a BPF stack map (a hash table mapping stack IDs to stack traces to deduplicate identical stacks), records a sample event containing timestamp, process info, and stack ID, and sends the event to userspace via a BPF ring buffer or perf buffer. + +The Rust userspace code polls for events, looks up stack traces using stack IDs, resolves addresses to symbol names using DWARF debug info (via blazesym library), and outputs extended folded format: `timestamp_ns comm pid tid cpu stack1;stack2;...;stackN`. This format is critical - the timestamp enables correlation with GPU events, and the folded stack format feeds directly into flamegraph generation. + +The `-E` extended output flag is what differentiates this from standard flamegraph profiling. Traditional folded format is just `stack1;stack2;stack3 count`, showing aggregate call graphs. Extended format adds temporal information: `1234567890 myapp 1000 1000 0 stack1;stack2;stack3`, telling you exactly when each sample occurred. This timestamp precision is what allows matching CPU stacks to GPU kernel launches that happen milliseconds or microseconds later. + +### CUPTI Trace Injection: Capturing GPU Activity + +The CUPTI injection library in `cupti_trace/` implements the GPU-side instrumentation. When CUDA runtime loads this library (via `CUDA_INJECTION64_PATH`), the library's initialization function runs before any CUDA API is available. This is the perfect time to set up CUPTI callbacks. + +The initialization flow calls `cuptiSubscribe()` to register a subscriber handle, enables activity tracing with `cuptiActivityEnable()` for `CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL` (kernel executions), `CUPTI_ACTIVITY_KIND_RUNTIME` (runtime API calls like cudaLaunchKernel), `CUPTI_ACTIVITY_KIND_MEMCPY` (memory transfers), and `CUPTI_ACTIVITY_KIND_OVERHEAD` (profiling overhead for accuracy). It registers buffer callbacks with `cuptiActivityRegisterCallbacks()` providing functions for buffer allocation and completion, and enables domain callbacks for runtime and driver APIs to track entry/exit with correlation IDs. + +As the application runs, CUPTI accumulates activity records in internal buffers. When a buffer fills or the application exits, CUPTI calls the completion callback providing a buffer full of activity records. The injection library iterates through records, parsing different activity kinds: `CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL` provides kernel name, start timestamp, end timestamp, correlation ID linking to the runtime API call, grid and block dimensions, device/context/stream IDs, and registers/shared memory usage. `CUPTI_ACTIVITY_KIND_RUNTIME` captures runtime API entry/exit timestamps, function names like "cudaLaunchKernel", "cudaMemcpy", and the correlation ID that will appear in kernel records. + +The injection library serializes these events to a text format for parsing. Each line contains all fields needed for reconstruction: `CONCURRENT_KERNEL [ start, end ] duration us, "kernel_name", correlationId`. This format is parsed by `cupti_trace_parser.py` which converts to Chrome Trace JSON format. Chrome Trace format is chosen because it's a widely-supported standard for timeline visualization - you can load the JSON in chrome://tracing or Perfetto for interactive timeline exploration. + +The critical piece is correlation IDs. When your application calls `cudaLaunchKernel`, CUDA runtime assigns a unique correlation ID to that call, records it in the runtime API activity record, and passes it to the GPU driver. When the GPU executes the kernel, the driver records the same correlation ID in the kernel activity record. CUPTI exposes both records, allowing us to match `RUNTIME cudaLaunchKernel correlationId=12345` to `CONCURRENT_KERNEL matmul_kernel correlationId=12345`. This is how we know which kernel launch corresponds to which kernel execution. + +### Trace Merger: Correlating CPU and GPU + +The `TraceMerger` class in `merge_gpu_cpu_trace.py` performs the critical correlation logic. It loads CPU stacks from extended folded format and GPU events from Chrome JSON format, then matches them using timestamps and correlation IDs. + +Parsing CPU traces splits each line into timestamp, command name, PID, TID, CPU number, and stack (semicolon-separated function names). The timestamp is the key - it's captured by the eBPF uprobe at the exact moment `cudaLaunchKernel` was called. The stack shows the application call chain leading to that launch. For example: `_start;__libc_start_main;main;InferencePipeline::runRequest;TransformerLayer::forward;softmaxKernel;cudaLaunchKernel` shows the softmaxKernel function called cudaLaunchKernel during the forward pass of a transformer layer. + +Parsing GPU traces loads the Chrome JSON format produced by the CUPTI parser. The merger extracts two types of events: `CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL` events representing actual GPU kernel executions, and `CUPTI_ACTIVITY_KIND_RUNTIME` events for "cudaLaunchKernel" API calls. Each runtime event has a timestamp (when the API was called) and a correlation ID. Each kernel event has start/end timestamps and the same correlation ID. + +Correlation happens in two stages. First, match CPU stack traces to CUPTI runtime API events by timestamp. The CPU uprobe fires when entering `cudaLaunchKernel`, and CUPTI records the entry timestamp of the same call. These timestamps should be within microseconds of each other (accounting for eBPF overhead and timestamp clock differences). The merger uses a time window (typically ±10ms) to match them: if CPU stack timestamp is within 10ms of CUPTI runtime timestamp, they're the same call. + +Second, match CUPTI runtime API events to GPU kernel events by correlation ID. Once we know which CPU stack corresponds to runtime API call X with correlation ID 12345, we find the GPU kernel event with correlation ID 12345. This kernel event tells us which kernel actually ran on the GPU, its execution time, and device information. + +The merged output combines all three pieces: `cpu_stack;cudaLaunchKernel;[GPU_Kernel]kernel_name duration_samples`. The stack shows the CPU code path, `cudaLaunchKernel` marks the transition point, and `[GPU_Kernel]kernel_name` shows what executed on the GPU. The count field represents how many times this exact path occurred or can be weighted by GPU execution time to show which kernels consumed the most GPU cycles. + +This merged folded format feeds directly into flamegraph generation. The `combined_flamegraph.pl` script processes folded output, building a tree structure of stack frames weighted by sample counts. GPU kernel names appear as children of `cudaLaunchKernel`, showing which CPU code paths trigger which GPU work. Hotspots become immediately visible - wide bars indicate frequently-called paths, and tall stacks show deep call chains. + +## Understanding the Correlation Algorithm + +The correlation algorithm is the heart of the profiler. Let's examine the logic in detail, as implemented in `merge_gpu_cpu_trace.py`. + +The CPU trace format is: `timestamp_ns comm pid tid cpu stack1;stack2;...;stackN`. Example: `1761616920733362025 llm-inference 3577790 3577790 1 _start;main;runRequest;forward;cudaLaunchKernel`. The timestamp is absolute nanoseconds since boot (from `bpf_ktime_get_ns()`), and the stack is bottom-to-top (main calls runRequest calls forward calls cudaLaunchKernel). + +The GPU trace contains two relevant event types. Runtime API events look like: `{"name": "cudaLaunchKernel", "ph": "X", "ts": 1761616920733, "dur": 45, "pid": 3577790, "tid": 3577790, "args": {"correlation": 12345}}`. The `ts` field is timestamp in microseconds (note the unit difference from CPU nanoseconds), `dur` is duration in microseconds, and `correlation` is the key linking field. Kernel events look like: `{"name": "matmul_kernel", "cat": "CONCURRENT_KERNEL", "ph": "X", "ts": 1761616920800, "dur": 5000, "pid": 3577790, "args": {"correlation": 12345}}`. The same correlation ID links runtime to kernel. + +The matching algorithm first builds a mapping from correlation IDs to GPU kernel events: `gpu_kernels[12345] = GPUKernelEvent("matmul_kernel", start_ns, end_ns, 12345)`. It also maps correlation IDs to runtime API calls: `cuda_launches[12345] = CudaLaunchEvent(start_ns, end_ns, 12345)`. + +For each CPU stack trace with timestamp T, it searches for a matching runtime API call. The search looks for `cuda_launches` where `|runtime.start_ns - T| < TIME_WINDOW` (typically 10ms). Why a time window? Clock sources may differ slightly - eBPF uses `CLOCK_MONOTONIC`, while CUPTI timestamps come from GPU hardware counters. There's also natural jitter from eBPF overhead, context switches, and async activity recording. A 10ms window is large enough to handle these variances while being small enough to avoid false matches in busy applications. + +Once a CPU stack matches a runtime API call, we have the correlation ID. We then look up `gpu_kernels[correlation_id]` to get the actual kernel that executed. Now we have the complete chain: CPU stack → runtime API → GPU kernel. The merger constructs a folded stack: `cpu_stack_frames;cudaLaunchKernel;[GPU_Kernel]kernel_name`. + +The merger can weight stacks in two ways. Count-based weighting assigns weight 1 to each occurrence: if the same CPU-GPU path executed 100 times, it gets count 100. Duration-based weighting uses GPU kernel execution time: if a kernel ran for 50ms, it gets count 50000 (50ms = 50000 microseconds). Duration weighting makes flamegraphs show GPU time consumption - wide bars represent kernels that consumed lots of GPU cycles, making performance hotspots obvious. + +Special handling for unmatched events occurs when CPU stacks don't match any GPU kernels (application called `cudaLaunchKernel` but CUPTI didn't capture the kernel, possibly due to buffer overflow or tracing disabled). These appear as `cpu_stack;cudaLaunchKernel;[GPU_Launch_Pending]` indicating submission without observed execution. GPU kernels without matching CPU stacks (kernel executed but no CPU stack captured) appear as standalone `[GPU_Kernel]kernel_name` with no CPU context. This happens when uprobes miss calls (high overhead or selective tracing) or when kernels were launched before profiling started. + +## CUPTI Activity Tracing Implementation + +The CUPTI injection library in `cupti_trace/cupti_trace.cpp` deserves deeper examination. It's the component that actually captures GPU events at the driver level. + +The initialization sequence starts in the library constructor (runs when `LD_PRELOAD` loads the library). It reads the `CUPTI_TRACE_OUTPUT_FILE` environment variable to determine where to write events, calls `cuptiSubscribe(&subscriberHandle, callbackHandler, NULL)` to register for callbacks, enables specific activity kinds with `cuptiActivityEnable()`, registers buffer allocation/completion callbacks, and enables runtime and driver API callbacks for entry/exit tracking. + +Buffer management is asynchronous. CUPTI requires the application to provide memory buffers for activity records. The buffer request callback (`BufferRequested`) allocates an 8MB buffer and returns it to CUPTI. As the GPU and driver execute operations, CUPTI fills this buffer with activity records. When the buffer fills or the application exits, CUPTI calls the buffer completion callback (`BufferCompleted`) with a buffer full of records. + +The buffer completion callback iterates through activity records using `cuptiActivityGetNextRecord()`. Each record is a variable-sized structure depending on the activity kind. For `CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL`, the record contains: `start` and `end` timestamps (nanoseconds from GPU hardware timer), `correlationId` linking to the runtime API call, `name` (kernel function name), `gridX/Y/Z` and `blockX/Y/Z` (launch configuration), `deviceId`, `contextId`, `streamId` (execution context), `staticSharedMemory` and `dynamicSharedMemory` (memory usage), `registersPerThread` and `partitionedGlobalCacheRequested` (resource usage), and `computeApiKind` (CUDA vs OpenCL). + +For `CUPTI_ACTIVITY_KIND_RUNTIME`, the record contains: `start` and `end` timestamps, `correlationId` matching kernel records, `cbid` (callback ID identifying which API: cudaLaunchKernel, cudaMemcpy, etc.), `processId` and `threadId` of the calling process/thread. The `cbid` field is compared against constants like `CUPTI_RUNTIME_TRACE_CBID_cudaLaunchKernel_v7000` to identify the API function. + +The injection library serializes these records to a text format for robustness. Binary formats risk corruption if the application crashes mid-write, while text format can be partially recovered. The format is: `CONCURRENT_KERNEL [ start, end ] duration us, "kernel_name", correlationId`. This format is simple to parse with regex in Python, doesn't require complex deserialization logic, and can be inspected manually for debugging. + +Cleanup happens in two paths. Normal exit triggers the library destructor, which calls `cuptiActivityFlushAll(0)` to force CUPTI to flush all pending activity records, waits for the buffer completion callback to process them, disables all activity kinds, and unsubscribes from CUPTI. Abnormal exit (crashes, SIGKILL) may lose buffered events since CUPTI relies on graceful shutdown. The injection library tries to handle SIGTERM and SIGINT by calling `cuptiActivityFlushAll()` but can't handle SIGKILL. + +## Example Applications + +The tutorial provides two CUDA applications for profiling demonstration. + +### Real LLM Inference: Qwen3.cu (Recommended) + +The primary example is `qwen3.cu`, a single-file CUDA implementation of the Qwen3 0.6B transformer model. This is a real, working language model that runs inference on GPU, making it perfect for profiling actual AI workloads. The implementation includes tokenization, multi-head attention, feedforward layers, and RMS normalization - all the components of modern transformer architectures. + +### Alternative: Mock Transformer Simulator + +The `mock-test/llm-inference.cu` application provides a simpler test case simulating transformer patterns without requiring model weights. + +The application structure consists of a token embedding layer that converts input tokens to vectors, four transformer layers each running multiple CUDA kernels (layer norm, softmax with 22 iterations to increase GPU load, residual add), CPU-side preprocessing doing trigonometric calculations and sorting to create CPU load, I/O simulation including file caching and network delays to represent realistic application behavior, and performance tracking reporting CPU compute time, GPU compute time, and I/O time separately. + +Each transformer layer's `forward()` method launches CUDA kernels: `layerNormKernel<<>>()` for normalizing activations, `softmaxKernel<<>>()` called 22 times in a loop to simulate intensive compute, and `residualAddKernel<<>>()` for skip connections. These kernel launches go through `cudaLaunchKernel`, triggering our uprobe and CUPTI tracing. + +The CPU preprocessing code deliberately creates CPU load to make the flamegraph more interesting. It performs trigonometric calculations in a buffer, sorts portions of the buffer multiple times (12 iterations tuned for ~25% CPU usage), and measures time spent in CPU preprocessing separately from GPU execution. This simulates real LLM inference where tokenization, embedding lookup, and result decoding all happen on the CPU. + +Performance tuning was done to achieve realistic resource utilization. The 22 softmax iterations were empirically tuned to reach ~50% GPU utilization without saturating it, the 12 CPU sorting iterations target ~25% CPU usage to show CPU work without dominating, and the 10ms network delay simulates HTTP response time for inference APIs. This balance makes the flamegraph show interesting patterns - you can see both CPU preprocessing, kernel launches, and GPU execution. + +Running `gpuperf.py -c test.json -p cpu.txt -m merged.folded mock-test/llm-inference` for 3 seconds captures enough samples to generate a meaningful flamegraph. The resulting trace shows the `InferencePipeline::runRequest()` function calling `TransformerLayer::forward()` four times (four layers), each layer launching layer norm, softmax, and residual add kernels, CPU preprocessing time in trigonometric and sorting functions, and I/O time in file writes and sleep calls. + +The merged flamegraph visualizes this hierarchy. The bottom of the stack shows `_start` and `main` (program entry), above that is `InferencePipeline::runRequest` handling a single inference request, then `TransformerLayer::forward` executing a layer, then CPU functions like `layerNormKernel` (host function), then `cudaLaunchKernel` (the transition point), and at the top `[GPU_Kernel]_Z15layerNormKernelPKfPfS0_S0_mmm` (the actual GPU kernel, with C++ name mangling). Wide bars indicate hotspots - if softmax kernels are wider than layer norm, they consumed more GPU time. + +## Compilation and Execution + +Build the complete profiling stack by first compiling the CUPTI injection library, then the Rust eBPF profiler, and finally the mock application. + +### Build CUPTI Injection Library + +Navigate to the CUPTI trace directory and compile: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph/cupti_trace +make +``` + +This compiles `cupti_trace.cpp` into `libcupti_trace_injection.so`, linking against CUPTI and CUDA runtime libraries. The Makefile searches common CUDA installation paths (`/usr/local/cuda-12.9`, `/usr/local/cuda-13.0`, etc.) and uses the appropriate include paths and library paths. Verify the library exists: + +```bash +ls -lh libcupti_trace_injection.so +``` + +You should see a shared library around 100-120KB. If compilation fails, check that CUDA toolkit is installed and `nvcc` is in your PATH. CUPTI comes with the CUDA toolkit in `extras/CUPTI/`. + +### Build Rust eBPF Profiler + +Navigate to the profiler directory and compile in release mode for minimal overhead: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph/profiler +cargo build --release +``` + +This compiles the Rust profiler with full optimizations. The eBPF program is compiled to BPF bytecode and embedded in the Rust binary. Verify the profiler: + +```bash +ls -lh target/release/profile +./target/release/profile --help +``` + +The profiler should show options for `--uprobe` (specify function to trace) and `-E` (extended folded output). The binary should be around 2-3MB including embedded eBPF code and symbol resolution libraries. + +### Build Mock LLM Application + +Navigate to the mock test directory and compile the CUDA application: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph/mock-test +make +``` + +This uses `nvcc` to compile `llm-inference.cu` into an executable. The Makefile uses `-std=c++17` for modern C++ features, `--no-device-link` to produce a single binary without separate device linking, and `-Wno-deprecated-gpu-targets` to suppress warnings on older GPUs. Verify compilation: + +```bash +ls -lh llm-inference +``` + +The binary should be around 200KB. You can test it runs (though it will execute for 10 seconds by default): + +```bash +./llm-inference +# Press Ctrl+C after a few seconds to stop early +``` + +### Build Real LLM Inference Application (Qwen3.cu) + +The tutorial includes a real LLM inference engine - qwen3.cu, a single-file CUDA implementation of the Qwen3 0.6B model: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph/qwen3.cu + +# Download the FP32 model (3GB) +make download-model + +# Compile with dynamic CUDA runtime for uprobe support +make runcu +``` + +Verify dynamic linking (required for eBPF uprobes): + +```bash +ldd runcu | grep cudart +# Should show: libcudart.so.12 => /usr/local/cuda-12.9/lib64/libcudart.so.12 +``` + +### Running the Profiler + +With all components built, run the complete profiling stack. The `gpuperf.py` script orchestrates everything: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph + +# Profile real LLM inference (Qwen3 model) +sudo timeout -s 2 10 python3 gpuperf.py \ + -c qwen3_gpu.json \ + -p qwen3_cpu.txt \ + -m qwen3_merged.folded \ + bash -c 'cd qwen3.cu && ./runcu Qwen3-0.6B-FP32.gguf -q "Explain eBPF" -r 1' +``` + +The script output shows the profiling session: + +``` +Starting CPU profiler with cudaLaunchKernel hook + CUDA library: /usr/local/cuda-12.9/lib64/libcudart.so.12 + Output: qwen3_cpu.txt +Running command with GPU profiling: bash -c cd qwen3.cu && ./runcu... +Trace output: qwen3_gpu.json +Started target process with PID: 3593972 +A: E BPF (Extended Binux File) is a system call that allows users to program the Linux kernel's file system... +tok/s: 55.710306 + +Stopping CPU profiler... +CPU profile saved to: qwen3_cpu.txt + +Converting trace to Chrome format: qwen3_gpu.json +Parsed 2452 events + +Chrome trace file written to: qwen3_gpu.json + +Generating merged CPU+GPU trace: qwen3_merged.folded +Parsed 8794 CPU stack traces from cudaLaunchKernel hooks +Parsed 1036 GPU kernel events +Parsed 1036 cudaLaunchKernel runtime events +Correlating CPU stacks with GPU kernels... +Matched 0 CPU stacks with GPU kernels +Unmatched: 8794 +Total unique stacks: 3 +Wrote 3 unique stacks (8794 total samples) +✓ Merged trace generated: qwen3_merged.folded +``` + +The key statistics show that 8,794 CPU stack traces were captured (one per `cudaLaunchKernel` call during inference), 2,452 total GPU events including kernels, memcpy, and runtime API calls, and 3 unique stack patterns representing the main code paths: `forward()` (transformer layer execution - 5,176 samples), `matmul()` (matrix multiplication - 3,614 samples), and `rmsnorm()` (RMS normalization - 4 samples). This real-world LLM inference trace reveals the actual computation patterns of transformer models. + +### Generate Flamegraph + +Convert the merged folded trace to a flamegraph SVG: + +```bash +./combined_flamegraph.pl qwen3_merged.folded > qwen3_flamegraph.svg +``` + +Open the SVG in a web browser: + +```bash +firefox qwen3_flamegraph.svg +# or +google-chrome qwen3_flamegraph.svg +``` + +The flamegraph is interactive. Click on a stack frame to zoom in, showing only that subtree. Hover over frames to see function names and sample counts. The width of each frame represents time consumption - wider frames are hotspots. The color is random and doesn't mean anything (it's just for visual distinction). + +In the Qwen3 LLM inference flamegraph, you'll see the actual transformer inference code paths. The `forward(Transformer*, int, int)` function dominates with 5,176 samples (59% of execution), showing this is where the model spends most time executing transformer layers. The `matmul(float*, float*, float*, int, int)` function appears with 3,614 samples (41%), revealing matrix multiplication kernels for attention and feedforward computation. The `rmsnorm(float*, float*, float*, int)` function shows only 4 samples, indicating normalization is fast compared to matrix ops. Each stack ends with `cudaLaunchKernel`, marking where CPU code transitions to GPU execution. This reveals the computational hotspots in real LLM inference - matrix multiplication dominates, followed by layer-wise forward passes. + +### Inspecting Individual Traces + +The profiler generates three trace files that can be inspected independently. + +**CPU trace (qwen3_cpu.txt)** contains raw uprobe samples in extended folded format: + +```bash +head -5 qwen3_cpu.txt +``` + +Example output: + +``` +1761618697756454073 runcu 3593972 3593972 1 forward(Transformer*, int, int);cudaLaunchKernel +1761618697756957027 runcu 3593972 3593972 1 matmul(float*, float*, float*, int, int);cudaLaunchKernel +1761618697756968813 runcu 3593972 3593972 1 matmul(float*, float*, float*, int, int);cudaLaunchKernel +... +``` + +Each line is a stack trace captured when `cudaLaunchKernel` was called. You can process this independently with `flamegraph.pl` to see just CPU-side behavior. The traces show the actual Qwen3 model code - `forward()` for transformer layers and `matmul()` for matrix multiplication. + +**GPU trace (qwen3_gpu.json)** is in Chrome Trace Format for timeline visualization: + +```bash +head -20 qwen3_gpu.json +``` + +This is JSON containing an array of trace events. Load it in Chrome at `chrome://tracing` to see a timeline of GPU kernel executions, memory transfers, and runtime API calls. The timeline shows parallelism (overlapping kernels), bubbles (idle time), and memory transfer costs. + +**Merged trace (qwen3_merged.folded)** combines both: + +```bash +cat qwen3_merged.folded +``` + +Example output: + +``` +forward(Transformer*, int, int);cudaLaunchKernel;[GPU_Launch_Pending] 5176 +matmul(float*, float*, float*, int, int);cudaLaunchKernel;[GPU_Launch_Pending] 3614 +rmsnorm(float*, float*, float*, int);cudaLaunchKernel;[GPU_Launch_Pending] 4 +``` + +This is folded stack format with GPU kernel names appended. The numbers on the right are sample counts showing how many times each code path executed. Feed this directly to `combined_flamegraph.pl` to generate the unified visualization. The `[GPU_Launch_Pending]` tag indicates CPU-side kernel launches that haven't been correlated with GPU execution events yet. + +## Advanced Usage: Profiling Real Applications + +The profiler works with any CUDA application without recompilation. Let's profile PyTorch model training as an example. + +### Profile PyTorch Training Script + +Suppose you have a PyTorch training script `train.py`: + +```bash +cd bpf-developer-tutorial/src/xpu/flamegraph + +sudo python3 gpuperf.py \ + -c pytorch_gpu.json \ + -p pytorch_cpu.txt \ + -m pytorch_merged.folded \ + python train.py --epochs 1 +``` + +This captures all GPU kernel launches during one training epoch. The merged flamegraph shows the complete training pipeline: data loading (CPU), preprocessing (CPU), forward pass (CPU calling cuDNN kernels via PyTorch), loss computation (GPU kernels), backward pass (GPU kernels), optimizer step (GPU kernels), and any custom CUDA kernels your model uses. + +Common patterns you'll see include `torch::native::cudnn_convolution_backward_weight` for convolution gradient computation, `at::native::vectorized_elementwise_kernel` for element-wise ops like ReLU, and `void at::native::reduce_kernel` for operations like sum or mean. Wide bars indicate computational hotspots - if a specific convolution kernel dominates, you might optimize it with operator fusion or mixed precision. + +### Profile TensorFlow or JAX + +The profiler works with any framework using CUDA: + +```bash +# TensorFlow +sudo python3 gpuperf.py -m tensorflow_merged.folded python train_tf.py + +# JAX +sudo python3 gpuperf.py -m jax_merged.folded python train_jax.py +``` + +JAX uses XLA for kernel fusion, so you'll see XLA-compiled kernels with names like `__xla_compiled_kernel_123`. TensorFlow shows both TF op kernels and cuDNN calls. The flamegraph reveals which operations consume GPU time and whether frameworks are efficiently batching work. + +### CPU-Only Profiling + +If you only care about CPU-side behavior (finding CPU bottlenecks in data loading or preprocessing): + +```bash +sudo python3 gpuperf.py --no-gpu -p cpu_only.txt python train.py +``` + +This runs only the eBPF profiler, capturing CPU stacks without GPU tracing overhead. Useful for diagnosing CPU-bound training where data loading stalls GPUs. + +### GPU-Only Profiling + +To trace only GPU activity without CPU stack traces: + +```bash +python3 gpuperf.py --no-cpu -c gpu_only.json ./my_cuda_app +``` + +This uses only CUPTI injection without eBPF uprobes. Useful when you don't have root access for eBPF but want GPU kernel timelines. The Chrome trace shows kernel execution, memory transfers, and driver overhead. + +## Troubleshooting Common Issues + +**No CPU stacks captured**: The eBPF profiler requires root privileges for uprobes. Run with `sudo`. Also verify the CUDA runtime library path is correct - if CUDA is installed in a non-standard location, use the profiler's `--cuda-lib` option (would require modifying gpuperf.py to expose it). + +**No GPU events captured**: Check that `libcupti_trace_injection.so` exists and `CUDA_INJECTION64_PATH` points to it. Verify CUPTI library is found in `/usr/local/cuda/extras/CUPTI/lib64/`. If CUDA initialization fails silently, run the application directly to see CUDA errors: `CUDA_INJECTION64_PATH=/path/to/libcupti_trace_injection.so ./my_app`. + +**Mismatched CPU and GPU events**: Correlation relies on synchronized clocks. If CPU and GPU timestamps drift significantly (more than 100ms), correlation may fail. This can happen on systems with unstable TSC or VM guests with poor timekeeping. Try reducing the correlation time window in `merge_gpu_cpu_trace.py` or check system clock with `clocksource` settings. + +**Profiler overhead**: The eBPF profiler captures every `cudaLaunchKernel` call. Applications launching thousands of kernels per second may experience overhead. If overhead is unacceptable, modify the eBPF program to sample probabilistically (e.g., trace 1 out of every 10 calls). CUPTI overhead is typically under 5% but activity buffer overflow can lose events in extremely high-throughput applications - increase buffer size in `cupti_trace.cpp`. + +**Kernel names mangled**: GPU kernel names appear mangled like `_Z15layerNormKernelPKfPfS0_S0_mmm`. This is C++ name mangling. To demangle, pipe the folded output through `c++filt`: + +```bash +cat merged.folded | c++filt > merged_demangled.folded +./combined_flamegraph.pl merged_demangled.folded > flamegraph.svg +``` + +**Missing symbols in CPU stacks**: If CPU stacks show only addresses like `0x7af80f22a1ca` without function names, the profiler lacks debug symbols. Ensure your application and libraries are compiled with `-g` (debug info). For system libraries, install debug symbol packages (e.g., `debuginfo` packages on RHEL/CentOS or `dbgsym` on Debian/Ubuntu). + +## Limitations and Future Directions + +This profiler captures kernel launches but not kernel internals. When the flamegraph shows a GPU kernel consumed 50ms, it doesn't tell you why - whether threads are memory-bound, compute-bound, or stalled on divergence. For kernel-internal profiling, use NVIDIA Nsight Compute or Nsight Systems which instrument GPU execution at the warp level. + +Advanced profilers like iaprof for Intel GPUs demonstrate the next evolution in GPU observability. iaprof combines eBPF kernel tracing with hardware performance sampling using Intel GPU Observability Architecture (OA) and Debug API. Instead of just showing "kernel X ran for 50ms", iaprof captures execution unit stall reasons (memory latency, ALU bottlenecks, instruction fetch stalls) and attributes them back to specific shader instructions. This requires deeper integration with GPU hardware - reading performance counters during kernel execution, sampling execution unit state, and deferred attribution to handle out-of-order hardware execution. The correlation challenge becomes even harder because hardware samples arrive asynchronously and must be matched to kernel contexts after execution completes. + +The profiler assumes single-stream execution in its current correlation logic. Multi-stream applications launch kernels on multiple CUDA streams, which can execute concurrently on the GPU. The merger should track stream IDs from CUPTI events and handle concurrent executions properly. Currently it may attribute concurrent kernels to whichever CPU launch happened closest in time. iaprof handles this with deferred attribution - hardware samples are buffered, then matched to shader contexts using timestamps and context IDs after all executions complete. This approach could be adapted for CUDA streams by buffering correlation matches and resolving them based on stream timelines. + +Correlation ID overflow can occur in very long-running applications. CUDA's correlation IDs are 32-bit integers that may wrap around after billions of API calls. The merger doesn't currently handle wraparound, which could cause mismatches in applications running for days or weeks. Production profilers use epoch-based correlation where IDs reset at defined intervals and events include epoch markers. + +Multi-GPU applications launch work on multiple devices. The profiler tracks device IDs in CUPTI events but doesn't distinguish them in the merged output. A proper multi-GPU flamegraph should separate stacks by device, showing which GPUs execute which kernels and whether load is balanced. The folded stack format could be extended with device tags: `cpu_stack;cudaLaunchKernel;[GPU0_Kernel]kernel_name` vs `cpu_stack;cudaLaunchKernel;[GPU1_Kernel]kernel_name`. + +Integration with higher-level profilers would be valuable. Combining this tool with NVIDIA Nsight Systems would provide both high-level code flow (from flamegraphs) and detailed kernel metrics (from Nsight). Similarly, integrating with perf or BPF-based full-system profilers would show GPU work in the context of system-wide resource usage (CPU scheduling, interrupts, memory pressure). The folded stack format is designed for this - you can merge CPU perf samples with GPU samples by concatenating stacks. + +For truly unified CPU+GPU observability, explore eBPF programs running on the GPU itself. The [bpftime GPU project](https://github.com/eunomia-bpf/bpftime/tree/master/example/gpu) compiles eBPF bytecode to PTX instructions, enabling instrumentation inside GPU kernels. This exposes thread-level metrics like memory coalescing efficiency, warp occupancy, and bank conflicts - data impossible to obtain from kernel-side tracing. Future directions could combine kernel-side CUPTI tracing with GPU-side eBPF instrumentation for complete visibility from application code to individual warp execution. + +## Summary + +GPU profiling requires bridging two execution domains: CPU code submitting work and GPU hardware executing it. This tutorial demonstrated a complete profiling stack combining eBPF for CPU stack traces, CUPTI for GPU activity tracing, and correlation logic to merge them into unified flamegraphs. The eBPF profiler captures CPU stacks at every `cudaLaunchKernel` call with nanosecond timestamps. CUPTI injection records GPU kernel executions with correlation IDs linking them back to CPU API calls. The trace merger matches events via timestamps and correlation IDs, producing folded stacks showing complete execution paths from application code through CUDA runtime to GPU kernels. The resulting flamegraphs visualize end-to-end execution, revealing hotspots across both CPU and GPU. + +This approach works without recompiling applications, supports any CUDA framework (PyTorch, TensorFlow, JAX, raw CUDA), and provides low overhead suitable for production profiling. The tools are modular - you can use eBPF profiling alone for CPU analysis, CUPTI injection alone for GPU timelines, or combine them for unified visibility. Apply these techniques to diagnose performance bottlenecks in ML training, GPU-accelerated applications, or any CUDA workload where understanding CPU-GPU interaction is critical. + +> If you'd like to dive deeper into eBPF, check out our tutorial repository at or visit our website at . + +## Cross-Vendor GPU Profiling Comparison + +This tutorial focuses on NVIDIA GPUs using CUPTI, but the architecture applies across GPU vendors with vendor-specific APIs replacing CUPTI. Understanding these alternatives helps you apply similar techniques to other GPU platforms. + +**Intel GPUs (iaprof approach)**: Intel's profiling architecture uses Level Zero API for GPU tracing and Intel GPU Observability Architecture (OA) for hardware performance monitoring. Instead of CUPTI injection, iaprof uses eBPF tracepoints on the i915/Xe kernel drivers to intercept GPU command submission. The EU Stall Collector samples execution unit performance counters during kernel execution, capturing memory stalls, ALU bottlenecks, and instruction fetch delays. The Debug Collector retrieves shader binaries and context metadata through Intel Debug API. Correlation happens through batch buffer parsing - iaprof extracts kernel contexts from GPU command buffers and matches them to eBPF CPU stack traces via timestamp proximity. The deferred attribution model handles out-of-order hardware samples by buffering them until kernel execution completes, then matching samples to shader contexts using context IDs and timestamps. This is more complex than CUPTI correlation because GPU hardware doesn't provide correlation IDs directly. + +**AMD GPUs (ROCm approach)**: AMD's ROCm stack provides rocProfiler for GPU tracing, similar to CUPTI's role in CUDA. The rocProfiler API enables activity callbacks for kernel dispatches, memory transfers, and hardware performance counters. eBPF profiling on AMD GPUs can attach uprobes to `hipLaunchKernel` or `hsa_queue_create` in the ROCm runtime. The HSA (Heterogeneous System Architecture) runtime assigns correlation IDs to kernel dispatches, analogous to CUDA correlation IDs. AMD GPUs expose hardware counters through rocProfiler that reveal memory bandwidth utilization, wavefront occupancy, and cache hit rates. The correlation mechanism is similar to CUPTI - match eBPF CPU stacks to rocProfiler runtime API events by timestamp, then match runtime events to kernel executions via correlation IDs. + +**Vendor-Neutral Approaches**: For applications using portable GPU APIs like OpenCL or SYCL, profiling must work across vendors. OpenCL provides `clSetEventCallback()` for event notification and `clGetEventProfilingInfo()` for kernel timing. eBPF uprobes can attach to `clEnqueueNDRangeKernel` to capture CPU stacks at kernel submission. SYCL queue profiling captures kernel execution times through `info::event::command_start` and `command_end` queries. The challenge is that vendor-neutral APIs don't expose hardware performance counters uniformly - each vendor requires platform-specific extensions. + +The key insight: all GPU profiling follows the same pattern - capture CPU context at API submission (eBPF uprobes), trace GPU execution with vendor APIs (CUPTI/OA/rocProfiler), and correlate via timestamps and IDs. The iaprof architecture demonstrates advanced techniques like deferred attribution and batch buffer parsing that can enhance CUDA profiling for complex multi-stream workloads. Studying cross-vendor approaches reveals common challenges (asynchronous hardware samples, out-of-order execution, clock synchronization) and solutions (buffered correlation, epoch-based ID tracking, timestamp windows). + +## References + +- **NVIDIA CUPTI Documentation**: +- **CUPTI Activity API**: +- **CUDA Profiling Guide**: +- **eBPF Stack Trace Helpers**: +- **Chrome Trace Format**: +- **Flamegraph Visualization**: +- **bpftime GPU eBPF**: +- **iaprof Intel GPU Profiling**: +- **Intel GPU Observability Architecture**: Intel Graphics documentation +- **AMD ROCm Profiler**: +- **Tutorial Repository**: + +Complete source code including the eBPF profiler, CUPTI injection library, trace merger, and test applications is available in the tutorial repository. Contributions and issue reports welcome! diff --git a/src/xpu/flamegraph/combined_flamegraph.pl b/src/xpu/flamegraph/combined_flamegraph.pl new file mode 100755 index 0000000..fbe4b88 --- /dev/null +++ b/src/xpu/flamegraph/combined_flamegraph.pl @@ -0,0 +1,1315 @@ +#!/usr/bin/perl -w +# +# flamegraph.pl flame stack grapher. +# +# This takes stack samples and renders a call graph, allowing hot functions +# and codepaths to be quickly identified. Stack samples can be generated using +# tools such as DTrace, perf, SystemTap, and Instruments. +# +# USAGE: ./flamegraph.pl [options] input.txt > graph.svg +# +# grep funcA input.txt | ./flamegraph.pl [options] > graph.svg +# +# Then open the resulting .svg in a web browser, for interactivity: mouse-over +# frames for info, click to zoom, and ctrl-F to search. +# +# Options are listed in the usage message (--help). +# +# The input is stack frames and sample counts formatted as single lines. Each +# frame in the stack is semicolon separated, with a space and count at the end +# of the line. These can be generated for Linux perf script output using +# stackcollapse-perf.pl, for DTrace using stackcollapse.pl, and for other tools +# using the other stackcollapse programs. Example input: +# +# swapper;start_kernel;rest_init;cpu_idle;default_idle;native_safe_halt 1 +# +# An optional extra column of counts can be provided to generate a differential +# flame graph of the counts, colored red for more, and blue for less. This +# can be useful when using flame graphs for non-regression testing. +# See the header comment in the difffolded.pl program for instructions. +# +# The input functions can optionally have annotations at the end of each +# function name, following a precedent by some tools (Linux perf's _[k]): +# _[k] for kernel +# _[i] for inlined +# _[j] for jit +# _[w] for waker +# Some of the stackcollapse programs support adding these annotations, eg, +# stackcollapse-perf.pl --kernel --jit. They are used merely for colors by +# some palettes, eg, flamegraph.pl --color=java. +# +# The output flame graph shows relative presence of functions in stack samples. +# The ordering on the x-axis has no meaning; since the data is samples, time +# order of events is not known. The order used sorts function names +# alphabetically. +# +# While intended to process stack samples, this can also process stack traces. +# For example, tracing stacks for memory allocation, or resource usage. You +# can use --title to set the title to reflect the content, and --countname +# to change "samples" to "bytes" etc. +# +# There are a few different palettes, selectable using --color. By default, +# the colors are selected at random (except for differentials). Functions +# called "-" will be printed gray, which can be used for stack separators (eg, +# between user and kernel stacks). +# +# HISTORY +# +# This was inspired by Neelakanth Nadgir's excellent function_call_graph.rb +# program, which visualized function entry and return trace events. As Neel +# wrote: "The output displayed is inspired by Roch's CallStackAnalyzer which +# was in turn inspired by the work on vftrace by Jan Boerhout". See: +# https://blogs.oracle.com/realneel/entry/visualizing_callstacks_via_dtrace_and +# +# Copyright 2016 Netflix, Inc. +# Copyright 2011 Joyent, Inc. All rights reserved. +# Copyright 2011 Brendan Gregg. All rights reserved. +# +# CDDL HEADER START +# +# The contents of this file are subject to the terms of the +# Common Development and Distribution License (the "License"). +# You may not use this file except in compliance with the License. +# +# You can obtain a copy of the license at docs/cddl1.txt or +# http://opensource.org/licenses/CDDL-1.0. +# See the License for the specific language governing permissions +# and limitations under the License. +# +# When distributing Covered Code, include this CDDL HEADER in each +# file and include the License file at docs/cddl1.txt. +# If applicable, add the following below this CDDL HEADER, with the +# fields enclosed by brackets "[]" replaced with your own identifying +# information: Portions Copyright [yyyy] [name of copyright owner] +# +# CDDL HEADER END +# +# 11-Oct-2014 Adrien Mahieux Added zoom. +# 21-Nov-2013 Shawn Sterling Added consistent palette file option +# 17-Mar-2013 Tim Bunce Added options and more tunables. +# 15-Dec-2011 Dave Pacheco Support for frames with whitespace. +# 10-Sep-2011 Brendan Gregg Created this. + +use strict; + +use Getopt::Long; + +use open qw(:std :utf8); + +# tunables +my $encoding; +my $fonttype = "Verdana"; +my $imagewidth = 1200; # max width, pixels +my $frameheight = 16; # max height is dynamic +my $fontsize = 12; # base text size +my $fontwidth = 0.59; # avg width relative to fontsize +my $minwidth = 0.1; # min function width, pixels or percentage of time +my $nametype = "Function:"; # what are the names in the data? +my $countname = "samples"; # what are the counts in the data? +my $colors = "hot"; # color theme +my $bgcolors = ""; # background color theme +my $nameattrfile; # file holding function attributes +my $timemax; # (override the) sum of the counts +my $factor = 1; # factor to scale counts by +my $hash = 0; # color by function name +my $rand = 0; # color randomly +my $palette = 0; # if we use consistent palettes (default off) +my %palette_map; # palette map hash +my $pal_file = "palette.map"; # palette map file name +my $stackreverse = 0; # reverse stack order, switching merge end +my $inverted = 0; # icicle graph +my $flamechart = 0; # produce a flame chart (sort by time, do not merge stacks) +my $negate = 0; # switch differential hues +my $titletext = ""; # centered heading +my $titledefault = "Flame Graph"; # overwritten by --title +my $titleinverted = "Icicle Graph"; # " " +my $searchcolor = "rgb(230,0,230)"; # color for search highlighting +my $notestext = ""; # embedded notes in SVG +my $subtitletext = ""; # second level title (optional) +my $help = 0; + +sub usage { + die < outfile.svg\n + --title TEXT # change title text + --subtitle TEXT # second level title (optional) + --width NUM # width of image (default 1200) + --height NUM # height of each frame (default 16) + --minwidth NUM # omit smaller functions. In pixels or use "%" for + # percentage of time (default 0.1 pixels) + --fonttype FONT # font type (default "Verdana") + --fontsize NUM # font size (default 12) + --countname TEXT # count type label (default "samples") + --nametype TEXT # name type label (default "Function:") + --colors PALETTE # set color palette. choices are: hot (default), mem, + # io, wakeup, chain, java, js, perl, red, green, blue, + # aqua, yellow, purple, orange + --bgcolors COLOR # set background colors. gradient choices are yellow + # (default), blue, green, grey; flat colors use "#rrggbb" + --hash # colors are keyed by function name hash + --random # colors are randomly generated + --cp # use consistent palette (palette.map) + --reverse # generate stack-reversed flame graph + --inverted # icicle graph + --flamechart # produce a flame chart (sort by time, do not merge stacks) + --negate # switch differential hues (blue<->red) + --notes TEXT # add notes comment in SVG (for debugging) + --help # this message + + eg, + $0 --title="Flame Graph: malloc()" trace.txt > graph.svg +USAGE_END +} + +GetOptions( + 'fonttype=s' => \$fonttype, + 'width=i' => \$imagewidth, + 'height=i' => \$frameheight, + 'encoding=s' => \$encoding, + 'fontsize=f' => \$fontsize, + 'fontwidth=f' => \$fontwidth, + 'minwidth=s' => \$minwidth, + 'title=s' => \$titletext, + 'subtitle=s' => \$subtitletext, + 'nametype=s' => \$nametype, + 'countname=s' => \$countname, + 'nameattr=s' => \$nameattrfile, + 'total=s' => \$timemax, + 'factor=f' => \$factor, + 'colors=s' => \$colors, + 'bgcolors=s' => \$bgcolors, + 'hash' => \$hash, + 'random' => \$rand, + 'cp' => \$palette, + 'reverse' => \$stackreverse, + 'inverted' => \$inverted, + 'flamechart' => \$flamechart, + 'negate' => \$negate, + 'notes=s' => \$notestext, + 'help' => \$help, +) or usage(); +$help && usage(); + +# internals +my $ypad1 = $fontsize * 3; # pad top, include title +my $ypad2 = $fontsize * 2 + 10; # pad bottom, include labels +my $ypad3 = $fontsize * 2; # pad top, include subtitle (optional) +my $xpad = 10; # pad lefm and right +my $framepad = 1; # vertical padding for frames +my $depthmax = 0; +my %Events; +my %nameattr; + +if ($flamechart && $titletext eq "") { + $titletext = "Flame Chart"; +} + +if ($titletext eq "") { + unless ($inverted) { + $titletext = $titledefault; + } else { + $titletext = $titleinverted; + } +} + +if ($nameattrfile) { + # The name-attribute file format is a function name followed by a tab then + # a sequence of tab separated name=value pairs. + open my $attrfh, $nameattrfile or die "Can't read $nameattrfile: $!\n"; + while (<$attrfh>) { + chomp; + my ($funcname, $attrstr) = split /\t/, $_, 2; + die "Invalid format in $nameattrfile" unless defined $attrstr; + $nameattr{$funcname} = { map { split /=/, $_, 2 } split /\t/, $attrstr }; + } +} + +if ($notestext =~ /[<>]/) { + die "Notes string can't contain < or >" +} + +# Ensure minwidth is a valid floating-point number, +# print usage string if not +my $minwidth_f; +if ($minwidth =~ /^([0-9.]+)%?$/) { + $minwidth_f = $1; +} else { + warn "Value '$minwidth' is invalid for minwidth, expected a float.\n"; + usage(); +} + +# background colors: +# - yellow gradient: default (hot, java, js, perl) +# - green gradient: mem +# - blue gradient: io, wakeup, chain +# - gray gradient: flat colors (red, green, blue, ...) +if ($bgcolors eq "") { + # choose a default + if ($colors eq "mem") { + $bgcolors = "green"; + } elsif ($colors =~ /^(io|wakeup|chain)$/) { + $bgcolors = "blue"; + } elsif ($colors =~ /^(red|green|blue|aqua|yellow|purple|orange)$/) { + $bgcolors = "grey"; + } else { + $bgcolors = "yellow"; + } +} +my ($bgcolor1, $bgcolor2); +if ($bgcolors eq "yellow") { + $bgcolor1 = "#eeeeee"; # background color gradient start + $bgcolor2 = "#eeeeb0"; # background color gradient stop +} elsif ($bgcolors eq "blue") { + $bgcolor1 = "#eeeeee"; $bgcolor2 = "#e0e0ff"; +} elsif ($bgcolors eq "green") { + $bgcolor1 = "#eef2ee"; $bgcolor2 = "#e0ffe0"; +} elsif ($bgcolors eq "grey") { + $bgcolor1 = "#f8f8f8"; $bgcolor2 = "#e8e8e8"; +} elsif ($bgcolors =~ /^#......$/) { + $bgcolor1 = $bgcolor2 = $bgcolors; +} else { + die "Unrecognized bgcolor option \"$bgcolors\"" +} + +# SVG functions +{ package SVG; + sub new { + my $class = shift; + my $self = {}; + bless ($self, $class); + return $self; + } + + sub header { + my ($self, $w, $h) = @_; + my $enc_attr = ''; + if (defined $encoding) { + $enc_attr = qq{ encoding="$encoding"}; + } + $self->{svg} .= < + + + + +SVG + } + + sub include { + my ($self, $content) = @_; + $self->{svg} .= $content; + } + + sub colorAllocate { + my ($self, $r, $g, $b) = @_; + return "rgb($r,$g,$b)"; + } + + sub group_start { + my ($self, $attr) = @_; + + my @g_attr = map { + exists $attr->{$_} ? sprintf(qq/$_="%s"/, $attr->{$_}) : () + } qw(id class); + push @g_attr, $attr->{g_extra} if $attr->{g_extra}; + if ($attr->{href}) { + my @a_attr; + push @a_attr, sprintf qq/xlink:href="%s"/, $attr->{href} if $attr->{href}; + # default target=_top else links will open within SVG + push @a_attr, sprintf qq/target="%s"/, $attr->{target} || "_top"; + push @a_attr, $attr->{a_extra} if $attr->{a_extra}; + $self->{svg} .= sprintf qq/\n/, join(' ', (@a_attr, @g_attr)); + } else { + $self->{svg} .= sprintf qq/\n/, join(' ', @g_attr); + } + + $self->{svg} .= sprintf qq/%s<\/title>/, $attr->{title} + if $attr->{title}; # should be first element within g container + } + + sub group_end { + my ($self, $attr) = @_; + $self->{svg} .= $attr->{href} ? qq/<\/a>\n/ : qq/<\/g>\n/; + } + + sub filledRectangle { + my ($self, $x1, $y1, $x2, $y2, $fill, $extra) = @_; + $x1 = sprintf "%0.1f", $x1; + $x2 = sprintf "%0.1f", $x2; + my $w = sprintf "%0.1f", $x2 - $x1; + my $h = sprintf "%0.1f", $y2 - $y1; + $extra = defined $extra ? $extra : ""; + $self->{svg} .= qq/\n/; + } + + sub stringTTF { + my ($self, $id, $x, $y, $str, $extra) = @_; + $x = sprintf "%0.2f", $x; + $id = defined $id ? qq/id="$id"/ : ""; + $extra ||= ""; + $self->{svg} .= qq/$str<\/text>\n/; + } + + sub svg { + my $self = shift; + return "$self->{svg}\n"; + } + 1; +} + +sub namehash { + # Generate a vector hash for the name string, weighting early over + # later characters. We want to pick the same colors for function + # names across different flame graphs. + my $name = shift; + my $vector = 0; + my $weight = 1; + my $max = 1; + my $mod = 10; + # if module name present, trunc to 1st char + $name =~ s/.(.*?)`//; + foreach my $c (split //, $name) { + my $i = (ord $c) % $mod; + $vector += ($i / ($mod++ - 1)) * $weight; + $max += 1 * $weight; + $weight *= 0.70; + last if $mod > 12; + } + return (1 - $vector / $max) +} + +sub sum_namehash { + my $name = shift; + return unpack("%32W*", $name); +} + +sub random_namehash { + # Generate a random hash for the name string. + # This ensures that functions with the same name have the same color, + # both within a flamegraph and across multiple flamegraphs without + # needing to set a palette and while preserving the original flamegraph + # optic, unlike what happens with --hash. + my $name = shift; + my $hash = sum_namehash($name); + srand($hash); + return rand(1) +} + +sub color { + my ($type, $hash, $name) = @_; + my ($v1, $v2, $v3); + + if ($hash) { + $v1 = namehash($name); + $v2 = $v3 = namehash(scalar reverse $name); + } elsif ($rand) { + $v1 = rand(1); + $v2 = rand(1); + $v3 = rand(1); + } else { + $v1 = random_namehash($name); + $v2 = random_namehash($name); + $v3 = random_namehash($name); + } + + # theme palettes + if (defined $type and $type eq "hot") { + my $r = 205 + int(50 * $v3); + my $g = 0 + int(230 * $v1); + my $b = 0 + int(55 * $v2); + return "rgb($r,$g,$b)"; + } + if (defined $type and $type eq "mem") { + my $r = 0; + my $g = 190 + int(50 * $v2); + my $b = 0 + int(210 * $v1); + return "rgb($r,$g,$b)"; + } + if (defined $type and $type eq "io") { + my $r = 80 + int(60 * $v1); + my $g = $r; + my $b = 190 + int(55 * $v2); + return "rgb($r,$g,$b)"; + } + + # multi palettes + if (defined $type and $type eq "java") { + # Handle both annotations (_[j], _[i], ...; which are + # accurate), as well as input that lacks any annotations, as + # best as possible. Without annotations, we get a little hacky + # and match on java|org|com, etc. + if ($name =~ m:_\[j\]$:) { # jit annotation + $type = "green"; + } elsif ($name =~ m:_\[i\]$:) { # inline annotation + $type = "aqua"; + } elsif ($name =~ m:^L?(java|javax|jdk|net|org|com|io|sun)/:) { # Java + $type = "green"; + } elsif ($name =~ /:::/) { # Java, typical perf-map-agent method separator + $type = "green"; + } elsif ($name =~ /::/) { # C++ + $type = "yellow"; + } elsif ($name =~ m:_\[k\]$:) { # kernel annotation + $type = "orange"; + } elsif ($name =~ /::/) { # C++ + $type = "yellow"; + } else { # system + $type = "red"; + } + # fall-through to color palettes + } + if (defined $type and $type eq "perl") { + if ($name =~ /::/) { # C++ + $type = "yellow"; + } elsif ($name =~ m:Perl: or $name =~ m:\.pl:) { # Perl + $type = "green"; + } elsif ($name =~ m:_\[k\]$:) { # kernel + $type = "orange"; + } else { # system + $type = "red"; + } + # fall-through to color palettes + } + if (defined $type and $type eq "js") { + # Handle both annotations (_[j], _[i], ...; which are + # accurate), as well as input that lacks any annotations, as + # best as possible. Without annotations, we get a little hacky, + # and match on a "/" with a ".js", etc. + if ($name =~ m:_\[j\]$:) { # jit annotation + if ($name =~ m:/:) { + $type = "green"; # source + } else { + $type = "aqua"; # builtin + } + } elsif ($name =~ /::/) { # C++ + $type = "yellow"; + } elsif ($name =~ m:/.*\.js:) { # JavaScript (match "/" in path) + $type = "green"; + } elsif ($name =~ m/:/) { # JavaScript (match ":" in builtin) + $type = "aqua"; + } elsif ($name =~ m/^ $/) { # Missing symbol + $type = "green"; + } elsif ($name =~ m:_\[k\]:) { # kernel + $type = "orange"; + } else { # system + $type = "red"; + } + # fall-through to color palettes + } + if (defined $type and $type eq "wakeup") { + $type = "aqua"; + # fall-through to color palettes + } + if (defined $type and $type eq "chain") { + if ($name =~ m:_\[w\]:) { # waker + $type = "aqua" + } else { # off-CPU + $type = "blue"; + } + # fall-through to color palettes + } + + + if (defined $type and $type eq "combined") { + if ($name =~ m:_\[c\]$:) { # CPU annotation (on-CPU) + $type = "red"; + } elsif ($name =~ m:_\[o\]$:) { # off-CPU annotation (I/O/blocking) + $type = "blue"; + } else { # default + $type = "yellow"; + } + # fall-through to color palettes + } + + # color palettes + if (defined $type and $type eq "red") { + my $r = 200 + int(55 * $v1); + my $x = 50 + int(80 * $v1); + return "rgb($r,$x,$x)"; + } + if (defined $type and $type eq "green") { + my $g = 200 + int(55 * $v1); + my $x = 50 + int(60 * $v1); + return "rgb($x,$g,$x)"; + } + if (defined $type and $type eq "blue") { + my $b = 205 + int(50 * $v1); + my $x = 80 + int(60 * $v1); + return "rgb($x,$x,$b)"; + } + if (defined $type and $type eq "yellow") { + my $x = 175 + int(55 * $v1); + my $b = 50 + int(20 * $v1); + return "rgb($x,$x,$b)"; + } + if (defined $type and $type eq "purple") { + my $x = 190 + int(65 * $v1); + my $g = 80 + int(60 * $v1); + return "rgb($x,$g,$x)"; + } + if (defined $type and $type eq "aqua") { + my $r = 50 + int(60 * $v1); + my $g = 165 + int(55 * $v1); + my $b = 165 + int(55 * $v1); + return "rgb($r,$g,$b)"; + } + if (defined $type and $type eq "orange") { + my $r = 190 + int(65 * $v1); + my $g = 90 + int(65 * $v1); + return "rgb($r,$g,0)"; + } + + return "rgb(0,0,0)"; +} + +sub color_scale { + my ($value, $max) = @_; + my ($r, $g, $b) = (255, 255, 255); + $value = -$value if $negate; + if ($value > 0) { + $g = $b = int(210 * ($max - $value) / $max); + } elsif ($value < 0) { + $r = $g = int(210 * ($max + $value) / $max); + } + return "rgb($r,$g,$b)"; +} + +sub color_map { + my ($colors, $func) = @_; + if (exists $palette_map{$func}) { + return $palette_map{$func}; + } else { + $palette_map{$func} = color($colors, $hash, $func); + return $palette_map{$func}; + } +} + +sub write_palette { + open(FILE, ">$pal_file"); + foreach my $key (sort keys %palette_map) { + print FILE $key."->".$palette_map{$key}."\n"; + } + close(FILE); +} + +sub read_palette { + if (-e $pal_file) { + open(FILE, $pal_file) or die "can't open file $pal_file: $!"; + while ( my $line = ) { + chomp($line); + (my $key, my $value) = split("->",$line); + $palette_map{$key}=$value; + } + close(FILE) + } +} + +my %Node; # Hash of merged frame data +my %Tmp; + +# flow() merges two stacks, storing the merged frames and value data in %Node. +sub flow { + my ($last, $this, $v, $d) = @_; + + my $len_a = @$last - 1; + my $len_b = @$this - 1; + + my $i = 0; + my $len_same; + for (; $i <= $len_a; $i++) { + last if $i > $len_b; + last if $last->[$i] ne $this->[$i]; + } + $len_same = $i; + + for ($i = $len_a; $i >= $len_same; $i--) { + my $k = "$last->[$i];$i"; + # a unique ID is constructed from "func;depth;etime"; + # func-depth isn't unique, it may be repeated later. + $Node{"$k;$v"}->{stime} = delete $Tmp{$k}->{stime}; + if (defined $Tmp{$k}->{delta}) { + $Node{"$k;$v"}->{delta} = delete $Tmp{$k}->{delta}; + } + delete $Tmp{$k}; + } + + for ($i = $len_same; $i <= $len_b; $i++) { + my $k = "$this->[$i];$i"; + $Tmp{$k}->{stime} = $v; + if (defined $d) { + $Tmp{$k}->{delta} += $i == $len_b ? $d : 0; + } + } + + return $this; +} + +# parse input +my @Data; +my @SortedData; +my $last = []; +my $time = 0; +my $delta = undef; +my $ignored = 0; +my $line; +my $maxdelta = 1; + +# reverse if needed +foreach (<>) { + chomp; + $line = $_; + if ($stackreverse) { + # there may be an extra samples column for differentials + # XXX todo: redo these REs as one. It's repeated below. + my($stack, $samples) = (/^(.*)\s+?(\d+(?:\.\d*)?)$/); + my $samples2 = undef; + if ($stack =~ /^(.*)\s+?(\d+(?:\.\d*)?)$/) { + $samples2 = $samples; + ($stack, $samples) = $stack =~ (/^(.*)\s+?(\d+(?:\.\d*)?)$/); + unshift @Data, join(";", reverse split(";", $stack)) . " $samples $samples2"; + } else { + unshift @Data, join(";", reverse split(";", $stack)) . " $samples"; + } + } else { + unshift @Data, $line; + } +} + +if ($flamechart) { + # In flame chart mode, just reverse the data so time moves from left to right. + @SortedData = reverse @Data; +} else { + @SortedData = sort @Data; +} + +# process and merge frames +foreach (@SortedData) { + chomp; + # process: folded_stack count + # eg: func_a;func_b;func_c 31 + my ($stack, $samples) = (/^(.*)\s+?(\d+(?:\.\d*)?)$/); + unless (defined $samples and defined $stack) { + ++$ignored; + next; + } + + # there may be an extra samples column for differentials: + my $samples2 = undef; + if ($stack =~ /^(.*)\s+?(\d+(?:\.\d*)?)$/) { + $samples2 = $samples; + ($stack, $samples) = $stack =~ (/^(.*)\s+?(\d+(?:\.\d*)?)$/); + } + $delta = undef; + if (defined $samples2) { + $delta = $samples2 - $samples; + $maxdelta = abs($delta) if abs($delta) > $maxdelta; + } + + # for chain graphs, annotate waker frames with "_[w]", for later + # coloring. This is a hack, but has a precedent ("_[k]" from perf). + if ($colors eq "chain") { + my @parts = split ";--;", $stack; + my @newparts = (); + $stack = shift @parts; + $stack .= ";--;"; + foreach my $part (@parts) { + $part =~ s/;/_[w];/g; + $part .= "_[w]"; + push @newparts, $part; + } + $stack .= join ";--;", @parts; + } + + # merge frames and populate %Node: + $last = flow($last, [ '', split ";", $stack ], $time, $delta); + + if (defined $samples2) { + $time += $samples2; + } else { + $time += $samples; + } +} +flow($last, [], $time, $delta); + +if ($countname eq "samples") { + # If $countname is used, it's likely that we're not measuring in stack samples + # (e.g. time could be the unit), so don't warn. + warn "Stack count is low ($time). Did something go wrong?\n" if $time < 100; +} + +warn "Ignored $ignored lines with invalid format\n" if $ignored; +unless ($time) { + warn "ERROR: No stack counts found\n"; + my $im = SVG->new(); + # emit an error message SVG, for tools automating flamegraph use + my $imageheight = $fontsize * 5; + $im->header($imagewidth, $imageheight); + $im->stringTTF(undef, int($imagewidth / 2), $fontsize * 2, + "ERROR: No valid input provided to flamegraph.pl."); + print $im->svg; + exit 2; +} +if ($timemax and $timemax < $time) { + warn "Specified --total $timemax is less than actual total $time, so ignored\n" + if $timemax/$time > 0.02; # only warn is significant (e.g., not rounding etc) + undef $timemax; +} +$timemax ||= $time; + +my $widthpertime = ($imagewidth - 2 * $xpad) / $timemax; + +# Treat as a percentage of time if the string ends in a "%". +my $minwidth_time; +if ($minwidth =~ /%$/) { + $minwidth_time = $timemax * $minwidth_f / 100; +} else { + $minwidth_time = $minwidth_f / $widthpertime; +} + +# prune blocks that are too narrow and determine max depth +while (my ($id, $node) = each %Node) { + my ($func, $depth, $etime) = split ";", $id; + my $stime = $node->{stime}; + die "missing start for $id" if not defined $stime; + + if (($etime-$stime) < $minwidth_time) { + delete $Node{$id}; + next; + } + $depthmax = $depth if $depth > $depthmax; +} + +# draw canvas, and embed interactive JavaScript program +my $imageheight = (($depthmax + 1) * $frameheight) + $ypad1 + $ypad2; +$imageheight += $ypad3 if $subtitletext ne ""; +my $titlesize = $fontsize + 5; +my $im = SVG->new(); +my ($black, $vdgrey, $dgrey) = ( + $im->colorAllocate(0, 0, 0), + $im->colorAllocate(160, 160, 160), + $im->colorAllocate(200, 200, 200), + ); +$im->header($imagewidth, $imageheight); +my $inc = < + + + + + + + +INC +$im->include($inc); +$im->filledRectangle(0, 0, $imagewidth, $imageheight, 'url(#background)'); +$im->stringTTF("title", int($imagewidth / 2), $fontsize * 2, $titletext); +$im->stringTTF("subtitle", int($imagewidth / 2), $fontsize * 4, $subtitletext) if $subtitletext ne ""; +$im->stringTTF("details", $xpad, $imageheight - ($ypad2 / 2), " "); +$im->stringTTF("unzoom", $xpad, $fontsize * 2, "Reset Zoom", 'class="hide"'); +$im->stringTTF("search", $imagewidth - $xpad - 100, $fontsize * 2, "Search"); +$im->stringTTF("ignorecase", $imagewidth - $xpad - 16, $fontsize * 2, "ic"); +$im->stringTTF("matched", $imagewidth - $xpad - 100, $imageheight - ($ypad2 / 2), " "); + +if ($palette) { + read_palette(); +} + +# draw frames +$im->group_start({id => "frames"}); +while (my ($id, $node) = each %Node) { + my ($func, $depth, $etime) = split ";", $id; + my $stime = $node->{stime}; + my $delta = $node->{delta}; + + $etime = $timemax if $func eq "" and $depth == 0; + + my $x1 = $xpad + $stime * $widthpertime; + my $x2 = $xpad + $etime * $widthpertime; + my ($y1, $y2); + unless ($inverted) { + $y1 = $imageheight - $ypad2 - ($depth + 1) * $frameheight + $framepad; + $y2 = $imageheight - $ypad2 - $depth * $frameheight; + } else { + $y1 = $ypad1 + $depth * $frameheight; + $y2 = $ypad1 + ($depth + 1) * $frameheight - $framepad; + } + + # Add commas per perlfaq5: + # https://perldoc.perl.org/perlfaq5#How-can-I-output-my-numbers-with-commas-added? + my $samples = sprintf "%.0f", ($etime - $stime) * $factor; + (my $samples_txt = $samples) + =~ s/(^[-+]?\d+?(?=(?>(?:\d{3})+)(?!\d))|\G\d{3}(?=\d))/$1,/g; + + my $info; + if ($func eq "" and $depth == 0) { + $info = "all ($samples_txt $countname, 100%)"; + } else { + my $pct = sprintf "%.2f", ((100 * $samples) / ($timemax * $factor)); + my $escaped_func = $func; + # clean up SVG breaking characters: + $escaped_func =~ s/&/&/g; + $escaped_func =~ s//>/g; + $escaped_func =~ s/"/"/g; + $escaped_func =~ s/_\[[kwij]\]$//; # strip any annotation + unless (defined $delta) { + $info = "$escaped_func ($samples_txt $countname, $pct%)"; + } else { + my $d = $negate ? -$delta : $delta; + my $deltapct = sprintf "%.2f", ((100 * $d) / ($timemax * $factor)); + $deltapct = $d > 0 ? "+$deltapct" : $deltapct; + $info = "$escaped_func ($samples_txt $countname, $pct%; $deltapct%)"; + } + } + + my $nameattr = { %{ $nameattr{$func}||{} } }; # shallow clone + $nameattr->{title} ||= $info; + $im->group_start($nameattr); + + my $color; + if ($func eq "--") { + $color = $vdgrey; + } elsif ($func eq "-") { + $color = $dgrey; + } elsif (defined $delta) { + $color = color_scale($delta, $maxdelta); + } elsif ($palette) { + $color = color_map($colors, $func); + } else { + $color = color($colors, $hash, $func); + } + $im->filledRectangle($x1, $y1, $x2, $y2, $color, 'rx="2" ry="2"'); + + my $chars = int( ($x2 - $x1) / ($fontsize * $fontwidth)); + my $text = ""; + if ($chars >= 3) { # room for one char plus two dots + $func =~ s/_\[[kwij]\]$//; # strip any annotation + $text = substr $func, 0, $chars; + substr($text, -2, 2) = ".." if $chars < length $func; + $text =~ s/&/&/g; + $text =~ s//>/g; + } + $im->stringTTF(undef, $x1 + 3, 3 + ($y1 + $y2) / 2, $text); + + $im->group_end($nameattr); +} +$im->group_end(); + +print $im->svg; + +if ($palette) { + write_palette(); +} + +# vim: ts=8 sts=8 sw=8 noexpandtab diff --git a/src/xpu/flamegraph/cupti_trace/.gitignore b/src/xpu/flamegraph/cupti_trace/.gitignore new file mode 100644 index 0000000..67cd698 --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace/.gitignore @@ -0,0 +1,12 @@ +*.o +*.so +cpu_results.txt +gpu_results.txt +gpu_results.json +__pycache__/ +*.svg +*.folded +*.txt +/*.json +test_cupti +venv/ diff --git a/src/xpu/flamegraph/cupti_trace/Makefile b/src/xpu/flamegraph/cupti_trace/Makefile new file mode 100644 index 0000000..01fb4cd --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace/Makefile @@ -0,0 +1,71 @@ +# +# Copyright 2021 NVIDIA Corporation. All rights reserved +# +ifndef OS + OS := $(shell uname) + HOST_ARCH := $(shell uname -m) +endif + +CUDA_INSTALL_PATH ?= /usr/local/cuda-12.9 +NVCC := "$(CUDA_INSTALL_PATH)/bin/nvcc" +CUPTI_INSTALL_PATH ?= $(CUDA_INSTALL_PATH)/targets/x86_64-linux +CUPTI_SAMPLES_PATH ?= $(CUDA_INSTALL_PATH)/extras/CUPTI/samples/common +INCLUDES := -I"$(CUDA_INSTALL_PATH)/include" -I$(CUPTI_INSTALL_PATH)/include -I$(CUPTI_SAMPLES_PATH) + +ifeq ($(OS),Windows_NT) + LIB_PATH ?= ..\..\lib64 +else + EXTRAS_LIB_PATH := $(CUPTI_INSTALL_PATH)/lib + LIB_PATH ?= $(CUDA_INSTALL_PATH)/lib64 +endif + +# Point to the necessary cross-compiler. +NVCCFLAGS := +ifeq ($(OS),Windows_NT) + export PATH := $(PATH):$(LIB_PATH) + LIBS= -L $(LIB_PATH) -lcuda -lcupti -ldetours + LIBNAME := libcupti_trace_injection.dll +else + ifeq ($(OS), Darwin) + export DYLD_LIBRARY_PATH := $(DYLD_LIBRARY_PATH):$(LIB_PATH) + LIBS= -Xlinker -framework -Xlinker cuda -L $(EXTRAS_LIB_PATH) -L $(LIB_PATH) -lcupti + else + export LD_LIBRARY_PATH := $(LD_LIBRARY_PATH):$(LIB_PATH) + LIBS = -L $(LIB_PATH) -lcuda -L $(EXTRAS_LIB_PATH) -lcupti + endif + LIBNAME := libcupti_trace_injection.so + NVCCFLAGS += -Xcompiler -fPIC +endif + +ifneq ($(TARGET_ARCH), $(HOST_ARCH)) + ifeq ($(TARGET_ARCH), aarch64) + ifeq ($(TARGET_OS), linux) + HOST_COMPILER ?= aarch64-linux-gnu-g++ + else ifeq ($(TARGET_OS),qnx) + ifeq ($(QNX_HOST),) + $(error ERROR - QNX_HOST must be passed to the QNX host toolchain) + endif + ifeq ($(QNX_TARGET),) + $(error ERROR - QNX_TARGET must be passed to the QNX target toolchain) + endif + HOST_COMPILER ?= $(QNX_HOST)/usr/bin/q++ + ifndef QPP_CONFIG_VERSION + QPP_CONFIG_VERSION = 12.2.0 + endif + $(info QPP_CONFIG_VERSION = $(QPP_CONFIG_VERSION)) + NVCCFLAGS += --qpp-config $(QPP_CONFIG_VERSION),gcc_ntoaarch64le -lsocket + endif + endif + + ifdef HOST_COMPILER + NVCC_COMPILER := -ccbin $(HOST_COMPILER) + endif +endif + +all: cupti_trace_injection + +cupti_trace_injection: cupti_trace_injection.cpp + $(NVCC) $(NVCC_COMPILER) $(NVCCFLAGS) $(INCLUDES) -o $(LIBNAME) -shared $< $(LIBS) --no-device-link + +clean: + rm -f $(LIBNAME) cupti_trace_injection.o *.o *.bak diff --git a/src/xpu/flamegraph/cupti_trace/cupti_trace_injection.cpp b/src/xpu/flamegraph/cupti_trace/cupti_trace_injection.cpp new file mode 100644 index 0000000..0d744a7 --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace/cupti_trace_injection.cpp @@ -0,0 +1,593 @@ +/* + * Copyright 2021-2024 NVIDIA Corporation. All rights reserved. + * + * CUPTI based tracing injection to trace any CUDA application. + * This sample demonstrates how to use activity + * and callback APIs in the injection code. + * Refer to the README.txt file for usage. + * + * Workflow in brief: + * + * After the initialization routine returns, the application resumes running, + * with the registered callbacks triggering as expected. + * Subscribed to ProfilerStart and ProfilerStop callbacks. These callbacks + * control the collection of profiling data. + * + * ProfilerStart callback: + * Start the collection by enabling activities. Also enable callback for + * the API cudaDeviceReset to flush activity buffers. + * + * ProfilerStop callback: + * Get all the activity buffers which have all the activity records completed + * by using cuptiActivityFlushAll() API and then disable cudaDeviceReset callback + * and all the activities to stop collection. + * + * AtExitHandler: + * Register to the atexit handler to get all the activity buffers including the ones + * which have incomplete activity records by using force flush API + * cuptiActivityFlushAll(1). + */ + +// System headers +#include +#include +#include +#include +#include + +// CUDA headers +#include + +// CUPTI headers +#include "helper_cupti_activity.h" + +// Detours for Windows +#ifdef _WIN32 +#include "detours.h" +#include +#else +#include +#include +#endif + +// Macros +#define IS_ACTIVITY_SELECTED(activitySelect, activityKind) \ + (activitySelect & (1LL << activityKind)) + +#define SELECT_ACTIVITY(activitySelect, activityKind) \ + (activitySelect |= (1LL << activityKind)) + +// Variable related to initialize injection. +std::mutex initializeInjectionMutex; + +// Data structures for graph node tracking +typedef struct ApiData_st +{ + const char *pFunctionName; + uint32_t correlationId; +} ApiData; + +typedef std::map NodeIdApiDataMap; +NodeIdApiDataMap nodeIdCorrelationMap; + +// Global Structure +typedef struct InjectionGlobals_st +{ + volatile uint32_t initialized; + CUpti_SubscriberHandle subscriberHandle; + int tracingEnabled; + uint64_t profileMode; +} InjectionGlobals; + +InjectionGlobals injectionGlobals; + +CUptiResult +DisableCuptiActivities( + CUcontext ctx); + +// Functions +static void +InitializeInjectionGlobals(void) +{ + injectionGlobals.initialized = 0; + injectionGlobals.subscriberHandle = NULL; + injectionGlobals.tracingEnabled = 0; + injectionGlobals.profileMode = 0; +} + +static void +AtExitHandler(void) +{ + CUPTI_API_CALL(cuptiGetLastError()); + + // Force flush the activity buffers. + if (injectionGlobals.tracingEnabled) + { + CUPTI_API_CALL(DisableCuptiActivities(NULL)); + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(1)); + } + + // Flush and close output file if it's not stdout + if (globals.pOutputFile && globals.pOutputFile != stdout && globals.pOutputFile != stderr) + { + fflush(globals.pOutputFile); + fclose(globals.pOutputFile); + globals.pOutputFile = NULL; + } +} + +#ifdef _WIN32 +typedef void(WINAPI *rtlExitUserProcess_t)(uint32_t exitCode); +rtlExitUserProcess_t Real_RtlExitUserProcess = NULL; + +// Detour_RtlExitUserProcess. +void WINAPI +Detour_RtlExitUserProcess( + uint32_t exitCode) +{ + AtExitHandler(); + + Real_RtlExitUserProcess(exitCode); +} +#endif + +void +RegisterAtExitHandler(void) +{ +#ifdef _WIN32 + { + // It's unsafe to use atexit(), static destructors, DllMain PROCESS_DETACH, etc. + // because there's no way to guarantee the CUDA driver is still in a valid state + // when you get to those, due to the undefined order of dynamic library tear-down + // during process destruction. + // Also, the first thing the Windows kernel does when any thread in a process + // calls exit() is to immediately terminate all other threads, without any kind of + // synchronization. + // So the only valid time to do any in-process cleanup at exit() is before control + // is passed to the kernel. Use Detours to intercept a low-level ntdll.dll + // function "RtlExitUserProcess". + int detourStatus = 0; + FARPROC proc; + + // ntdll.dll will always be loaded, no need to load the library. + HMODULE ntDll = GetModuleHandle(TEXT("ntdll.dll")); + if (!ntDll) + { + detourStatus = 1; + goto DetourError; + } + + proc = GetProcAddress(ntDll, "RtlExitUserProcess"); + if (!proc) + { + detourStatus = 1; + goto DetourError; + } + Real_RtlExitUserProcess = (rtlExitUserProcess_t)proc; + + // Begin a detour transaction + if (DetourTransactionBegin() != ERROR_SUCCESS) + { + detourStatus = 1; + goto DetourError; + } + + if (DetourUpdateThread(GetCurrentThread()) != ERROR_SUCCESS) + { + detourStatus = 1; + goto DetourError; + } + + DetourSetIgnoreTooSmall(TRUE); + + if (DetourAttach((void **)&Real_RtlExitUserProcess, + (void *)Detour_RtlExitUserProcess) != ERROR_SUCCESS) + { + detourStatus = 1; + goto DetourError; + } + + // Commit the transaction + if (DetourTransactionCommit() != ERROR_SUCCESS) + { + detourStatus = 1; + goto DetourError; + } + DetourError: + if (detourStatus != 0) + { + atexit(&AtExitHandler); + } + } +#else + atexit(&AtExitHandler); +#endif +} + +static CUptiResult +SelectActivities() +{ + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_DRIVER); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_RUNTIME); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_OVERHEAD); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MEMSET); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MEMCPY); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MEMCPY2); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MEMORY2); + // Enable activities to capture the NVTX annotations - markers, ranges and resource naming. + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_NAME); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MARKER); + SELECT_ACTIVITY(injectionGlobals.profileMode, CUPTI_ACTIVITY_KIND_MARKER_DATA); + + return CUPTI_SUCCESS; +} + +void +GraphTraceRecords( + CUpti_Activity *pRecord) +{ + switch (pRecord->kind) + { + case CUPTI_ACTIVITY_KIND_MEMCPY: + { + CUpti_ActivityMemcpy6 *pMemcpyRecord = (CUpti_ActivityMemcpy6 *) pRecord; + + // Retrieve the information of the API used to create the node. + NodeIdApiDataMap::iterator it = nodeIdCorrelationMap.find(pMemcpyRecord->graphNodeId); + if (it != nodeIdCorrelationMap.end()) + { + fprintf(globals.pOutputFile, "Graph node was created using API %s with correlationId %u\n", + it->second.pFunctionName, it->second.correlationId); + } + break; + } + case CUPTI_ACTIVITY_KIND_KERNEL: + case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: + { + CUpti_ActivityKernel9 *pKernelRecord = (CUpti_ActivityKernel9 *) pRecord; + + // Retrieve the information of the API used to create the node. + NodeIdApiDataMap::iterator it = nodeIdCorrelationMap.find(pKernelRecord->graphNodeId); + if (it != nodeIdCorrelationMap.end()) + { + fprintf(globals.pOutputFile, "Graph node was created using API %s with correlationId %u\n", + it->second.pFunctionName, it->second.correlationId); + } + break; + } + default: + break; + } +} + +static CUptiResult +EnableCuptiActivities( + CUcontext context) +{ + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_RUNTIME_API, CUPTI_RUNTIME_TRACE_CBID_cudaDeviceReset_v3020)); + + CUPTI_API_CALL(SelectActivities()); + + for (int i = 0; i < CUPTI_ACTIVITY_KIND_COUNT; ++i) + { + CUptiResult result = CUPTI_SUCCESS; + + if (IS_ACTIVITY_SELECTED(injectionGlobals.profileMode, i)) + { + // If context is NULL activities are being enabled after CUDA initialization. + // Else the activities are being enabled on cudaProfilerStart API. + if (context == NULL) + { + std::cout << "Enabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << ".\n"; + CUPTI_API_CALL(cuptiActivityEnable((CUpti_ActivityKind)i)); + } + else + { + // Since some activities are not supported at context mode, + // enable them in global mode if context mode fails. + std::cout << "Enabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << " for a context.\n"; + result = cuptiActivityEnableContext(context, (CUpti_ActivityKind)i); + + if (result == CUPTI_ERROR_INVALID_KIND) + { + cuptiGetLastError(); + std::cout << "Enabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << ".\n"; + CUPTI_API_CALL_VERBOSE(cuptiActivityEnable((CUpti_ActivityKind)i)); + } + else if (result != CUPTI_SUCCESS) + { + CUPTI_API_CALL(result); + } + } + } + } + + return CUPTI_SUCCESS; +} + +CUptiResult +DisableCuptiActivities( + CUcontext context) +{ + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(0, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_RUNTIME_API, CUPTI_RUNTIME_TRACE_CBID_cudaDeviceReset_v3020)); + + for (int i = 0; i < CUPTI_ACTIVITY_KIND_COUNT; ++i) + { + CUptiResult result = CUPTI_SUCCESS; + + if (IS_ACTIVITY_SELECTED(injectionGlobals.profileMode, i)) + { + if (context == NULL) + { + std::cout << "Disabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << ".\n"; + CUPTI_API_CALL(cuptiActivityDisable((CUpti_ActivityKind)i)); + } + else + { + // Since some activities are not supported at context mode, + // disable them in global mode if context mode fails. + std::cout << "Disabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << " for a context.\n"; + result = cuptiActivityDisableContext(context, (CUpti_ActivityKind)i); + + if (result == CUPTI_ERROR_INVALID_KIND) + { + cuptiGetLastError(); + std::cout << "Disabling CUPTI_ACTIVITY_KIND_" << GetActivityKindString((CUpti_ActivityKind)i) << ".\n"; + CUPTI_API_CALL(cuptiActivityDisable((CUpti_ActivityKind)i)); + } + else if (result != CUPTI_SUCCESS) + { + CUPTI_API_CALL(result); + } + } + } + } + + return CUPTI_SUCCESS; +} + +static CUptiResult +OnCudaDeviceReset(void) +{ + // Flush all activity buffers. + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(0)); + + return CUPTI_SUCCESS; +} + +static CUptiResult +OnProfilerStart( + CUcontext context) +{ + if (context == NULL) + { + // Don't do anything if context is NULL. + return CUPTI_SUCCESS; + } + + CUPTI_API_CALL(EnableCuptiActivities(context)); + + return CUPTI_SUCCESS; +} + +static CUptiResult +OnProfilerStop( + CUcontext context) +{ + if (context == NULL) + { + // Don't do anything if context is NULL. + return CUPTI_SUCCESS; + } + + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(0)); + CUPTI_API_CALL(DisableCuptiActivities(context)); + + return CUPTI_SUCCESS; +} + +void CUPTIAPI +InjectionCallbackHandler( + void *pUserData, + CUpti_CallbackDomain domain, + CUpti_CallbackId callbackId, + void *pCallbackData) +{ + static const char *s_pFunctionName; + static uint32_t s_correlationId; + + const CUpti_CallbackData *pCallbackInfo = (CUpti_CallbackData *)pCallbackData; + + // Clear any previous CUPTI errors. cuptiGetLastError() retrieves and clears the last error. + // We don't treat this as fatal since it's just clearing state from previous operations. + CUptiResult _status = cuptiGetLastError(); + if (_status != CUPTI_SUCCESS && _status != CUPTI_ERROR_NOT_INITIALIZED) + { + const char *pErrorString; + cuptiGetResultString(_status, &pErrorString); + // Log but don't exit - this is just informational + std::cerr << "Warning: Cleared previous CUPTI error(" << _status << "): " << pErrorString << "\n"; + } + + switch (domain) + { + case CUPTI_CB_DOMAIN_STATE: + HandleDomainStateCallback(callbackId, (CUpti_StateData *)pCallbackData); + break; + case CUPTI_CB_DOMAIN_RESOURCE: + { + CUpti_ResourceData *pResourceData = (CUpti_ResourceData *)pCallbackData; + switch (callbackId) + { + case CUPTI_CBID_RESOURCE_GRAPHNODE_CREATED: + { + // Do not store info for the nodes that are created during graph instantiate. + if (s_pFunctionName && !strncmp(s_pFunctionName, "cudaGraphInstantiate", strlen("cudaGraphInstantiate"))) + { + break; + } + CUpti_GraphData *callbackData = (CUpti_GraphData *) pResourceData->resourceDescriptor; + uint64_t nodeId; + + // Query the graph node ID and store the API correlation id and function name. + CUPTI_API_CALL(cuptiGetGraphNodeId(callbackData->node, &nodeId)); + ApiData apiData; + apiData.correlationId = s_correlationId; + apiData.pFunctionName = s_pFunctionName; + nodeIdCorrelationMap[nodeId] = apiData; + break; + } + case CUPTI_CBID_RESOURCE_GRAPHNODE_CLONED: + { + CUpti_GraphData *callbackData = (CUpti_GraphData *) pResourceData->resourceDescriptor; + uint64_t nodeId, originalNodeId; + + // Overwrite the map entry with node ID of the cloned graph node. + CUPTI_API_CALL(cuptiGetGraphNodeId(callbackData->originalNode, &originalNodeId)); + NodeIdApiDataMap::iterator it = nodeIdCorrelationMap.find(originalNodeId); + if (it != nodeIdCorrelationMap.end()) + { + CUPTI_API_CALL(cuptiGetGraphNodeId(callbackData->node, &nodeId)); + ApiData apiData = it->second; + nodeIdCorrelationMap.erase(it); + nodeIdCorrelationMap[nodeId] = apiData; + } + break; + } + default: + break; + } + break; + } + case CUPTI_CB_DOMAIN_DRIVER_API: + { + switch (callbackId) + { + case CUPTI_DRIVER_TRACE_CBID_cuProfilerStart: + { + // We start profiling collection on exit of the API. + if (pCallbackInfo->callbackSite == CUPTI_API_EXIT) + { + OnProfilerStart(pCallbackInfo->context); + } + break; + } + case CUPTI_DRIVER_TRACE_CBID_cuProfilerStop: + { + // We stop profiling collection on entry of the API. + if (pCallbackInfo->callbackSite == CUPTI_API_ENTER) + { + OnProfilerStop(pCallbackInfo->context); + } + break; + } + default: + break; + } + break; + } + case CUPTI_CB_DOMAIN_RUNTIME_API: + { + if (pCallbackInfo->callbackSite == CUPTI_API_ENTER) + { + s_correlationId = pCallbackInfo->correlationId; + s_pFunctionName = pCallbackInfo->functionName; + } + + switch (callbackId) + { + case CUPTI_RUNTIME_TRACE_CBID_cudaDeviceReset_v3020: + { + if (pCallbackInfo->callbackSite == CUPTI_API_ENTER) + { + CUPTI_API_CALL(OnCudaDeviceReset()); + } + break; + } + default: + break; + } + break; + } + default: + break; + } +} + +static void +SetupCupti(void) +{ + UserData *pUserData = (UserData *)malloc(sizeof(UserData)); + MEMORY_ALLOCATION_CALL(pUserData); + + memset(pUserData, 0, sizeof(UserData)); + pUserData->pPostProcessActivityRecords = GraphTraceRecords; + pUserData->printActivityRecords = 1; + + // Common CUPTI Initialization. + // Configure output file from environment variable or use default + const char *outputPath = getenv("CUPTI_TRACE_OUTPUT_FILE"); + if (!outputPath) { + outputPath = "cupti_trace_output.txt"; // Default filename + } + + FILE *outputFile = stdout; // Default to stdout + if (strcmp(outputPath, "stdout") != 0) { + outputFile = fopen(outputPath, "w"); + if (!outputFile) { + std::cerr << "Failed to open output file '" << outputPath << "', falling back to stdout\n"; + outputFile = stdout; + } else { + std::cout << "CUPTI trace output will be written to: " << outputPath << "\n"; + } + } + InitCuptiTrace(pUserData, (void *)InjectionCallbackHandler, outputFile); + + injectionGlobals.subscriberHandle = globals.subscriberHandle; + + // Subscribe Driver callback to call OnProfilerStart/OnProfilerStop function. + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuProfilerStart)); + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_DRIVER_API, CUPTI_DRIVER_TRACE_CBID_cuProfilerStop)); + + // Enable callbacks for CUDA graph node tracking. + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_RESOURCE, CUPTI_CBID_RESOURCE_GRAPHNODE_CREATED)); + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_RESOURCE, CUPTI_CBID_RESOURCE_GRAPHNODE_CLONED)); + CUPTI_API_CALL_VERBOSE(cuptiEnableDomain(1, injectionGlobals.subscriberHandle, CUPTI_CB_DOMAIN_RUNTIME_API)); + + // Enable CUPTI activities. + CUPTI_API_CALL(EnableCuptiActivities(NULL)); +} + +#ifdef _WIN32 +extern "C" __declspec(dllexport) int +InitializeInjection(void) +#else +extern "C" int +InitializeInjection(void) +#endif +{ + if (injectionGlobals.initialized) + { + // Return 1 to indicate that the injection is already successfully initialized. + return 1; + } + + initializeInjectionMutex.lock(); + + // Initialize injection global options. + InitializeInjectionGlobals(); + + RegisterAtExitHandler(); + + // Initialize CUPTI. + SetupCupti(); + + injectionGlobals.tracingEnabled = 1; + injectionGlobals.initialized = 1; + + initializeInjectionMutex.unlock(); + + // Return 1 to indicate that the injection is successfully initialized. + return 1; +} diff --git a/src/xpu/flamegraph/cupti_trace/helper_cupti.h b/src/xpu/flamegraph/cupti_trace/helper_cupti.h new file mode 100644 index 0000000..13f6725 --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace/helper_cupti.h @@ -0,0 +1,184 @@ +/** + * Copyright 2022-2024 NVIDIA Corporation. All rights reserved. + * + * Please refer to the NVIDIA end user license agreement (EULA) associated + * with this source code for terms and conditions that govern your use of + * this software. Any use, reproduction, disclosure, or distribution of + * this software and related documentation outside the terms of the EULA + * is strictly prohibited. + * + */ + +//////////////////////////////////////////////////////////////////////////////// + +#ifndef HELPER_CUPTI_H_ +#define HELPER_CUPTI_H_ + +#pragma once + +#include + +#ifndef EXIT_WAIVED +#define EXIT_WAIVED 2 +#endif + +#if defined(WIN32) || defined(_WIN32) +#define stricmp _stricmp +#else +#define stricmp strcasecmp +#endif + +#define CUDA_MAX_DEVICES 256 // consider theoretical max devices as 256 +#define DEV_NAME_LEN 256 + +#ifndef DRIVER_API_CALL +#define DRIVER_API_CALL(apiFunctionCall) \ +do \ +{ \ + CUresult _status = apiFunctionCall; \ + if (_status != CUDA_SUCCESS) \ + { \ + const char *pErrorString; \ + cuGetErrorString(_status, &pErrorString); \ + \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error(" << _status << "): " \ + << pErrorString << ".\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef RUNTIME_API_CALL +#define RUNTIME_API_CALL(apiFunctionCall) \ +do \ +{ \ + cudaError_t _status = apiFunctionCall; \ + if (_status != cudaSuccess) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error(" << _status << "): " \ + << cudaGetErrorString(_status) << ".\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef CUPTI_API_CALL +#define CUPTI_API_CALL(apiFunctionCall) \ +do \ +{ \ + CUptiResult _status = apiFunctionCall; \ + if (_status != CUPTI_SUCCESS) \ + { \ + const char *pErrorString; \ + cuptiGetResultString(_status, &pErrorString); \ + \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error(" << _status << "): " \ + << pErrorString << ".\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef CUPTI_API_CALL_VERBOSE +#define CUPTI_API_CALL_VERBOSE(apiFunctionCall) \ +do \ +{ \ + std::cout << "Calling CUPTI API: " << #apiFunctionCall << "\n"; \ + \ + CUptiResult _status = apiFunctionCall; \ + if (_status != CUPTI_SUCCESS) \ + { \ + const char *pErrorString; \ + cuptiGetResultString(_status, &pErrorString); \ + \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error(" << _status << "): " \ + << pErrorString << ".\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef CUPTI_UTIL_CALL +#define CUPTI_UTIL_CALL(apiFunctionCall) \ +do \ +{ \ + CUptiUtilResult _status = apiFunctionCall; \ + if (_status != CUPTI_UTIL_SUCCESS) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error: " << _status << "\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef NVPW_API_CALL +#define NVPW_API_CALL(apiFunctionCall) \ +do \ +{ \ + NVPA_Status _status = apiFunctionCall; \ + if (_status != NVPA_STATUS_SUCCESS) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Function " \ + << #apiFunctionCall << " failed with error: " << _status << "\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef MEMORY_ALLOCATION_CALL +#define MEMORY_ALLOCATION_CALL(variable) \ +do \ +{ \ + if (variable == NULL) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << \ + " Memory allocation failed.\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef CHECK_CONDITION +#define CHECK_CONDITION(condition) \ +do \ +{ \ + if (!(condition)) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Condition " \ + << #condition << " failed.\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#ifndef CHECK_INTEGER_CONDITION +#define CHECK_INTEGER_CONDITION(argument1, operator, argument2) \ +do \ +{ \ + if (!(argument1 operator argument2)) \ + { \ + std::cerr << "\n\nError: " << __FILE__ << ":" << __LINE__ << ": Condition " \ + << #argument1 << " " << #operator << " " << #argument2 << " fails. " << \ + #argument1 << " = " << argument1 << ", " << #argument2 << " = " << \ + argument2 << "\n\n"; \ + \ + exit(EXIT_FAILURE); \ + } \ +} while (0) +#endif + +#endif // HELPER_CUPTI_H_ + diff --git a/src/xpu/flamegraph/cupti_trace/helper_cupti_activity.h b/src/xpu/flamegraph/cupti_trace/helper_cupti_activity.h new file mode 100644 index 0000000..fb88b50 --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace/helper_cupti_activity.h @@ -0,0 +1,2152 @@ +/** + * Copyright 2022-2024 NVIDIA Corporation. All rights reserved. + * + * Please refer to the NVIDIA end user license agreement (EULA) associated + * with this source code for terms and conditions that govern your use of + * this software. Any use, reproduction, disclosure, or distribution of + * this software and related documentation outside the terms of the EULA + * is strictly prohibited. + * + */ + +//////////////////////////////////////////////////////////////////////////////// + +#ifndef HELPER_CUPTI_ACTIVITY_H_ +#define HELPER_CUPTI_ACTIVITY_H_ + +#pragma once + +// System headers +#include +#include +#include +#include + +// CUPTI headers +#include +#include +#include +#include +#include + + +// Macros +#define LINE_SIZE 2048 + +// CUPTI buffer size 32 MB +#define BUF_SIZE (32 * 1024 * 1024) + +// 8-byte alignment for the buffers +#define ALIGN_SIZE (8) +#define ALIGN_BUFFER(buffer, align) \ + (((uintptr_t) (buffer) & ((align)-1)) ? ((buffer) + (align) - ((uintptr_t) (buffer) & ((align)-1))) : (buffer)) + +typedef uint64_t HashMapKey; + +// Data structures + +// Global state +typedef struct GlobalState_st +{ + CUpti_SubscriberHandle subscriberHandle; // CUPTI subcriber handle to subcribe to CUPTI callbacks. + size_t activityBufferSize; // CUPTI activity buffer size. + FILE *pOutputFile; // File handle to print the CUPTI activity records. default = stdout. + void *pUserData; // User data used to initialize CUPTI trace. Refer UserData structure. + uint64_t buffersRequested; // Requested buffers by CUPTI. + uint64_t buffersCompleted; // Completed buffers by received from CUPTI. +} GlobalState; + +// User data provided by the application using InitCuptiTrace() +// User need to allocate memory for this structure in the sample. +// Set the options according to the workloads requirement. +typedef struct UserData_st +{ + size_t activityBufferSize; // CUPTI activity buffer size. + size_t deviceBufferSize; // CUPTI device buffer size. + uint8_t flushAtStreamSync; // Flush CUPTI activity records at stream syncronization. + uint8_t flushAtCtxSync; // Flush CUPTI activity records at context syncronization. + uint8_t printCallbacks; // Print callbacks enabled in CUPTI. + uint8_t printActivityRecords; // Print CUPTI activity records. + uint8_t skipCuptiSubscription; // Check if the user application wants to skip subscription in CUPTI. + void (*pPostProcessActivityRecords)(CUpti_Activity *pRecord); // Provide function pointer in the user application for CUPTI records for post processing. +} UserData; + +// Global variables +static GlobalState globals = { 0 }; + +// Helper Functions +static const char * +GetActivityKindString( + CUpti_ActivityKind activityKind) +{ + switch (activityKind) + { + case CUPTI_ACTIVITY_KIND_MEMCPY: + return "MEMCPY"; + case CUPTI_ACTIVITY_KIND_MEMSET: + return "MEMSET"; + case CUPTI_ACTIVITY_KIND_KERNEL: + return "KERNEL"; + case CUPTI_ACTIVITY_KIND_DRIVER: + return "DRIVER"; + case CUPTI_ACTIVITY_KIND_RUNTIME: + return "RUNTIME"; + case CUPTI_ACTIVITY_KIND_DEVICE: + return "DEVICE"; + case CUPTI_ACTIVITY_KIND_CONTEXT: + return "CONTEXT"; + case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: + return "CONCURRENT_KERNEL"; + case CUPTI_ACTIVITY_KIND_NAME: + return "NAME"; + case CUPTI_ACTIVITY_KIND_MARKER: + return "MARKER"; + case CUPTI_ACTIVITY_KIND_MARKER_DATA: + return "MARKER_DATA"; + case CUPTI_ACTIVITY_KIND_SOURCE_LOCATOR: + return "SOURCE_LOCATOR"; + case CUPTI_ACTIVITY_KIND_GLOBAL_ACCESS: + return "GLOBAL_ACCESS"; + case CUPTI_ACTIVITY_KIND_BRANCH: + return "BRANCH"; + case CUPTI_ACTIVITY_KIND_OVERHEAD: + return "OVERHEAD"; + case CUPTI_ACTIVITY_KIND_CDP_KERNEL: + return "CDP_KERNEL"; + case CUPTI_ACTIVITY_KIND_PREEMPTION: + return "PREEMPTION"; + case CUPTI_ACTIVITY_KIND_ENVIRONMENT: + return "ENVIRONMENT"; + case CUPTI_ACTIVITY_KIND_MEMCPY2: + return "MEMCPY2"; + case CUPTI_ACTIVITY_KIND_INSTRUCTION_EXECUTION: + return "INSTRUCTION_EXECUTION"; + case CUPTI_ACTIVITY_KIND_UNIFIED_MEMORY_COUNTER: + return "UNIFIED_MEMORY_COUNTER"; + case CUPTI_ACTIVITY_KIND_FUNCTION: + return "FUNCTION"; + case CUPTI_ACTIVITY_KIND_MODULE: + return "MODULE"; + case CUPTI_ACTIVITY_KIND_DEVICE_ATTRIBUTE: + return "DEVICE_ATTRIBUTE"; + case CUPTI_ACTIVITY_KIND_SHARED_ACCESS: + return "SHARED_ACCESS"; + case CUPTI_ACTIVITY_KIND_PC_SAMPLING: + return "PC_SAMPLING"; + case CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO: + return "PC_SAMPLING_RECORD_INFO"; + case CUPTI_ACTIVITY_KIND_INSTRUCTION_CORRELATION: + return "INSTRUCTION_CORRELATION"; + case CUPTI_ACTIVITY_KIND_OPENACC_DATA: + return "OPENACC_DATA"; + case CUPTI_ACTIVITY_KIND_OPENACC_LAUNCH: + return "OPENACC_LAUNCH"; + case CUPTI_ACTIVITY_KIND_OPENACC_OTHER: + return "OPENACC_OTHER"; + case CUPTI_ACTIVITY_KIND_CUDA_EVENT: + return "CUDA_EVENT"; + case CUPTI_ACTIVITY_KIND_STREAM: + return "STREAM"; + case CUPTI_ACTIVITY_KIND_SYNCHRONIZATION: + return "SYNCHRONIZATION"; + case CUPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION: + return "EXTERNAL_CORRELATION"; + case CUPTI_ACTIVITY_KIND_NVLINK: + return "NVLINK"; + case CUPTI_ACTIVITY_KIND_MEMORY: + return "MEMORY"; + case CUPTI_ACTIVITY_KIND_PCIE: + return "PCIE"; + case CUPTI_ACTIVITY_KIND_OPENMP: + return "OPENMP"; + case CUPTI_ACTIVITY_KIND_INTERNAL_LAUNCH_API: + return "INTERNAL_LAUNCH_API"; + case CUPTI_ACTIVITY_KIND_MEMORY2: + return "MEMORY2"; + case CUPTI_ACTIVITY_KIND_MEMORY_POOL: + return "MEMORY_POOL"; + case CUPTI_ACTIVITY_KIND_GRAPH_TRACE: + return "GRAPH_TRACE"; + case CUPTI_ACTIVITY_KIND_JIT: + return "JIT"; + case CUPTI_ACTIVITY_KIND_MEM_DECOMPRESS: + return "MEM_DECOMPRESS"; + default: + return ""; + } +} + +static CUpti_ActivityKind +GetActivityKindFromString( + const char *pActivityKindString) +{ + if (!pActivityKindString) + { + std::cerr << "\n\nError: NULL string.\n\n"; + exit(-1); + } + + if (!stricmp(pActivityKindString, "MEMCPY")) + { + return CUPTI_ACTIVITY_KIND_MEMCPY; + } + else if (!stricmp(pActivityKindString, "MEMSET")) + { + return CUPTI_ACTIVITY_KIND_MEMSET; + } + else if (!stricmp(pActivityKindString, "KERNEL")) + { + return CUPTI_ACTIVITY_KIND_KERNEL; + } + else if (!stricmp(pActivityKindString, "DRIVER")) + { + return CUPTI_ACTIVITY_KIND_DRIVER; + } + else if (!stricmp(pActivityKindString, "RUNTIME")) + { + return CUPTI_ACTIVITY_KIND_RUNTIME; + } + else if (!stricmp(pActivityKindString, "DEVICE")) + { + return CUPTI_ACTIVITY_KIND_DEVICE; + } + else if (!stricmp(pActivityKindString, "CONTEXT")) + { + return CUPTI_ACTIVITY_KIND_CONTEXT; + } + else if (!stricmp(pActivityKindString, "CONCURRENT_KERNEL")) + { + return CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL; + } + else if (!stricmp(pActivityKindString, "NAME")) + { + return CUPTI_ACTIVITY_KIND_NAME; + } + else if (!stricmp(pActivityKindString, "MARKER")) + { + return CUPTI_ACTIVITY_KIND_MARKER; + } + else if (!stricmp(pActivityKindString, "MARKER_DATA")) + { + return CUPTI_ACTIVITY_KIND_MARKER_DATA; + } + else if (!stricmp(pActivityKindString, "SOURCE_LOCATOR")) + { + return CUPTI_ACTIVITY_KIND_SOURCE_LOCATOR; + } + else if (!stricmp(pActivityKindString, "GLOBAL_ACCESS")) + { + return CUPTI_ACTIVITY_KIND_GLOBAL_ACCESS; + } + else if (!stricmp(pActivityKindString, "BRANCH")) + { + return CUPTI_ACTIVITY_KIND_BRANCH; + } + else if (!stricmp(pActivityKindString, "OVERHEAD")) + { + return CUPTI_ACTIVITY_KIND_OVERHEAD; + } + else if (!stricmp(pActivityKindString, "CDP_KERNEL")) + { + return CUPTI_ACTIVITY_KIND_CDP_KERNEL; + } + else if (!stricmp(pActivityKindString, "PREEMPTION")) + { + return CUPTI_ACTIVITY_KIND_PREEMPTION; + } + else if (!stricmp(pActivityKindString, "ENVIRONMENT")) + { + return CUPTI_ACTIVITY_KIND_ENVIRONMENT; + } + else if (!stricmp(pActivityKindString, "MEMCPY2")) + { + return CUPTI_ACTIVITY_KIND_MEMCPY2; + } + else if (!stricmp(pActivityKindString, "INSTRUCTION_EXECUTION")) + { + return CUPTI_ACTIVITY_KIND_INSTRUCTION_EXECUTION; + } + else if (!stricmp(pActivityKindString, "UNIFIED_MEMORY_COUNTER")) + { + return CUPTI_ACTIVITY_KIND_UNIFIED_MEMORY_COUNTER; + } + else if (!stricmp(pActivityKindString, "FUNCTION")) + { + return CUPTI_ACTIVITY_KIND_FUNCTION; + } + else if (!stricmp(pActivityKindString, "MODULE")) + { + return CUPTI_ACTIVITY_KIND_MODULE; + } + else if (!stricmp(pActivityKindString, "DEVICE_ATTRIBUTE")) + { + return CUPTI_ACTIVITY_KIND_DEVICE_ATTRIBUTE; + } + else if (!stricmp(pActivityKindString, "SHARED_ACCESS")) + { + return CUPTI_ACTIVITY_KIND_SHARED_ACCESS; + } + else if (!stricmp(pActivityKindString, "PC_SAMPLING")) + { + return CUPTI_ACTIVITY_KIND_PC_SAMPLING; + } + else if (!stricmp(pActivityKindString, "PC_SAMPLING_RECORD_INFO")) + { + return CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO; + } + else if (!stricmp(pActivityKindString, "INSTRUCTION_CORRELATION")) + { + return CUPTI_ACTIVITY_KIND_INSTRUCTION_CORRELATION; + } + else if (!stricmp(pActivityKindString, "OPENACC_DATA")) + { + return CUPTI_ACTIVITY_KIND_OPENACC_DATA; + } + else if (!stricmp(pActivityKindString, "OPENACC_LAUNCH")) + { + return CUPTI_ACTIVITY_KIND_OPENACC_LAUNCH; + } + else if (!stricmp(pActivityKindString, "OPENACC_OTHER")) + { + return CUPTI_ACTIVITY_KIND_OPENACC_OTHER; + } + else if (!stricmp(pActivityKindString, "CUDA_EVENT")) + { + return CUPTI_ACTIVITY_KIND_CUDA_EVENT; + } + else if (!stricmp(pActivityKindString, "STREAM")) + { + return CUPTI_ACTIVITY_KIND_STREAM; + } + else if (!stricmp(pActivityKindString, "SYNCHRONIZATION")) + { + return CUPTI_ACTIVITY_KIND_SYNCHRONIZATION; + } + else if (!stricmp(pActivityKindString, "EXTERNAL_CORRELATION")) + { + return CUPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION; + } + else if (!stricmp(pActivityKindString, "NVLINK")) + { + return CUPTI_ACTIVITY_KIND_NVLINK; + } + else if (!stricmp(pActivityKindString, "MEMORY")) + { + return CUPTI_ACTIVITY_KIND_MEMORY; + } + else if (!stricmp(pActivityKindString, "PCIE")) + { + return CUPTI_ACTIVITY_KIND_PCIE; + } + else if (!stricmp(pActivityKindString, "OPENMP")) + { + return CUPTI_ACTIVITY_KIND_OPENMP; + } + else if (!stricmp(pActivityKindString, "INTERNAL_LAUNCH_API")) + { + return CUPTI_ACTIVITY_KIND_INTERNAL_LAUNCH_API; + } + else if (!stricmp(pActivityKindString, "MEMORY2")) + { + return CUPTI_ACTIVITY_KIND_MEMORY2; + } + else if (!stricmp(pActivityKindString, "MEMORY_POOL")) + { + return CUPTI_ACTIVITY_KIND_MEMORY_POOL; + } + else if (!stricmp(pActivityKindString, "GRAPH_TRACE")) + { + return CUPTI_ACTIVITY_KIND_GRAPH_TRACE; + } + else if (!stricmp(pActivityKindString, "JIT")) + { + return CUPTI_ACTIVITY_KIND_JIT; + } +else if (!stricmp(pActivityKindString, "MEM_DECOMPRESS")) + { + return CUPTI_ACTIVITY_KIND_MEM_DECOMPRESS; + } + else { + std::cerr << "\n\nError: Invalid string " << pActivityKindString << " cannot be converted to CUPTI Activity Kind.\n\n"; + exit(-1); + } +} + + + +static const char * +GetActivityObjectKindString( + CUpti_ActivityObjectKind objectKind) +{ + switch (objectKind) + { + case CUPTI_ACTIVITY_OBJECT_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_OBJECT_PROCESS: + return "PROCESS"; + case CUPTI_ACTIVITY_OBJECT_THREAD: + return "THREAD"; + case CUPTI_ACTIVITY_OBJECT_DEVICE: + return "DEVICE"; + case CUPTI_ACTIVITY_OBJECT_CONTEXT: + return "CONTEXT"; + case CUPTI_ACTIVITY_OBJECT_STREAM: + return "STREAM"; + default: + return ""; + } +} + +static uint32_t +GetActivityObjectKindId( + CUpti_ActivityObjectKind objectKind, + CUpti_ActivityObjectKindId *pObjectKindId) +{ + switch (objectKind) + { + case CUPTI_ACTIVITY_OBJECT_UNKNOWN: + return 0xffffffff; + case CUPTI_ACTIVITY_OBJECT_PROCESS: + return pObjectKindId->pt.processId; + case CUPTI_ACTIVITY_OBJECT_THREAD: + return pObjectKindId->pt.threadId; + case CUPTI_ACTIVITY_OBJECT_DEVICE: + return pObjectKindId->dcs.deviceId; + case CUPTI_ACTIVITY_OBJECT_CONTEXT: + return pObjectKindId->dcs.contextId; + case CUPTI_ACTIVITY_OBJECT_STREAM: + return pObjectKindId->dcs.streamId; + default: + return 0xffffffff; + } +} + +static const char * +GetActivityOverheadKindString( + CUpti_ActivityOverheadKind overheadKind) +{ + switch (overheadKind) + { + case CUPTI_ACTIVITY_OVERHEAD_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_OVERHEAD_DRIVER_COMPILER: + return "DRIVER_COMPILER"; + case CUPTI_ACTIVITY_OVERHEAD_CUPTI_BUFFER_FLUSH: + return "CUPTI_BUFFER_FLUSH"; + case CUPTI_ACTIVITY_OVERHEAD_CUPTI_INSTRUMENTATION: + return "CUPTI_INSTRUMENTATION"; + case CUPTI_ACTIVITY_OVERHEAD_CUPTI_RESOURCE: + return "CUPTI_RESOURCE"; + case CUPTI_ACTIVITY_OVERHEAD_RUNTIME_TRIGGERED_MODULE_LOADING: + return "RUNTIME_TRIGGERED_MODULE_LOADING"; + case CUPTI_ACTIVITY_OVERHEAD_LAZY_FUNCTION_LOADING: + return "LAZY_FUNCTION_LOADING"; + case CUPTI_ACTIVITY_OVERHEAD_COMMAND_BUFFER_FULL: + return "COMMAND_BUFFER_FULL"; + case CUPTI_ACTIVITY_OVERHEAD_ACTIVITY_BUFFER_REQUEST: + return "ACTIVITY_BUFFER_REQUEST"; + case CUPTI_ACTIVITY_OVERHEAD_UVM_ACTIVITY_INIT: + return "UVM_ACTIVITY_INIT"; + default: + return ""; + } +} + +static const char * +GetComputeApiKindString( + CUpti_ActivityComputeApiKind computeApiKind) +{ + switch (computeApiKind) + { + case CUPTI_ACTIVITY_COMPUTE_API_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_COMPUTE_API_CUDA: + return "CUDA"; + case CUPTI_ACTIVITY_COMPUTE_API_CUDA_MPS: + return "CUDA_MPS"; + default: + return ""; + } +} + +static const char * +GetStallReasonString( + CUpti_ActivityPCSamplingStallReason pcSamplingStallReason) +{ + switch (pcSamplingStallReason) + { + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_NONE: + return "NONE"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_INST_FETCH: + return "INSTRUCTION_FETCH"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_EXEC_DEPENDENCY: + return "EXECUTION_DEPENDENCY"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_MEMORY_DEPENDENCY: + return "MEMORY_DEPENDENCY"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_TEXTURE: + return "TEXTURE"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_SYNC: + return "SYNC"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_CONSTANT_MEMORY_DEPENDENCY: + return "CONSTANT_MEMORY_DEPENDENCY"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_PIPE_BUSY: + return "PIPE_BUSY"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_MEMORY_THROTTLE: + return "MEMORY_THROTTLE"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_NOT_SELECTED: + return "SELECTED"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_OTHER: + return "OTHER"; + case CUPTI_ACTIVITY_PC_SAMPLING_STALL_SLEEPING: + return "SLEEPING"; + default: + return ""; + } +} + +static const char * +GetMemcpyKindString( + CUpti_ActivityMemcpyKind memcpyKind) +{ + switch (memcpyKind) + { + case CUPTI_ACTIVITY_MEMCPY_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_MEMCPY_KIND_HTOD: + return "HtoD"; + case CUPTI_ACTIVITY_MEMCPY_KIND_DTOH: + return "DtoH"; + case CUPTI_ACTIVITY_MEMCPY_KIND_HTOA: + return "HtoA"; + case CUPTI_ACTIVITY_MEMCPY_KIND_ATOH: + return "AtoH"; + case CUPTI_ACTIVITY_MEMCPY_KIND_ATOA: + return "AtoA"; + case CUPTI_ACTIVITY_MEMCPY_KIND_ATOD: + return "AtoD"; + case CUPTI_ACTIVITY_MEMCPY_KIND_DTOA: + return "DtoA"; + case CUPTI_ACTIVITY_MEMCPY_KIND_DTOD: + return "DtoD"; + case CUPTI_ACTIVITY_MEMCPY_KIND_HTOH: + return "HtoH"; + case CUPTI_ACTIVITY_MEMCPY_KIND_PTOP: + return "PtoP"; + default: + return ""; + } +} + +static const char * +GetMemoryKindString( + CUpti_ActivityMemoryKind memoryKind) +{ + switch (memoryKind) + { + case CUPTI_ACTIVITY_MEMORY_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_MEMORY_KIND_PAGEABLE: + return "PAGEABLE"; + case CUPTI_ACTIVITY_MEMORY_KIND_PINNED: + return "PINNED"; + case CUPTI_ACTIVITY_MEMORY_KIND_DEVICE: + return "DEVICE"; + case CUPTI_ACTIVITY_MEMORY_KIND_ARRAY: + return "ARRAY"; + case CUPTI_ACTIVITY_MEMORY_KIND_MANAGED: + return "MANAGED"; + case CUPTI_ACTIVITY_MEMORY_KIND_DEVICE_STATIC: + return "DEVICE_STATIC"; + case CUPTI_ACTIVITY_MEMORY_KIND_MANAGED_STATIC: + return "MANAGED_STATIC"; + default: + return ""; + } +} + +static const char * +GetPreemptionKindString( + CUpti_ActivityPreemptionKind preemptionKind) +{ + switch (preemptionKind) + { + case CUPTI_ACTIVITY_PREEMPTION_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_PREEMPTION_KIND_SAVE: + return "SAVE"; + case CUPTI_ACTIVITY_PREEMPTION_KIND_RESTORE: + return "RESTORE"; + default: + return ""; + } +} + +static const char * +GetActivityEnvironmentKindString( + CUpti_ActivityEnvironmentKind environmentKind) +{ + switch (environmentKind) + { + case CUPTI_ACTIVITY_ENVIRONMENT_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_ENVIRONMENT_SPEED: + return "SPEED"; + case CUPTI_ACTIVITY_ENVIRONMENT_TEMPERATURE: + return "TEMPERATURE"; + case CUPTI_ACTIVITY_ENVIRONMENT_POWER: + return "POWER"; + case CUPTI_ACTIVITY_ENVIRONMENT_COOLING: + return "COOLING"; + default: + return ""; + } +} + +static const char * +GetUvmCounterScopeString( + CUpti_ActivityUnifiedMemoryCounterScope unifiedMemoryCounterScope) +{ + switch (unifiedMemoryCounterScope) + { + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_SCOPE_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_SCOPE_PROCESS_SINGLE_DEVICE: + return "PROCESS_SINGLE_DEVICE"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_SCOPE_PROCESS_ALL_DEVICES: + return "PROCESS_ALL_DEVICES"; + default: + return ""; + } +} + +static const char * +GetUvmCounterKindString( + CUpti_ActivityUnifiedMemoryCounterKind unifiedMemoryCounterKind) +{ + switch (unifiedMemoryCounterKind) + { + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_BYTES_TRANSFER_HTOD: + return "BYTES_TRANSFER_HTOD"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_BYTES_TRANSFER_DTOH: + return "BYTES_TRANSFER_DTOH"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_CPU_PAGE_FAULT_COUNT: + return "CPU_PAGE_FAULT_COUNT"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_GPU_PAGE_FAULT: + return "GPU_PAGE_FAULT"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_THRASHING: + return "THRASHING"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_THROTTLING: + return "THROTTLING"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_REMOTE_MAP: + return "REMOTE_MAP"; + case CUPTI_ACTIVITY_UNIFIED_MEMORY_COUNTER_KIND_BYTES_TRANSFER_DTOD: + return "BYTES_TRANSFER_DTOD"; + default: + return ""; + } +} + +static const char * +GetSynchronizationType( + CUpti_ActivitySynchronizationType syncronizationType) +{ + switch (syncronizationType) + { + case CUPTI_ACTIVITY_SYNCHRONIZATION_TYPE_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_SYNCHRONIZATION_TYPE_EVENT_SYNCHRONIZE: + return "EVENT_SYNCHRONIZE"; + case CUPTI_ACTIVITY_SYNCHRONIZATION_TYPE_STREAM_WAIT_EVENT: + return "STREAM_WAIT_EVENT"; + case CUPTI_ACTIVITY_SYNCHRONIZATION_TYPE_STREAM_SYNCHRONIZE: + return "STREAM_SYNCHRONIZE"; + case CUPTI_ACTIVITY_SYNCHRONIZATION_TYPE_CONTEXT_SYNCHRONIZE: + return "CONTEXT_SYNCHRONIZE"; + default: + return ""; + } +} + +static const char * +GetStreamType( + CUpti_ActivityStreamFlag streamFlag) +{ + switch (streamFlag) + { + case CUPTI_ACTIVITY_STREAM_CREATE_FLAG_UNKNOWN: + return "UNKNOWN"; + case CUPTI_ACTIVITY_STREAM_CREATE_FLAG_DEFAULT: + return "DEFAULT_STREAM"; + case CUPTI_ACTIVITY_STREAM_CREATE_FLAG_NON_BLOCKING: + return "NON_BLOCKING_STREAM"; + case CUPTI_ACTIVITY_STREAM_CREATE_FLAG_NULL: + return "NULL_STREAM"; + default: + return ""; + } +} + +static const char * +GetMemoryOperationTypeString( + CUpti_ActivityMemoryOperationType memoryOperationType) +{ + switch (memoryOperationType) + { + case CUPTI_ACTIVITY_MEMORY_OPERATION_TYPE_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_MEMORY_OPERATION_TYPE_ALLOCATION: + return "ALLOCATE"; + case CUPTI_ACTIVITY_MEMORY_OPERATION_TYPE_RELEASE: + return "RELEASE"; + default: + return ""; + } +} + +static const char * +GetMemoryPoolTypeString( + CUpti_ActivityMemoryPoolType memoryPoolType) +{ + switch (memoryPoolType) + { + case CUPTI_ACTIVITY_MEMORY_POOL_TYPE_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_MEMORY_POOL_TYPE_LOCAL: + return "LOCAL"; + case CUPTI_ACTIVITY_MEMORY_POOL_TYPE_IMPORTED: + return "IMPORTED"; + default: + return ""; + } +} + +static const char * +GetMemoryPoolOperationTypeString( + CUpti_ActivityMemoryPoolOperationType memoryPoolOperationType) +{ + switch (memoryPoolOperationType) + { + case CUPTI_ACTIVITY_MEMORY_POOL_OPERATION_TYPE_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_MEMORY_POOL_OPERATION_TYPE_CREATED: + return "MEM_POOL_CREATED"; + case CUPTI_ACTIVITY_MEMORY_POOL_OPERATION_TYPE_DESTROYED: + return "MEM_POOL_DESTROYED"; + case CUPTI_ACTIVITY_MEMORY_POOL_OPERATION_TYPE_TRIMMED: + return "MEM_POOL_TRIMMED"; + default: + return ""; + } +} + +static const char * +GetChannelType( + CUpti_ChannelType channelType) +{ + switch (channelType) + { + case CUPTI_CHANNEL_TYPE_INVALID: + return "INVALID"; + case CUPTI_CHANNEL_TYPE_COMPUTE: + return "COMPUTE"; + case CUPTI_CHANNEL_TYPE_ASYNC_MEMCPY: + return "ASYNC_MEMCPY"; + case CUPTI_CHANNEL_TYPE_DECOMP: + return "DECOMP"; + default: + return ""; + } +} + +static const char * +GetJitEntryType( + CUpti_ActivityJitEntryType jitEntryType) +{ + switch (jitEntryType) + { + case CUPTI_ACTIVITY_JIT_ENTRY_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_JIT_ENTRY_PTX_TO_CUBIN: + return "PTX_TO_CUBIN"; + case CUPTI_ACTIVITY_JIT_ENTRY_NVVM_IR_TO_PTX: + return "NVVM_IR_TO_PTX"; + default: + return ""; + } +} + +static const char * +GetJitOperationType( + CUpti_ActivityJitOperationType jitOperationType) +{ + switch (jitOperationType) + { + case CUPTI_ACTIVITY_JIT_OPERATION_INVALID: + return "INVALID"; + case CUPTI_ACTIVITY_JIT_OPERATION_CACHE_LOAD: + return "CACHE_LOAD"; + case CUPTI_ACTIVITY_JIT_OPERATION_CACHE_STORE: + return "CACHE_STORE"; + case CUPTI_ACTIVITY_JIT_OPERATION_COMPILE: + return "COMPILE"; + default: + return ""; + } +} + +static const char * +GetName( + const char *pName) +{ + if (pName == NULL) + { + return ""; + } + + return pName; +} + +static const char * +GetDomainName( + const char *pName) +{ + if (pName == NULL) + { + return ""; + } + + return pName; +} + +static const char * +GetOpenAccConstructString( + CUpti_OpenAccConstructKind openAccConstructKind) +{ + switch (openAccConstructKind) + { + case CUPTI_OPENACC_CONSTRUCT_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_OPENACC_CONSTRUCT_KIND_PARALLEL: + return "PARALLEL"; + case CUPTI_OPENACC_CONSTRUCT_KIND_KERNELS: + return "KERNELS"; + case CUPTI_OPENACC_CONSTRUCT_KIND_LOOP: + return "LOOP"; + case CUPTI_OPENACC_CONSTRUCT_KIND_DATA: + return "DATA"; + case CUPTI_OPENACC_CONSTRUCT_KIND_ENTER_DATA: + return "ENTER_DATA"; + case CUPTI_OPENACC_CONSTRUCT_KIND_EXIT_DATA: + return "EXIT_DATA"; + case CUPTI_OPENACC_CONSTRUCT_KIND_HOST_DATA: + return "HOST_DATA"; + case CUPTI_OPENACC_CONSTRUCT_KIND_ATOMIC: + return "ATOMIC"; + case CUPTI_OPENACC_CONSTRUCT_KIND_DECLARE: + return "DECLARE"; + case CUPTI_OPENACC_CONSTRUCT_KIND_INIT: + return "INIT"; + case CUPTI_OPENACC_CONSTRUCT_KIND_SHUTDOWN: + return "SHUTDOWN"; + case CUPTI_OPENACC_CONSTRUCT_KIND_SET: + return "SET"; + case CUPTI_OPENACC_CONSTRUCT_KIND_UPDATE: + return "UPDATE"; + case CUPTI_OPENACC_CONSTRUCT_KIND_ROUTINE: + return "ROUTINE"; + case CUPTI_OPENACC_CONSTRUCT_KIND_WAIT: + return "WAIT"; + case CUPTI_OPENACC_CONSTRUCT_KIND_RUNTIME_API: + return "RUNTIME_API"; + default: + return NULL; + } +} + +static const char * +GetExternalCorrelationKindString( + CUpti_ExternalCorrelationKind externalCorrelationKind) +{ + switch (externalCorrelationKind) + { + case CUPTI_EXTERNAL_CORRELATION_KIND_INVALID: + return "INVALID"; + case CUPTI_EXTERNAL_CORRELATION_KIND_UNKNOWN: + return "UNKNOWN"; + case CUPTI_EXTERNAL_CORRELATION_KIND_OPENACC: + return "OPENACC"; + case CUPTI_EXTERNAL_CORRELATION_KIND_CUSTOM0: + return "CUSTOM0"; + case CUPTI_EXTERNAL_CORRELATION_KIND_CUSTOM1: + return "CUSTOM1"; + case CUPTI_EXTERNAL_CORRELATION_KIND_CUSTOM2: + return "CUSTOM2"; + default: + return ""; + } +} + +static const char * +GetDevTypeNvlink( + CUpti_DevType devType) +{ + switch (devType) + { + case CUPTI_DEV_TYPE_INVALID: + return "INVALID"; + case CUPTI_DEV_TYPE_GPU: + return "GPU"; + case CUPTI_DEV_TYPE_NPU: + return "CPU"; + default: + return ""; + } +} + +static uint32_t +GetCorrelationId( + CUpti_Activity *pRecord) +{ + switch (pRecord->kind) + { + case CUPTI_ACTIVITY_KIND_MEMCPY: + return ((CUpti_ActivityMemcpy6 *)pRecord)->correlationId; + case CUPTI_ACTIVITY_KIND_MEMSET: + return ((CUpti_ActivityMemset4 *)pRecord)->correlationId; + case CUPTI_ACTIVITY_KIND_KERNEL: + case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: + return ((CUpti_ActivityKernel9 *)pRecord)->correlationId; + case CUPTI_ACTIVITY_KIND_DRIVER: + case CUPTI_ACTIVITY_KIND_RUNTIME: + return ((CUpti_ActivityAPI *)pRecord)->correlationId; + case CUPTI_ACTIVITY_KIND_CDP_KERNEL: + return ((CUpti_ActivityCdpKernel *)pRecord)->correlationId; + case CUPTI_ACTIVITY_KIND_MEMCPY2: + return ((CUpti_ActivityMemcpyPtoP4 *)pRecord)->correlationId; + default: + return 0; + } +} + +static void +PrintOpenaccCommon( + FILE *pFileHandle, + CUpti_ActivityOpenAcc *pOpenAcc) +{ + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, eventKind %u, parentConstruct %s, version %u, implicit %u, deviceType %u, deviceNumber %u, threadId %u,\n" + " async %llu, asyncMap %llu, lineNo %u, endLineNo %u, funcLineNo %u, endFuncLineNo %u,\n" + " cuDeviceId %u, cuContextId %u, cuStreamId %u, cuProcessId %u, cuThreadId %u, externalId %llu", + GetActivityKindString(pOpenAcc->kind), + (unsigned long long)pOpenAcc->start, + (unsigned long long)pOpenAcc->end, + (unsigned long long)(pOpenAcc->end - pOpenAcc->start), + pOpenAcc->eventKind, + GetOpenAccConstructString((CUpti_OpenAccConstructKind)pOpenAcc->parentConstruct), + pOpenAcc->version, + pOpenAcc->implicit, + pOpenAcc->deviceType, + pOpenAcc->deviceNumber, + pOpenAcc->threadId, + (unsigned long long)pOpenAcc->async, + (unsigned long long)pOpenAcc->asyncMap, + pOpenAcc->lineNo, + pOpenAcc->endLineNo, + pOpenAcc->funcLineNo, + pOpenAcc->funcEndLineNo, + pOpenAcc->cuDeviceId, + pOpenAcc->cuContextId, + pOpenAcc->cuStreamId, + pOpenAcc->cuProcessId, + pOpenAcc->cuThreadId, + (unsigned long long)pOpenAcc->externalId); + + fprintf(pFileHandle, ", srcFile %s", pOpenAcc->srcFile ? pOpenAcc->srcFile : "?"); + fprintf(pFileHandle, ", funcName %s", pOpenAcc->funcName ? pOpenAcc->funcName : "?"); + +} + +static void +PrintActivity( + CUpti_Activity *pRecord, + FILE *pFileHandle) +{ + CUpti_ActivityKind activityKind = pRecord->kind; + + switch (activityKind) + { + case CUPTI_ACTIVITY_KIND_MEMCPY: + { + CUpti_ActivityMemcpy6 *pMemcpyRecord = (CUpti_ActivityMemcpy6 *)pRecord; + + fprintf(pFileHandle, "%s \"%s\" [ %llu, %llu ] duration %llu, size %llu, copyCount %llu, srcKind %s, dstKind %s, correlationId %u\n" + "\tdeviceId %u, contextId %u, streamId %u, graphId %u, graphNodeId %llu, channelId %u, channelType %s\n", + GetActivityKindString(pMemcpyRecord->kind), + GetMemcpyKindString((CUpti_ActivityMemcpyKind)pMemcpyRecord->copyKind), + (unsigned long long)pMemcpyRecord->start, + (unsigned long long)pMemcpyRecord->end, + (unsigned long long)(pMemcpyRecord->end - pMemcpyRecord->start), + (unsigned long long)pMemcpyRecord->bytes, + (unsigned long long)pMemcpyRecord->copyCount, + GetMemoryKindString((CUpti_ActivityMemoryKind)pMemcpyRecord->srcKind), + GetMemoryKindString((CUpti_ActivityMemoryKind)pMemcpyRecord->dstKind), + pMemcpyRecord->correlationId, + pMemcpyRecord->deviceId, + pMemcpyRecord->contextId, + pMemcpyRecord->streamId, + pMemcpyRecord->graphId, + (unsigned long long)pMemcpyRecord->graphNodeId, + pMemcpyRecord->channelID, + GetChannelType(pMemcpyRecord->channelType)); + + break; + } + case CUPTI_ACTIVITY_KIND_MEMSET: + { + CUpti_ActivityMemset4 *pMemsetRecord = (CUpti_ActivityMemset4 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, value %u, size %llu, correlationId %u\n" + "\tdeviceId %u, contextId %u, streamId %u, graphId %u, graphNodeId %llu, channelId %u, channelType %s\n", + GetActivityKindString(pMemsetRecord->kind), + (unsigned long long)pMemsetRecord->start, + (unsigned long long)pMemsetRecord->end, + (unsigned long long)(pMemsetRecord->end - pMemsetRecord->start), + pMemsetRecord->value, + (unsigned long long)pMemsetRecord->bytes, + pMemsetRecord->correlationId, + pMemsetRecord->deviceId, + pMemsetRecord->contextId, + pMemsetRecord->streamId, + pMemsetRecord->graphId, + (unsigned long long)pMemsetRecord->graphNodeId, + pMemsetRecord->channelID, + GetChannelType(pMemsetRecord->channelType)); + + break; + } + case CUPTI_ACTIVITY_KIND_KERNEL: + case CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL: + { + CUpti_ActivityKernel9 *pKernelRecord = (CUpti_ActivityKernel9 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, \"%s\", correlationId %u, cacheConfigRequested %d, cacheConfigExecuted %d\n" + "\tgrid [ %u, %u, %u ], block [ %u, %u, %u ], cluster [ %u, %u, %u ], sharedMemory (static %u, dynamic %u)\n" + "\tdeviceId %u, contextId %u, streamId %u, graphId %u, graphNodeId %llu, channelId %u, channelType %s\n", + GetActivityKindString(pKernelRecord->kind), + (unsigned long long)pKernelRecord->start, + (unsigned long long)pKernelRecord->end, + (unsigned long long)(pKernelRecord->end - pKernelRecord->start), + GetName(pKernelRecord->name), + pKernelRecord->correlationId, + pKernelRecord->cacheConfig.config.requested, + pKernelRecord->cacheConfig.config.executed, + pKernelRecord->gridX, + pKernelRecord->gridY, + pKernelRecord->gridZ, + pKernelRecord->blockX, + pKernelRecord->blockY, + pKernelRecord->blockZ, + pKernelRecord->clusterX, + pKernelRecord->clusterY, + pKernelRecord->clusterZ, + pKernelRecord->staticSharedMemory, + pKernelRecord->dynamicSharedMemory, + pKernelRecord->deviceId, + pKernelRecord->contextId, + pKernelRecord->streamId, + pKernelRecord->graphId, + (unsigned long long)pKernelRecord->graphNodeId, + pKernelRecord->channelID, + GetChannelType(pKernelRecord->channelType)); + + break; + } + case CUPTI_ACTIVITY_KIND_DRIVER: + case CUPTI_ACTIVITY_KIND_RUNTIME: + case CUPTI_ACTIVITY_KIND_INTERNAL_LAUNCH_API: + { + CUpti_ActivityAPI *pApiRecord = (CUpti_ActivityAPI *)pRecord; + const char* pName = NULL; + + if (pApiRecord->kind == CUPTI_ACTIVITY_KIND_DRIVER) + { + cuptiGetCallbackName(CUPTI_CB_DOMAIN_DRIVER_API, pApiRecord->cbid, &pName); + } + else if (pApiRecord->kind == CUPTI_ACTIVITY_KIND_RUNTIME) + { + cuptiGetCallbackName(CUPTI_CB_DOMAIN_RUNTIME_API, pApiRecord->cbid, &pName); + } + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, \"%s\", cbid %u, processId %u, threadId %u, correlationId %u\n", + GetActivityKindString(pApiRecord->kind), + (unsigned long long)pApiRecord->start, + (unsigned long long)pApiRecord->end, + (unsigned long long)(pApiRecord->end - pApiRecord->start), + GetName(pName), + pApiRecord->cbid, + pApiRecord->processId, + pApiRecord->threadId, + pApiRecord->correlationId); + + break; + } + case CUPTI_ACTIVITY_KIND_DEVICE: + { + CUpti_ActivityDevice5 *pDeviceRecord = (CUpti_ActivityDevice5 *)pRecord; + + fprintf(pFileHandle, "%s %s [ %u ]\n", + GetActivityKindString(pDeviceRecord->kind), + GetName(pDeviceRecord->name), + pDeviceRecord->id); + + break; + } + case CUPTI_ACTIVITY_KIND_CONTEXT: + { + CUpti_ActivityContext3 *pContextRecord = (CUpti_ActivityContext3 *)pRecord; + + fprintf(pFileHandle, "%s computeApiKind %s, contextId %u, deviceId %u, nullStreamId %d, CIG mode %d\n", + GetActivityKindString(pContextRecord->kind), + GetComputeApiKindString((CUpti_ActivityComputeApiKind) pContextRecord->computeApiKind), + pContextRecord->contextId, + pContextRecord->deviceId, + (int)pContextRecord->nullStreamId, + pContextRecord->cigMode); + + break; + } + case CUPTI_ACTIVITY_KIND_NAME: + { + CUpti_ActivityName *pNameRecord = (CUpti_ActivityName *)pRecord; + + switch(pNameRecord->objectKind) + { + case CUPTI_ACTIVITY_OBJECT_CONTEXT: + { + fprintf(pFileHandle, "%s %s %u %s id %u, name %s\n", + GetActivityKindString(pNameRecord->kind), + GetActivityObjectKindString(pNameRecord->objectKind), + GetActivityObjectKindId(pNameRecord->objectKind, &pNameRecord->objectId), + GetActivityObjectKindString(CUPTI_ACTIVITY_OBJECT_DEVICE), + GetActivityObjectKindId(CUPTI_ACTIVITY_OBJECT_DEVICE, &pNameRecord->objectId), + GetName(pNameRecord->name)); + + break; + } + case CUPTI_ACTIVITY_OBJECT_STREAM: + { + fprintf(pFileHandle, "%s %s %u %s %u %s id %u, name %s\n", + GetActivityKindString(pNameRecord->kind), + GetActivityObjectKindString(pNameRecord->objectKind), + GetActivityObjectKindId(pNameRecord->objectKind, &pNameRecord->objectId), + GetActivityObjectKindString(CUPTI_ACTIVITY_OBJECT_CONTEXT), + GetActivityObjectKindId(CUPTI_ACTIVITY_OBJECT_CONTEXT, &pNameRecord->objectId), + GetActivityObjectKindString(CUPTI_ACTIVITY_OBJECT_DEVICE), + GetActivityObjectKindId(CUPTI_ACTIVITY_OBJECT_DEVICE, &pNameRecord->objectId), + GetName(pNameRecord->name)); + + break; + } + default: + { + fprintf(pFileHandle, "%s %s id %u, name %s\n", + GetActivityKindString(pNameRecord->kind), + GetActivityObjectKindString(pNameRecord->objectKind), + GetActivityObjectKindId(pNameRecord->objectKind, &pNameRecord->objectId), + GetName(pNameRecord->name)); + break; + } + } + + break; + } + case CUPTI_ACTIVITY_KIND_MARKER: + { + CUpti_ActivityMarker2 *pMarkerRecord = (CUpti_ActivityMarker2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu ] id %u, domain %s, name %s\n", + GetActivityKindString(pMarkerRecord->kind), + (unsigned long long)pMarkerRecord->timestamp, + pMarkerRecord->id, + GetDomainName(pMarkerRecord->domain), + GetName(pMarkerRecord->name)); + + break; + } + case CUPTI_ACTIVITY_KIND_MARKER_DATA: + { + CUpti_ActivityMarkerData *pMarkerDataRecord = (CUpti_ActivityMarkerData *)pRecord; + + fprintf(pFileHandle, "%s id %u, color 0x%x, category %u, payload %llu/%f\n", + GetActivityKindString(pMarkerDataRecord->kind), + pMarkerDataRecord->id, + pMarkerDataRecord->color, + pMarkerDataRecord->category, + (unsigned long long)pMarkerDataRecord->payload.metricValueUint64, + pMarkerDataRecord->payload.metricValueDouble); + + break; + } + case CUPTI_ACTIVITY_KIND_SOURCE_LOCATOR: + { + CUpti_ActivitySourceLocator *pSourceLocatorRecord = (CUpti_ActivitySourceLocator *)pRecord; + + char line[LINE_SIZE]; + FILE *pLocalFileHandle = NULL; + + if ((pLocalFileHandle = fopen(pSourceLocatorRecord->fileName, "rt")) == NULL) + { + fprintf(pFileHandle, "Failed to open source file: %s\n", pSourceLocatorRecord->fileName); + } + else + { + uint32_t temp = 0; + + while (pSourceLocatorRecord->lineNumber > temp) + { + if (fgets(line, LINE_SIZE, pLocalFileHandle) == NULL) + { + fprintf(pFileHandle, "Line %d could not be found in file %s.\n", + pSourceLocatorRecord->lineNumber, pSourceLocatorRecord->fileName); + break; + } + + temp++; + } + fprintf(pFileHandle, "%d, %s", pSourceLocatorRecord-> id, line); + fclose(pLocalFileHandle); + } + + break; + } + case CUPTI_ACTIVITY_KIND_GLOBAL_ACCESS: + { + CUpti_ActivityGlobalAccess3 *pGlobalAccessRecord = (CUpti_ActivityGlobalAccess3 *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%llx, correlationId %u, operation %s, isCached %s, size %u,\n" + " executed %u, threadsExecuted %llu, transactions %llu, optimizedTransactions %llu\n", + GetActivityKindString(pGlobalAccessRecord->kind), + pGlobalAccessRecord->sourceLocatorId, + pGlobalAccessRecord->functionId, + (unsigned long long)pGlobalAccessRecord->pcOffset, + pGlobalAccessRecord->correlationId, + ((pGlobalAccessRecord->flags & CUPTI_ACTIVITY_FLAG_GLOBAL_ACCESS_KIND_LOAD) ? "Load" : "Store"), + ((pGlobalAccessRecord->flags & CUPTI_ACTIVITY_FLAG_GLOBAL_ACCESS_KIND_CACHED) ? "Yes" : "No"), + (uint32_t)(pGlobalAccessRecord->flags & CUPTI_ACTIVITY_FLAG_GLOBAL_ACCESS_KIND_SIZE_MASK), + pGlobalAccessRecord->executed, + (unsigned long long)pGlobalAccessRecord->threadsExecuted, + (unsigned long long)pGlobalAccessRecord->l2_transactions, + (unsigned long long)pGlobalAccessRecord->theoreticalL2Transactions); + + break; + } + case CUPTI_ACTIVITY_KIND_BRANCH: + { + CUpti_ActivityBranch2 *pBranchRecord = (CUpti_ActivityBranch2 *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%x, correlationId %u,\n" + " executed %u, threadsExecuted %llu, diverged %u\n", + GetActivityKindString(pBranchRecord->kind), + pBranchRecord->sourceLocatorId, + pBranchRecord->functionId, + pBranchRecord->pcOffset, + pBranchRecord->correlationId, + pBranchRecord->executed, + (unsigned long long)pBranchRecord->threadsExecuted, + pBranchRecord->diverged); + + break; + } + case CUPTI_ACTIVITY_KIND_OVERHEAD: + { + CUpti_ActivityOverhead3 *pOverheadRecord = (CUpti_ActivityOverhead3 *)pRecord; + + fprintf(pFileHandle, "%s %s [ %llu, %llu ] duration %llu, %s, id %u, correlation id %lu\n", + GetActivityKindString(pOverheadRecord->kind), + GetActivityOverheadKindString(pOverheadRecord->overheadKind), + (unsigned long long)pOverheadRecord->start, + (unsigned long long)pOverheadRecord->end, + (unsigned long long)(pOverheadRecord->end - pOverheadRecord->start), + GetActivityObjectKindString(pOverheadRecord->objectKind), + GetActivityObjectKindId(pOverheadRecord->objectKind, &pOverheadRecord->objectId), + (unsigned long)pOverheadRecord->correlationId); + if (pOverheadRecord->overheadData) + { + switch (pOverheadRecord->overheadKind) + { + case CUPTI_ACTIVITY_OVERHEAD_COMMAND_BUFFER_FULL: + { + CUpti_ActivityOverheadCommandBufferFullData* pCommandBufferData = (CUpti_ActivityOverheadCommandBufferFullData*)pOverheadRecord->overheadData; + fprintf(pFileHandle, "CUpti_ActivityOverheadCommandBufferFullData : commandBufferLength %d channelID %d channelType %d\n", + pCommandBufferData->commandBufferLength, + pCommandBufferData->channelID, + pCommandBufferData->channelType); + break; + } + default: + { + break; + } + } + + } + + break; + } + case CUPTI_ACTIVITY_KIND_CDP_KERNEL: + { + CUpti_ActivityCdpKernel *pCdpKernelRecord = (CUpti_ActivityCdpKernel *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, \"%s\", deviceId %u, contextId %u, streamId %u, gridId %lld, correlationId %u,\n" + "\tgrid [ %u, %u, %u ], block [ %u, %u, %u ], registersPerThread %u, sharedMemory (static %u, dynamic %u), parentGridId %lld, parentBlockId [ %u, %u, %u ]\n", + GetActivityKindString(pCdpKernelRecord->kind), + (unsigned long long)pCdpKernelRecord->start, + (unsigned long long)pCdpKernelRecord->end, + (unsigned long long)(pCdpKernelRecord->end - pCdpKernelRecord->start), + GetName(pCdpKernelRecord->name), + pCdpKernelRecord->deviceId, + pCdpKernelRecord->contextId, + pCdpKernelRecord->streamId, + (long long)pCdpKernelRecord->gridId, + pCdpKernelRecord->correlationId, + pCdpKernelRecord->gridX, + pCdpKernelRecord->gridY, + pCdpKernelRecord->gridZ, + pCdpKernelRecord->blockX, + pCdpKernelRecord->blockY, + pCdpKernelRecord->blockZ, + pCdpKernelRecord->registersPerThread, + pCdpKernelRecord->staticSharedMemory, + pCdpKernelRecord->dynamicSharedMemory, + (long long)pCdpKernelRecord->parentGridId, + pCdpKernelRecord->parentBlockX, + pCdpKernelRecord->parentBlockY, + pCdpKernelRecord->parentBlockZ); + + break; + } + case CUPTI_ACTIVITY_KIND_PREEMPTION: + { + CUpti_ActivityPreemption *pPreemptionRecord = (CUpti_ActivityPreemption *)pRecord; + + fprintf(pFileHandle, "%s preemptionKind %s [ %llu ] gridId %lld, block [ %u, %u, %u ]\n", + GetActivityKindString(pPreemptionRecord->kind), + GetPreemptionKindString(pPreemptionRecord->preemptionKind), + (unsigned long long)pPreemptionRecord->timestamp, + (long long)pPreemptionRecord->gridId, + pPreemptionRecord->blockX, + pPreemptionRecord->blockY, + pPreemptionRecord->blockZ); + + break; + } + case CUPTI_ACTIVITY_KIND_ENVIRONMENT: + { + CUpti_ActivityEnvironment *pEnvironmentRecord = (CUpti_ActivityEnvironment *)pRecord; + + switch (pEnvironmentRecord->environmentKind) + { + case CUPTI_ACTIVITY_ENVIRONMENT_SPEED: + { + fprintf(pFileHandle, "%s: kind=SPEED, deviceId %u, timestamp %llu, memoryClock %u, smClock %u, pcieLinkGen %u, pcieLinkWidth %u, clocksThrottleReasons %u\n", + GetActivityKindString(pEnvironmentRecord->kind), + pEnvironmentRecord->deviceId, + (unsigned long long)pEnvironmentRecord->timestamp, + pEnvironmentRecord->data.speed.memoryClock, + pEnvironmentRecord->data.speed.smClock, + pEnvironmentRecord->data.speed.pcieLinkGen, + pEnvironmentRecord->data.speed.pcieLinkWidth, + pEnvironmentRecord->data.speed.clocksThrottleReasons); + + break; + } + case CUPTI_ACTIVITY_ENVIRONMENT_TEMPERATURE: + { + fprintf(pFileHandle, "%s: kind=TEMPERATURE, deviceId %u, timestamp %llu, gpuTemperature %u\n", + GetActivityKindString(pEnvironmentRecord->kind), + pEnvironmentRecord->deviceId, + (unsigned long long)pEnvironmentRecord->timestamp, + pEnvironmentRecord->data.temperature.gpuTemperature); + + break; + } + case CUPTI_ACTIVITY_ENVIRONMENT_POWER: + { + fprintf(pFileHandle, "%s: kind=POWER, deviceId %u, timestamp %llu, power %u, powerLimit %u\n", + GetActivityKindString(pEnvironmentRecord->kind), + pEnvironmentRecord->deviceId, + (unsigned long long)pEnvironmentRecord->timestamp, + pEnvironmentRecord->data.power.power, + pEnvironmentRecord->data.power.powerLimit); + + break; + } + case CUPTI_ACTIVITY_ENVIRONMENT_COOLING: + { + fprintf(pFileHandle, "%s: kind=COOLING, deviceId %u, timestamp %llu, fanSpeed %u\n", + GetActivityKindString(pEnvironmentRecord->kind), + pEnvironmentRecord->deviceId, + (unsigned long long)pEnvironmentRecord->timestamp, + pEnvironmentRecord->data.cooling.fanSpeed); + + break; + } + default: + break; + } + + break; + } + case CUPTI_ACTIVITY_KIND_MEMCPY2: + { + CUpti_ActivityMemcpyPtoP4 *pMemcpyPtoPRecord = (CUpti_ActivityMemcpyPtoP4 *)pRecord; + + fprintf(pFileHandle, "%s \"%s\" [ %llu, %llu ] duration %llu, size %llu, srcKind %s, dstKind %s, correlationId %u,\n" + "\tdeviceId %u, contextId %u, streamId %u, graphId %u, graphNodeId %llu, channelId %u, channelType %s\n" + "\tsrcDeviceId %u, srcContextId %u, dstDeviceId %u, dstContextId %u\n", + GetActivityKindString(pMemcpyPtoPRecord->kind), + GetMemcpyKindString((CUpti_ActivityMemcpyKind)pMemcpyPtoPRecord->copyKind), + (unsigned long long)pMemcpyPtoPRecord->start, + (unsigned long long)pMemcpyPtoPRecord->end, + (unsigned long long)(pMemcpyPtoPRecord->end - pMemcpyPtoPRecord->start), + (unsigned long long)pMemcpyPtoPRecord->bytes, + GetMemoryKindString((CUpti_ActivityMemoryKind)pMemcpyPtoPRecord->srcKind), + GetMemoryKindString((CUpti_ActivityMemoryKind)pMemcpyPtoPRecord->dstKind), + pMemcpyPtoPRecord->correlationId, + pMemcpyPtoPRecord->deviceId, + pMemcpyPtoPRecord->contextId, + pMemcpyPtoPRecord->streamId, + pMemcpyPtoPRecord->graphId, + (unsigned long long)pMemcpyPtoPRecord->graphNodeId, + pMemcpyPtoPRecord->channelID, + GetChannelType(pMemcpyPtoPRecord->channelType), + pMemcpyPtoPRecord->srcDeviceId, + pMemcpyPtoPRecord->srcContextId, + pMemcpyPtoPRecord->dstDeviceId, + pMemcpyPtoPRecord->dstContextId); + + break; + } + case CUPTI_ACTIVITY_KIND_INSTRUCTION_EXECUTION: + { + CUpti_ActivityInstructionExecution *pInstructionExecutionRecord = (CUpti_ActivityInstructionExecution *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%x, correlationId %u,\n" + " valid %s, executed %u, threadsExecuted %llu, notPredOffThreadsExecuted %llu\n", + GetActivityKindString(pInstructionExecutionRecord->kind), + pInstructionExecutionRecord->sourceLocatorId, + pInstructionExecutionRecord->functionId, + pInstructionExecutionRecord->pcOffset, + pInstructionExecutionRecord->correlationId, + ((pInstructionExecutionRecord->flags & CUPTI_ACTIVITY_FLAG_INSTRUCTION_VALUE_INVALID) ? "no" : "yes"), + pInstructionExecutionRecord->executed, + (unsigned long long)pInstructionExecutionRecord->threadsExecuted, + (unsigned long long)pInstructionExecutionRecord->notPredOffThreadsExecuted); + + break; + } + case CUPTI_ACTIVITY_KIND_UNIFIED_MEMORY_COUNTER: + { + CUpti_ActivityUnifiedMemoryCounter2 *pUnifiedMemoryCounterRecord = (CUpti_ActivityUnifiedMemoryCounter2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, counterKind %s, value %llu, address %llx, srcId %u, dstId %u, processId %u\n", + GetActivityKindString(pUnifiedMemoryCounterRecord->kind), + (unsigned long long)pUnifiedMemoryCounterRecord->start, + (unsigned long long)pUnifiedMemoryCounterRecord->end, + (unsigned long long)(pUnifiedMemoryCounterRecord->end - pUnifiedMemoryCounterRecord->start), + GetUvmCounterKindString(pUnifiedMemoryCounterRecord->counterKind), + (unsigned long long)pUnifiedMemoryCounterRecord->value, + (unsigned long long)pUnifiedMemoryCounterRecord->address, + pUnifiedMemoryCounterRecord->srcId, + pUnifiedMemoryCounterRecord->dstId, + pUnifiedMemoryCounterRecord->processId); + + break; + } + case CUPTI_ACTIVITY_KIND_FUNCTION: + { + CUpti_ActivityFunction *pFunctionRecord = (CUpti_ActivityFunction *)pRecord; + + fprintf(pFileHandle, "%s id %u, contextId %u, moduleId %u, functionIndex %u, name %s\n", + GetActivityKindString(pFunctionRecord->kind), + pFunctionRecord->id, + pFunctionRecord->contextId, + pFunctionRecord->moduleId, + pFunctionRecord->functionIndex, + GetName(pFunctionRecord->name)); + + break; + } + case CUPTI_ACTIVITY_KIND_MODULE: + { + CUpti_ActivityModule *pModuleRecord = (CUpti_ActivityModule *)pRecord; + + fprintf(pFileHandle, "%s contextId %u, id %d, cubinSize %d\n", + GetActivityKindString(pModuleRecord->kind), + pModuleRecord->contextId, + pModuleRecord->id, + pModuleRecord->cubinSize); + + break; + } + case CUPTI_ACTIVITY_KIND_DEVICE_ATTRIBUTE: + { + CUpti_ActivityDeviceAttribute *pDeviceAttributeRecord = (CUpti_ActivityDeviceAttribute *)pRecord; + + fprintf(pFileHandle, "%s %u, deviceId %u, value 0x%llx\n", + GetActivityKindString(pDeviceAttributeRecord->kind), + pDeviceAttributeRecord->attribute.cupti, + pDeviceAttributeRecord->deviceId, + (unsigned long long)pDeviceAttributeRecord->value.vUint64); + + break; + } + case CUPTI_ACTIVITY_KIND_SHARED_ACCESS: + { + CUpti_ActivitySharedAccess *pSharedAccessRecord = (CUpti_ActivitySharedAccess *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%x, correlationId %u,\n" + " op %s, size %u, executed %u, threadsExecuted %llu, sharedTransactions %llu, optimizedTransactions %llu\n", + GetActivityKindString(pSharedAccessRecord->kind), + pSharedAccessRecord->sourceLocatorId, + pSharedAccessRecord->functionId, + pSharedAccessRecord->pcOffset, + pSharedAccessRecord->correlationId, + ((pSharedAccessRecord->flags & CUPTI_ACTIVITY_FLAG_SHARED_ACCESS_KIND_LOAD) ? "Load" : "Store"), + (uint32_t)(pSharedAccessRecord->flags & CUPTI_ACTIVITY_FLAG_SHARED_ACCESS_KIND_SIZE_MASK), + pSharedAccessRecord->executed, + (unsigned long long)pSharedAccessRecord->threadsExecuted, + (unsigned long long)pSharedAccessRecord->sharedTransactions, + (unsigned long long)pSharedAccessRecord->theoreticalSharedTransactions); + + break; + } + case CUPTI_ACTIVITY_KIND_PC_SAMPLING: + { + CUpti_ActivityPCSampling3 *pPcSamplingRecord = (CUpti_ActivityPCSampling3 *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%llx, correlationId %u, samples %u, latencySamples %u, stallReason %s\n", + GetActivityKindString(pPcSamplingRecord->kind), + pPcSamplingRecord->sourceLocatorId, + pPcSamplingRecord->functionId, + (unsigned long long)pPcSamplingRecord->pcOffset, + pPcSamplingRecord->correlationId, + pPcSamplingRecord->samples, + pPcSamplingRecord->latencySamples, + GetStallReasonString(pPcSamplingRecord->stallReason)); + + break; + } + case CUPTI_ACTIVITY_KIND_PC_SAMPLING_RECORD_INFO: + { + CUpti_ActivityPCSamplingRecordInfo *pPcSamplingRecordInfo = (CUpti_ActivityPCSamplingRecordInfo *)pRecord; + + fprintf(pFileHandle, "%s correlationId %u, totalSamples %llu, droppedSamples %llu, samplingPeriodInCycles %llu\n", + GetActivityKindString(pPcSamplingRecordInfo->kind), + pPcSamplingRecordInfo->correlationId, + (unsigned long long)pPcSamplingRecordInfo->totalSamples, + (unsigned long long)pPcSamplingRecordInfo->droppedSamples, + (unsigned long long)pPcSamplingRecordInfo->samplingPeriodInCycles); + + break; + } + case CUPTI_ACTIVITY_KIND_INSTRUCTION_CORRELATION: + { + CUpti_ActivityInstructionCorrelation *pInstructionCorrelationRecord = (CUpti_ActivityInstructionCorrelation *)pRecord; + + fprintf(pFileHandle, "%s sourceLocatorId %u, functionId %u, pcOffset 0x%x\n", + GetActivityKindString(pInstructionCorrelationRecord->kind), + pInstructionCorrelationRecord->sourceLocatorId, + pInstructionCorrelationRecord->functionId, + pInstructionCorrelationRecord->pcOffset); + + break; + } + case CUPTI_ACTIVITY_KIND_OPENACC_DATA: + { + CUpti_ActivityOpenAccData *pOpenaccDataRecord = (CUpti_ActivityOpenAccData *)pRecord; + + PrintOpenaccCommon(pFileHandle, (CUpti_ActivityOpenAcc*)pOpenaccDataRecord); + + fprintf(pFileHandle, ", bytes %llu, varName %s\n", + (long long unsigned)pOpenaccDataRecord->bytes, + pOpenaccDataRecord->varName ? pOpenaccDataRecord->varName : "?"); + + break; + } + case CUPTI_ACTIVITY_KIND_OPENACC_LAUNCH: + { + CUpti_ActivityOpenAccLaunch *pOpenaccLaunchRecord = (CUpti_ActivityOpenAccLaunch *)pRecord; + + PrintOpenaccCommon(pFileHandle, (CUpti_ActivityOpenAcc*)pOpenaccLaunchRecord); + + fprintf(pFileHandle, ", numGangs %llu, numWorkers %llu, vectorLength %llu, kernelName %s\n", + (long long unsigned)pOpenaccLaunchRecord->numGangs, + (long long unsigned)pOpenaccLaunchRecord->numWorkers, + (long long unsigned)pOpenaccLaunchRecord->vectorLength, + pOpenaccLaunchRecord->kernelName ? pOpenaccLaunchRecord->kernelName : "?"); + + break; + } + case CUPTI_ACTIVITY_KIND_OPENACC_OTHER: + { + CUpti_ActivityOpenAccOther *pOpenaccOtherRecord = (CUpti_ActivityOpenAccOther *)pRecord; + + PrintOpenaccCommon(pFileHandle, (CUpti_ActivityOpenAcc*)pOpenaccOtherRecord); + printf("\n"); + + break; + } + case CUPTI_ACTIVITY_KIND_CUDA_EVENT: + { + CUpti_ActivityCudaEvent2 *pCudaEventRecord = (CUpti_ActivityCudaEvent2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu ] contextId %u, streamId %u, correlationId %u, eventId %u, cudaEventSyncId %llu\n", + GetActivityKindString(pCudaEventRecord->kind), + (long long unsigned)pCudaEventRecord->deviceTimestamp, + pCudaEventRecord->contextId, + pCudaEventRecord->streamId, + pCudaEventRecord->correlationId, + pCudaEventRecord->eventId, + (long long unsigned)pCudaEventRecord->cudaEventSyncId); + + break; + } + case CUPTI_ACTIVITY_KIND_STREAM: + { + CUpti_ActivityStream *pStreamRecord = (CUpti_ActivityStream *)pRecord; + + fprintf(pFileHandle, "%s type %s, priority %u, contextId %u, streamId %u, correlationId %u\n", + GetActivityKindString(pStreamRecord->kind), + GetStreamType(pStreamRecord->flag), + pStreamRecord->priority, + pStreamRecord->contextId, + pStreamRecord->streamId, + pStreamRecord->correlationId); + break; + } + case CUPTI_ACTIVITY_KIND_SYNCHRONIZATION: + { + CUpti_ActivitySynchronization2 *pSynchronizationRecord = (CUpti_ActivitySynchronization2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, type %s, contextId %u, streamId %d, correlationId %u, eventId %d, cudaEventSyncId %llu\n", + GetActivityKindString(pSynchronizationRecord->kind), + (unsigned long long)pSynchronizationRecord->start, + (unsigned long long)pSynchronizationRecord->end, + (unsigned long long)(pSynchronizationRecord->end - pSynchronizationRecord->start), + GetSynchronizationType(pSynchronizationRecord->type), + pSynchronizationRecord->contextId, + (int32_t)pSynchronizationRecord->streamId, + pSynchronizationRecord->correlationId, + (int32_t)pSynchronizationRecord->cudaEventId, + (long long unsigned)pSynchronizationRecord->cudaEventSyncId); + + break; + } + case CUPTI_ACTIVITY_KIND_EXTERNAL_CORRELATION: + { + CUpti_ActivityExternalCorrelation *pExternalCorrelationRecord = (CUpti_ActivityExternalCorrelation *)pRecord; + + fprintf(pFileHandle, "%s externalKind %s, correlationId %llu, externalId %llu\n", + GetActivityKindString(pExternalCorrelationRecord->kind), + GetExternalCorrelationKindString(pExternalCorrelationRecord->externalKind), + (long long unsigned)pExternalCorrelationRecord->correlationId, + (long long unsigned)pExternalCorrelationRecord->externalId); + + break; + } + case CUPTI_ACTIVITY_KIND_NVLINK: + { + CUpti_ActivityNvLink4 *pNvLinkRecord = (CUpti_ActivityNvLink4 *)pRecord; + unsigned int i = 0; + + fprintf(pFileHandle, "%s typeDev0 %s, typeDev1 %s, sysmem %d, peer %d, physicalNvLinkCount %d, ", + GetActivityKindString(pNvLinkRecord->kind), + GetDevTypeNvlink(pNvLinkRecord->typeDev0), + GetDevTypeNvlink(pNvLinkRecord->typeDev1), + ((pNvLinkRecord->flag & CUPTI_LINK_FLAG_SYSMEM_ACCESS) ? 1 : 0), + ((pNvLinkRecord->flag & CUPTI_LINK_FLAG_PEER_ACCESS) ? 1 : 0), + pNvLinkRecord->physicalNvLinkCount); + + fprintf(pFileHandle, "portDev0 "); + for (i = 0 ; i < pNvLinkRecord->physicalNvLinkCount ; i++ ) + { + fprintf(pFileHandle, "%d, ", pNvLinkRecord->portDev0[i]); + } + + + fprintf(pFileHandle, "portDev1 "); + for (i = 0 ; i < pNvLinkRecord->physicalNvLinkCount ; i++ ) + { + fprintf(pFileHandle, "%d, ", pNvLinkRecord->portDev1[i]); + } + + fprintf(pFileHandle, "bandwidth %llu\n", (long long unsigned int)pNvLinkRecord->bandwidth); + + break; + } + case CUPTI_ACTIVITY_KIND_MEMORY: + { + CUpti_ActivityMemory *pMemoryRecord = (CUpti_ActivityMemory *)(void *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, size %llu bytes, address %llu, memoryKind %s, deviceId %u, contextId %u, processId %u\n", + GetActivityKindString(pMemoryRecord->kind), + (unsigned long long)pMemoryRecord->start, + (unsigned long long)pMemoryRecord->end, + (unsigned long long)(pMemoryRecord->end - pMemoryRecord->start), + (unsigned long long)pMemoryRecord->bytes, + (unsigned long long)pMemoryRecord->address, + GetMemoryKindString(pMemoryRecord->memoryKind), + pMemoryRecord->deviceId, + pMemoryRecord->contextId, + pMemoryRecord->processId); + + break; + } + case CUPTI_ACTIVITY_KIND_PCIE: + { + CUpti_ActivityPcie *pPcieRecord = (CUpti_ActivityPcie *)pRecord; + + if (pPcieRecord->type == CUPTI_PCIE_DEVICE_TYPE_GPU) + { + fprintf(pFileHandle, "%s GPU %u, domain %u, upstreamBus %u, link rate %u GT/s, link width %u bits.\n", + GetActivityKindString(pPcieRecord->kind), + pPcieRecord->id.devId, + pPcieRecord->domain, + pPcieRecord->upstreamBus, + pPcieRecord->linkRate, + pPcieRecord->linkWidth); + } + else if (pPcieRecord->type == CUPTI_PCIE_DEVICE_TYPE_BRIDGE) + { + fprintf(pFileHandle, "%s bridgeId %u, domain %u, upstream Bus %u, downstream Bus %u, link rate %u GT/s, link width %u bits.\n", + GetActivityKindString(pPcieRecord->kind), + pPcieRecord->id.bridgeId, + pPcieRecord->domain, + pPcieRecord->upstreamBus, + pPcieRecord->attr.bridgeAttr.secondaryBus, + pPcieRecord->linkRate, + pPcieRecord->linkWidth); + } + + break; + } + case CUPTI_ACTIVITY_KIND_OPENMP: + { + CUpti_ActivityOpenMp *pOpenMpRecord = (CUpti_ActivityOpenMp *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, eventKind %u, cuProcessId %u, cuThreadId %u\n", + GetActivityKindString(pOpenMpRecord->kind), + (unsigned long long)pOpenMpRecord->start, + (unsigned long long)pOpenMpRecord->end, + (unsigned long long)(pOpenMpRecord->end - pOpenMpRecord->start), + pOpenMpRecord->eventKind, + pOpenMpRecord->cuProcessId, + pOpenMpRecord->cuThreadId); + + break; + } + case CUPTI_ACTIVITY_KIND_MEMORY2: + { + CUpti_ActivityMemory4 *pMemory2Record = (CUpti_ActivityMemory4 *)(void *)pRecord; + + fprintf(pFileHandle, "%s [ %llu ] memoryOperationType %s, memoryKind %s, size %llu, address %llu, pc %llu,\n" + " deviceId %u, contextId %u, streamId %u, processId %u, correlationId %u, isAsync %u,\n" + " memoryPool %s, memoryPoolAddress %llu, memoryPoolThreshold %llu\n" + " source %s\n" + , + GetActivityKindString(pMemory2Record->kind), + (unsigned long long)pMemory2Record->timestamp, + GetMemoryOperationTypeString(pMemory2Record->memoryOperationType), + GetMemoryKindString(pMemory2Record->memoryKind), + (unsigned long long)pMemory2Record->bytes, + (unsigned long long)pMemory2Record->address, + (unsigned long long)pMemory2Record->PC, + pMemory2Record->deviceId, + pMemory2Record->contextId, + pMemory2Record->streamId, + pMemory2Record->processId, + pMemory2Record->correlationId, + pMemory2Record->isAsync, + GetMemoryPoolTypeString(pMemory2Record->memoryPoolConfig.memoryPoolType), + (unsigned long long)pMemory2Record->memoryPoolConfig.address, + (unsigned long long)pMemory2Record->memoryPoolConfig.releaseThreshold + ,pMemory2Record->source + ); + + if (pMemory2Record->memoryPoolConfig.memoryPoolType == CUPTI_ACTIVITY_MEMORY_POOL_TYPE_LOCAL) + { + fprintf(pFileHandle, ", memoryPoolSize: %llu, memoryPoolUtilizedSize: %llu", + (unsigned long long)pMemory2Record->memoryPoolConfig.pool.size, + (unsigned long long)pMemory2Record->memoryPoolConfig.utilizedSize); + } + else if (pMemory2Record->memoryPoolConfig.memoryPoolType == CUPTI_ACTIVITY_MEMORY_POOL_TYPE_IMPORTED) + { + fprintf(pFileHandle, ", memoryPoolProcessId: %llu", + (unsigned long long)pMemory2Record->memoryPoolConfig.pool.processId); + } + + fprintf(pFileHandle, "\n"); + + break; + } + case CUPTI_ACTIVITY_KIND_MEMORY_POOL: + { + CUpti_ActivityMemoryPool2 *pMemoryPoolRecord = (CUpti_ActivityMemoryPool2 *)(void *)pRecord; + + fprintf(pFileHandle, "%s [ %llu ] memoryPoolOperation %s, memoryPool %s, address %llu, size %llu, utilizedSize %llu, releaseThreshold %llu,\n" + " deviceId %u, processId %u, correlationId %u\n", + GetActivityKindString(pMemoryPoolRecord->kind), + (unsigned long long)pMemoryPoolRecord->timestamp, + GetMemoryPoolOperationTypeString(pMemoryPoolRecord->memoryPoolOperationType), + GetMemoryPoolTypeString(pMemoryPoolRecord->memoryPoolType), + (unsigned long long)pMemoryPoolRecord->address, + (unsigned long long)pMemoryPoolRecord->size, + (unsigned long long)pMemoryPoolRecord->utilizedSize, + (unsigned long long)pMemoryPoolRecord->releaseThreshold, + pMemoryPoolRecord->deviceId, + pMemoryPoolRecord->processId, + pMemoryPoolRecord->correlationId); + + break; + } + case CUPTI_ACTIVITY_KIND_GRAPH_TRACE: + { + CUpti_ActivityGraphTrace2 *pGraphTraceRecord = (CUpti_ActivityGraphTrace2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, correlationId %u\n deviceId %u, contextId %u, streamId %u, graphId %u\n", + GetActivityKindString(pGraphTraceRecord->kind), + (unsigned long long)pGraphTraceRecord->start, + (unsigned long long)pGraphTraceRecord->end, + (unsigned long long)(pGraphTraceRecord->end - pGraphTraceRecord->start), + pGraphTraceRecord->correlationId, + pGraphTraceRecord->deviceId, + pGraphTraceRecord->contextId, + pGraphTraceRecord->streamId, + pGraphTraceRecord->graphId); + + break; + } + case CUPTI_ACTIVITY_KIND_JIT: + { + CUpti_ActivityJit2 *pJitRecord = (CUpti_ActivityJit2 *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, deviceId %u, correlationId %u, processId %u, threadId %u\n" + "jitEntryType %s, jitOperationType %s, jitOperationCorrelationId %llu\n cacheSize %llu, cachePath %s\n", + GetActivityKindString(pJitRecord->kind), + (unsigned long long)pJitRecord->start, + (unsigned long long)pJitRecord->end, + (unsigned long long)(pJitRecord->end - pJitRecord->start), + pJitRecord->deviceId, + pJitRecord->correlationId, + pJitRecord->processId, + pJitRecord->threadId, + GetJitEntryType(pJitRecord->jitEntryType), + GetJitOperationType(pJitRecord->jitOperationType), + (unsigned long long)pJitRecord->jitOperationCorrelationId, + (unsigned long long)pJitRecord->cacheSize, + GetName(pJitRecord->cachePath)); + + break; + } + case CUPTI_ACTIVITY_KIND_MEM_DECOMPRESS: + { + CUpti_ActivityMemDecompress *pMemDecompress = (CUpti_ActivityMemDecompress *)pRecord; + + fprintf(pFileHandle, "%s [ %llu, %llu ] duration %llu, deviceId %u, contextId %u, streamId %u, correlationId %u\n" + "channelId %u, channelType %s, numberOfOperations %u, sourceBytes %llu\n", + GetActivityKindString(pMemDecompress->kind), + (unsigned long long)pMemDecompress->start, + (unsigned long long)pMemDecompress->end, + (unsigned long long)(pMemDecompress->end - pMemDecompress->start), + pMemDecompress->deviceId, + pMemDecompress->contextId, + pMemDecompress->streamId, + pMemDecompress->correlationId, + pMemDecompress->channelID, + GetChannelType(pMemDecompress->channelType), + pMemDecompress->numberOfOperations, + (unsigned long long)pMemDecompress->sourceBytes); + + break; + } + default: + fprintf(pFileHandle, " \n"); + break; + } +} + +static void +PrintActivityBuffer( + uint8_t *pBuffer, + size_t validBytes, + FILE *pFileHandle, + void *pUserData) +{ + CUpti_Activity *pRecord = NULL; + CUptiResult status = CUPTI_SUCCESS; + + do + { + status = cuptiActivityGetNextRecord(pBuffer, validBytes, &pRecord); + if (status == CUPTI_SUCCESS) + { + if (!pUserData || + (pUserData && ((UserData *)pUserData)->printActivityRecords)) + { + PrintActivity(pRecord, pFileHandle); + } + + if (pUserData && + ((UserData *)pUserData)->pPostProcessActivityRecords) + { + ((UserData *)pUserData)->pPostProcessActivityRecords(pRecord); + } + } + else if (status == CUPTI_ERROR_MAX_LIMIT_REACHED) + { + break; + } + else if (status == CUPTI_ERROR_INVALID_KIND) + { + break; + } + else + { + CUPTI_API_CALL(status); + } + } while (1); +} + +// Buffer Management Functions +static void CUPTIAPI +BufferRequested( + uint8_t **ppBuffer, + size_t *pSize, + size_t *pMaxNumRecords) +{ + uint8_t *pBuffer = (uint8_t *) malloc(globals.activityBufferSize + ALIGN_SIZE); + MEMORY_ALLOCATION_CALL(pBuffer); + + *pSize = globals.activityBufferSize; + *ppBuffer = ALIGN_BUFFER(pBuffer, ALIGN_SIZE); + *pMaxNumRecords = 0; + + globals.buffersRequested++; +} + +static void CUPTIAPI +BufferCompleted( + CUcontext context, + uint32_t streamId, + uint8_t *pBuffer, + size_t size, + size_t validSize) +{ + if (validSize > 0) + { + FILE *pOutputFile = globals.pOutputFile; + if (!pOutputFile) + { + pOutputFile = stdout; + } + + PrintActivityBuffer(pBuffer, validSize, pOutputFile, globals.pUserData); + } + + globals.buffersCompleted++; + free(pBuffer); +} + +// CUPTI callback functions +static void +HandleSyncronizationCallbacks( + CUpti_CallbackId callbackId, + const CUpti_SynchronizeData *pSynchronizeData, + void *pUserData) +{ + // Flush the CUPTI activity records buffer on context synchronization + if (callbackId == CUPTI_CBID_SYNCHRONIZE_CONTEXT_SYNCHRONIZED && + ((UserData *)pUserData)->flushAtCtxSync) + { + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(0)); + } + // Flush the CUPTI activity records buffer on stream synchronization + else if (callbackId == CUPTI_CBID_SYNCHRONIZE_STREAM_SYNCHRONIZED && + ((UserData *)pUserData)->flushAtStreamSync) + { + uint32_t streamId = 0; + CUPTI_API_CALL_VERBOSE(cuptiGetStreamId(pSynchronizeData->context, pSynchronizeData->stream, &streamId)); + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(0)); + } +} + +static void +HandleDomainStateCallback( + CUpti_CallbackId callbackId, + const CUpti_StateData *pStateData) +{ + switch (callbackId) + { + case CUPTI_CBID_STATE_FATAL_ERROR: + { + const char *errorString = NULL; + cuptiGetResultString(pStateData->notification.result, &errorString); + + fprintf(globals.pOutputFile, "\nCUPTI encountered fatal error: %s\n", errorString); + fprintf(globals.pOutputFile, "Error message: %s\n", pStateData->notification.message); + + // Exiting the application if fatal error encountered in CUPTI + // If there is a CUPTI fatal error, it means CUPTI has stopped profiling the application. + exit(EXIT_FAILURE); + } + default: + break; + } +} + +static void CUPTIAPI +CuptiCallbackHandler( + void *pUserData, + CUpti_CallbackDomain domain, + CUpti_CallbackId callbackId, + const void *pCallbackData) +{ + CUPTI_API_CALL(cuptiGetLastError()); + + if (((UserData *)pUserData)->printCallbacks && + globals.pOutputFile != NULL) + { + fprintf(globals.pOutputFile, "CUPTI Callback: Domain %d CbId %d\n", domain, callbackId); + fflush(globals.pOutputFile); + } + + const CUpti_CallbackData *pCallabckInfo = (CUpti_CallbackData *)pCallbackData; + + switch (domain) + { + case CUPTI_CB_DOMAIN_STATE: + HandleDomainStateCallback(callbackId, (CUpti_StateData *)pCallbackData); + break; + case CUPTI_CB_DOMAIN_RUNTIME_API: + switch (callbackId) + { + case CUPTI_RUNTIME_TRACE_CBID_cudaDeviceReset_v3020: + if (pCallabckInfo->callbackSite == CUPTI_API_ENTER) + { + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(0)); + } + break; + default: + break; + } + break; + case CUPTI_CB_DOMAIN_SYNCHRONIZE: + HandleSyncronizationCallbacks(callbackId, (CUpti_SynchronizeData *)pCallbackData, pUserData); + break; + default: + break; + } +} + +// CUPTI Trace Setup +static void +InitCuptiTrace( + void *pUserData, + void *pTraceCallback, + FILE *pFileHandle) +{ + if (!pUserData) + { + std::cerr << "Invalid parameter pUserData.\n"; + exit(EXIT_FAILURE); + } + + globals.pOutputFile = pFileHandle; + globals.pUserData = pUserData; + + // Subscribe to CUPTI + if (((UserData *)pUserData)->skipCuptiSubscription == 0) + { + // If the user provides function pointer, subscribe CUPTI to that function pointer (pTraceCallback). + // Else subscribe CUPTI to the common CuptiCallbackHandler. + if (pTraceCallback) + { + CUPTI_API_CALL_VERBOSE(cuptiSubscribe(&globals.subscriberHandle, (CUpti_CallbackFunc)pTraceCallback, pUserData)); + } + else + { + CUPTI_API_CALL_VERBOSE(cuptiSubscribe(&globals.subscriberHandle, (CUpti_CallbackFunc)CuptiCallbackHandler, pUserData)); + } + + + // Enable CUPTI callback on context syncronization + if (((UserData *)pUserData)->flushAtCtxSync) + { + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, globals.subscriberHandle, CUPTI_CB_DOMAIN_SYNCHRONIZE, CUPTI_CBID_SYNCHRONIZE_CONTEXT_SYNCHRONIZED)); + } + + // Enable CUPTI callback on stream syncronization + if (((UserData *)pUserData)->flushAtStreamSync) + { + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, globals.subscriberHandle, CUPTI_CB_DOMAIN_SYNCHRONIZE, CUPTI_CBID_SYNCHRONIZE_STREAM_SYNCHRONIZED)); + } + + // Enable CUPTI callback on CUDA device reset by default + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, globals.subscriberHandle, CUPTI_CB_DOMAIN_RUNTIME_API, CUPTI_RUNTIME_TRACE_CBID_cudaDeviceReset_v3020)); + + // Enable CUPTI callback on fatal errors by default + CUPTI_API_CALL_VERBOSE(cuptiEnableCallback(1, globals.subscriberHandle, CUPTI_CB_DOMAIN_STATE, CUPTI_CBID_STATE_FATAL_ERROR)); + } + + // Register callbacks for buffer requests and for buffers completed by CUPTI. + globals.buffersRequested = 0; + globals.buffersCompleted = 0; + CUPTI_API_CALL_VERBOSE(cuptiActivityRegisterCallbacks(BufferRequested, BufferCompleted)); + + // Optionally get and set activity attributes. + // Attributes can be set by the CUPTI client to change behavior of the activity API. + // Some attributes require to be set before any CUDA context is created to be effective, + // E.g. To be applied to all device buffer allocations (see documentation). + if ((((UserData *)pUserData))->deviceBufferSize != 0) + { + size_t attrValue = (((UserData *)pUserData))->deviceBufferSize; + size_t attrValueSize = sizeof(size_t); + CUPTI_API_CALL_VERBOSE(cuptiActivitySetAttribute(CUPTI_ACTIVITY_ATTR_DEVICE_BUFFER_SIZE, &attrValueSize, &attrValue)); + std::cout << "CUPTI_ACTIVITY_ATTR_DEVICE_BUFFER_SIZE = " << attrValue << " bytes.\n"; + } + + if ((((UserData *)pUserData))->activityBufferSize != 0) + { + globals.activityBufferSize = (((UserData *)pUserData))->activityBufferSize; + } + else + { + globals.activityBufferSize = BUF_SIZE; + } + + std::cout << "Activity buffer size = " << globals.activityBufferSize << " bytes.\n"; +} + +static void +DeInitCuptiTrace(void) +{ + CUPTI_API_CALL(cuptiGetLastError()); + + if (globals.subscriberHandle) + { + CUPTI_API_CALL_VERBOSE(cuptiUnsubscribe(globals.subscriberHandle)); + } + + CUPTI_API_CALL_VERBOSE(cuptiActivityFlushAll(1)); + + if (globals.pUserData != NULL) + { + free(globals.pUserData); + } +} + +#endif // HELPER_CUPTI_ACTIVITY_H_ diff --git a/src/xpu/flamegraph/cupti_trace_parser.py b/src/xpu/flamegraph/cupti_trace_parser.py new file mode 100644 index 0000000..0ea4a73 --- /dev/null +++ b/src/xpu/flamegraph/cupti_trace_parser.py @@ -0,0 +1,314 @@ + +#!/usr/bin/env python3 +""" +CUPTI Trace Parser Module +Parses CUPTI trace data and converts to Chrome Trace Format +""" + +import re +import json +from typing import List, Dict, Any + + +class CuptiTraceParser: + """Parser for CUPTI trace data""" + + def __init__(self): + # Regular expressions for different trace line formats + self.runtime_pattern = r'RUNTIME \[ (\d+), (\d+) \] duration (\d+), "([^"]+)", cbid (\d+), processId (\d+), threadId (\d+), correlationId (\d+)' + self.driver_pattern = r'DRIVER \[ (\d+), (\d+) \] duration (\d+), "([^"]+)", cbid (\d+), processId (\d+), threadId (\d+), correlationId (\d+)' + self.kernel_pattern = r'CONCURRENT_KERNEL \[ (\d+), (\d+) \] duration (\d+), "([^"]+)", correlationId (\d+)' + self.overhead_pattern = r'OVERHEAD ([A-Z_]+) \[ (\d+), (\d+) \] duration (\d+), (\w+), id (\d+), correlation id (\d+)' + self.memory_pattern = r'MEMORY2 \[ (\d+) \] memoryOperationType (\w+), memoryKind (\w+), size (\d+), address (\d+)' + self.memcpy_pattern = r'MEMCPY "([^"]+)" \[ (\d+), (\d+) \] duration (\d+), size (\d+), copyCount (\d+), srcKind (\w+), dstKind (\w+), correlationId (\d+)' + self.grid_pattern = r'\s+grid \[ (\d+), (\d+), (\d+) \], block \[ (\d+), (\d+), (\d+) \]' + self.device_pattern = r'\s+deviceId (\d+), contextId (\d+), streamId (\d+)' + + def parse_file(self, filename: str) -> List[Dict[str, Any]]: + """Parse CUPTI trace file and return list of events""" + with open(filename, 'r') as f: + lines = f.readlines() + + return self.parse_lines(lines) + + def parse_lines(self, lines: List[str]) -> List[Dict[str, Any]]: + """Parse CUPTI trace lines and return list of events""" + events = [] + i = 0 + + while i < len(lines): + line = lines[i].strip() + + # Skip empty lines or non-trace lines + if not line or self._should_skip_line(line): + i += 1 + continue + + # Try parsing different event types + event = None + lines_consumed = 1 + + # Parse RUNTIME events + match = re.search(self.runtime_pattern, line) + if match: + event = self._parse_runtime_event(match) + else: + # Parse DRIVER events + match = re.search(self.driver_pattern, line) + if match: + event = self._parse_driver_event(match) + else: + # Parse CONCURRENT_KERNEL events + match = re.search(self.kernel_pattern, line) + if match: + event, lines_consumed = self._parse_kernel_event(match, lines, i) + else: + # Parse OVERHEAD events + match = re.search(self.overhead_pattern, line) + if match: + event = self._parse_overhead_event(match) + else: + # Parse MEMCPY events + match = re.search(self.memcpy_pattern, line) + if match: + event, lines_consumed = self._parse_memcpy_event(match, lines, i) + else: + # Parse MEMORY2 events + match = re.search(self.memory_pattern, line) + if match: + event = self._parse_memory_event(match) + + if event: + events.append(event) + + i += lines_consumed + + return events + + def _should_skip_line(self, line: str) -> bool: + """Check if line should be skipped""" + skip_prefixes = [ + 'Calling CUPTI', 'Enabling', 'Disabling', 'Found', + 'Configuring', 'It took', 'Activity buffer', 'CUPTI trace output', + 'Running command', 'Trace output:', 'Started target', + 'Starting CPU', 'Stopping CPU', 'CPU profile' + ] + return any(line.startswith(prefix) for prefix in skip_prefixes) + + def _parse_runtime_event(self, match) -> Dict[str, Any]: + """Parse RUNTIME event""" + start_time = int(match.group(1)) + duration = int(match.group(3)) + name = match.group(4) + cbid = match.group(5) + process_id = int(match.group(6)) + thread_id = int(match.group(7)) + correlation_id = int(match.group(8)) + + return { + "name": f"Runtime: {name}", + "ph": "X", # Complete event + "ts": start_time / 1000, # Convert ns to µs + "dur": duration / 1000, + "tid": thread_id, + "pid": process_id, + "cat": "CUDA_Runtime", + "args": { + "cbid": cbid, + "correlationId": correlation_id + } + } + + def _parse_driver_event(self, match) -> Dict[str, Any]: + """Parse DRIVER event""" + start_time = int(match.group(1)) + duration = int(match.group(3)) + name = match.group(4) + cbid = match.group(5) + process_id = int(match.group(6)) + thread_id = int(match.group(7)) + correlation_id = int(match.group(8)) + + return { + "name": f"Driver: {name}", + "ph": "X", + "ts": start_time / 1000, + "dur": duration / 1000, + "tid": thread_id, + "pid": process_id, + "cat": "CUDA_Driver", + "args": { + "cbid": cbid, + "correlationId": correlation_id + } + } + + def _parse_kernel_event(self, match, lines: List[str], current_index: int) -> tuple: + """Parse CONCURRENT_KERNEL event with optional additional info""" + start_time = int(match.group(1)) + duration = int(match.group(3)) + name = match.group(4) + correlation_id = int(match.group(5)) + + kernel_info = { + "name": f"Kernel: {name}", + "ph": "X", + "ts": start_time / 1000, + "dur": duration / 1000, + "cat": "GPU_Kernel", + "args": { + "correlationId": correlation_id + } + } + + lines_consumed = 1 + + # Check next lines for additional kernel info + if current_index + 1 < len(lines): + next_line = lines[current_index + 1].strip() + grid_match = re.search(self.grid_pattern, next_line) + if grid_match: + kernel_info["args"]["grid"] = [ + int(grid_match.group(1)), + int(grid_match.group(2)), + int(grid_match.group(3)) + ] + kernel_info["args"]["block"] = [ + int(grid_match.group(4)), + int(grid_match.group(5)), + int(grid_match.group(6)) + ] + lines_consumed += 1 + + if current_index + lines_consumed < len(lines): + next_line = lines[current_index + lines_consumed].strip() + device_match = re.search(self.device_pattern, next_line) + if device_match: + device_id = int(device_match.group(1)) + context_id = int(device_match.group(2)) + stream_id = int(device_match.group(3)) + + kernel_info["tid"] = f"GPU{device_id}_Stream{stream_id}" + kernel_info["pid"] = f"Device_{device_id}" + kernel_info["args"]["deviceId"] = device_id + kernel_info["args"]["contextId"] = context_id + kernel_info["args"]["streamId"] = stream_id + lines_consumed += 1 + + return kernel_info, lines_consumed + + def _parse_overhead_event(self, match) -> Dict[str, Any]: + """Parse OVERHEAD event""" + overhead_type = match.group(1) + start_time = int(match.group(2)) + duration = int(match.group(4)) + overhead_target = match.group(5) + overhead_id = int(match.group(6)) + correlation_id = int(match.group(7)) + + return { + "name": f"Overhead: {overhead_type}", + "ph": "X", + "ts": start_time / 1000, + "dur": duration / 1000, + "tid": overhead_id, + "pid": "CUPTI_Overhead", + "cat": "Overhead", + "args": { + "type": overhead_type, + "target": overhead_target, + "correlationId": correlation_id + } + } + + def _parse_memcpy_event(self, match, lines: List[str], current_index: int) -> tuple: + """Parse MEMCPY event with optional device info""" + copy_type = match.group(1) + start_time = int(match.group(2)) + duration = int(match.group(4)) + size = int(match.group(5)) + copy_count = int(match.group(6)) + src_kind = match.group(7) + dst_kind = match.group(8) + correlation_id = int(match.group(9)) + + memcpy_info = { + "name": f"MemCopy: {copy_type}", + "ph": "X", + "ts": start_time / 1000, + "dur": duration / 1000, + "cat": "MemCopy", + "args": { + "type": copy_type, + "size": size, + "copyCount": copy_count, + "srcKind": src_kind, + "dstKind": dst_kind, + "correlationId": correlation_id + } + } + + lines_consumed = 1 + + # Check next line for device info + if current_index + 1 < len(lines): + next_line = lines[current_index + 1].strip() + device_match = re.search(self.device_pattern, next_line) + if device_match: + device_id = int(device_match.group(1)) + context_id = int(device_match.group(2)) + stream_id = int(device_match.group(3)) + + memcpy_info["tid"] = f"GPU{device_id}_Stream{stream_id}" + memcpy_info["pid"] = f"Device_{device_id}" + memcpy_info["args"]["deviceId"] = device_id + memcpy_info["args"]["contextId"] = context_id + memcpy_info["args"]["streamId"] = stream_id + lines_consumed += 1 + else: + memcpy_info["tid"] = "MemCopy_Operations" + memcpy_info["pid"] = "MemCopy" + + return memcpy_info, lines_consumed + + def _parse_memory_event(self, match) -> Dict[str, Any]: + """Parse MEMORY2 event""" + timestamp = int(match.group(1)) + operation = match.group(2) + memory_kind = match.group(3) + size = int(match.group(4)) + address = int(match.group(5)) + + return { + "name": f"Memory: {operation} ({memory_kind})", + "ph": "i", # Instant event + "ts": timestamp / 1000, + "tid": "Memory_Operations", + "pid": "Memory", + "cat": "Memory", + "s": "g", # Global scope + "args": { + "operation": operation, + "kind": memory_kind, + "size": size, + "address": hex(address) + } + } + + def to_chrome_trace(self, events: List[Dict[str, Any]], metadata: Dict[str, Any] = None) -> Dict[str, Any]: + """Convert events to Chrome Trace Format""" + trace_data = { + "traceEvents": events, + "displayTimeUnit": "ms", + "metadata": metadata or { + "tool": "CUPTI Trace Parser", + "format": "Chrome Trace Format" + } + } + return trace_data + + def save_chrome_trace(self, events: List[Dict[str, Any]], output_file: str, metadata: Dict[str, Any] = None): + """Save events as Chrome Trace Format JSON""" + trace_data = self.to_chrome_trace(events, metadata) + with open(output_file, 'w') as f: + json.dump(trace_data, f, indent=2) \ No newline at end of file diff --git a/src/xpu/flamegraph/gpuperf.py b/src/xpu/flamegraph/gpuperf.py new file mode 100755 index 0000000..d848f7a --- /dev/null +++ b/src/xpu/flamegraph/gpuperf.py @@ -0,0 +1,415 @@ +#!/usr/bin/env python3 + +import os +import sys +import argparse +import subprocess +import tempfile +import atexit +import time +import json +from pathlib import Path +from cupti_trace_parser import CuptiTraceParser +from merge_gpu_cpu_trace import TraceMerger + +class GPUPerf: + def __init__(self): + self.script_dir = Path(__file__).parent.absolute() + self.injection_lib = self.script_dir / "cupti_trace/libcupti_trace_injection.so" + self.output_file = None + self.temp_trace_file = None + self.profiler_proc = None + self.profiler_output = None + self.parser = CuptiTraceParser() # Initialize the parser + + # Path to CPU profiler + script_dir = Path(__file__).parent.resolve() + self.cpu_profiler = script_dir / "profiler/target/release/profile" + if not self.cpu_profiler.exists(): + print(f"Warning: CPU profiler not found at {self.cpu_profiler}", file=sys.stderr) + self.cpu_profiler = None + + # Find CUPTI library path + cuda_paths = [ + "/usr/local/cuda-13.0/extras/CUPTI/lib64", + "/usr/local/cuda/extras/CUPTI/lib64", + "/usr/local/cuda-12.0/extras/CUPTI/lib64", + ] + + self.cupti_lib = None + for path in cuda_paths: + cupti_path = Path(path) / "libcupti.so" + if cupti_path.exists(): + self.cupti_lib = str(cupti_path) + self.cupti_lib_dir = str(Path(path)) + break + + if not self.cupti_lib: + print("Warning: Could not find CUPTI library. NVTX annotations may not work.", file=sys.stderr) + + def parse_cupti_trace(self, filename): + """Parse CUPTI trace data using the parser module""" + return self.parser.parse_file(filename) + + def start_cpu_profiler(self, pid=None, cpu_output_file=None, cuda_lib_path=None): + """Start CPU profiler with cudaLaunchKernel uprobe""" + if not self.cpu_profiler: + return None + + if not cpu_output_file: + cpu_output_file = f"cpu_profile_{pid if pid else 'cuda'}.txt" + + self.profiler_output = cpu_output_file + + # Find CUDA runtime library if not specified + if not cuda_lib_path: + cuda_paths = [ + "/usr/local/cuda-12.9/lib64/libcudart.so.12", + "/usr/local/cuda-13.0/lib64/libcudart.so.12", + "/usr/local/cuda/lib64/libcudart.so.12", + "/usr/local/cuda-12.8/lib64/libcudart.so.12", + ] + for path in cuda_paths: + if Path(path).exists(): + cuda_lib_path = path + break + + if not cuda_lib_path: + print("Warning: Could not find CUDA runtime library for uprobe", file=sys.stderr) + return None + + print(f"Starting CPU profiler with cudaLaunchKernel hook") + print(f" CUDA library: {cuda_lib_path}") + print(f" Output: {cpu_output_file}") + + try: + # Run profiler with cudaLaunchKernel uprobe in extended folded format + # Format: timestamp_ns comm pid tid cpu stack1;stack2;stack3 + cmd = ["sudo", str(self.cpu_profiler), + "--uprobe", f"{cuda_lib_path}:cudaLaunchKernel", + "-E"] # -E for extended folded format with timestamps + + self.profiler_proc = subprocess.Popen( + cmd, + stdout=open(cpu_output_file, 'w'), + stderr=subprocess.PIPE + ) + # Give it a moment to attach + time.sleep(1.0) + return self.profiler_proc + except Exception as e: + print(f"Warning: Failed to start CPU profiler: {e}", file=sys.stderr) + return None + + def stop_cpu_profiler(self): + """Stop the CPU profiler gracefully""" + if self.profiler_proc and self.profiler_proc.poll() is None: + print("Stopping CPU profiler...") + self.profiler_proc.terminate() + try: + self.profiler_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + self.profiler_proc.kill() + self.profiler_proc.wait() + + if self.profiler_output and os.path.exists(self.profiler_output): + print(f"CPU profile saved to: {self.profiler_output}") + + def run_with_trace(self, command, output_trace=None, chrome_trace=None, cpu_profile=None, merged_trace=None, no_merge=False): + """Run a command with CUPTI tracing and optional CPU profiling enabled""" + + # Determine if we're doing GPU profiling + do_gpu_profiling = output_trace is not None or chrome_trace is not None + + # Check if injection library exists (only if we're doing GPU profiling) + if do_gpu_profiling and not self.injection_lib.exists(): + print(f"Error: CUPTI injection library not found at {self.injection_lib}", file=sys.stderr) + print("Please build it first using 'make' in the cupti_trace directory", file=sys.stderr) + return 1 + + # Set up trace output file for GPU profiling + trace_file = None + if do_gpu_profiling: + if output_trace: + trace_file = output_trace + else: + # Create temporary file for trace output + fd, trace_file = tempfile.mkstemp(suffix=".txt", prefix="gpuperf_trace_") + os.close(fd) + self.temp_trace_file = trace_file + atexit.register(self.cleanup_temp_files) + + # Set up environment variables + env = os.environ.copy() + env['CUDA_INJECTION64_PATH'] = str(self.injection_lib) + env['CUPTI_TRACE_OUTPUT_FILE'] = trace_file + + if self.cupti_lib: + env['NVTX_INJECTION64_PATH'] = self.cupti_lib + if 'LD_LIBRARY_PATH' in env: + env['LD_LIBRARY_PATH'] = f"{self.cupti_lib_dir}:{env['LD_LIBRARY_PATH']}" + else: + env['LD_LIBRARY_PATH'] = self.cupti_lib_dir + + print(f"Running command with GPU profiling: {' '.join(command)}") + print(f"Trace output: {trace_file}") + + # Start the target process + target_proc = None + + try: + # Start CPU profiler FIRST if available and requested + if cpu_profile and self.cpu_profiler: + # Start profiler BEFORE target process to catch all kernel launches + self.start_cpu_profiler(cpu_output_file=cpu_profile) + + # Then start the target process + target_proc = subprocess.Popen(command, env=env) + target_pid = target_proc.pid + print(f"Started target process with PID: {target_pid}") + + # Wait for the target process to complete + return_code = target_proc.wait() + + except KeyboardInterrupt: + print("\nInterrupted by user") + if target_proc: + target_proc.terminate() + try: + target_proc.wait(timeout=5) + except subprocess.TimeoutExpired: + target_proc.kill() + return_code = 130 + except Exception as e: + print(f"Error running command: {e}", file=sys.stderr) + return_code = 1 + finally: + # Stop CPU profiler if running + self.stop_cpu_profiler() + + # Convert to Chrome trace if requested + if chrome_trace and os.path.exists(trace_file): + print(f"\nConverting trace to Chrome format: {chrome_trace}") + try: + events = self.parse_cupti_trace(trace_file) + print(f"Parsed {len(events)} events") + + metadata = { + "tool": "gpuperf - GPU Performance Profiler", + "format": "Chrome Trace Format", + "command": ' '.join(command) + } + + self.parser.save_chrome_trace(events, chrome_trace, metadata) + + print(f"\nChrome trace file written to: {chrome_trace}") + print("\nTo visualize the trace:") + print("1. Open Chrome or Edge browser") + print("2. Navigate to chrome://tracing or edge://tracing") + print("3. Click 'Load' and select the generated JSON file") + print("\nAlternatively, visit https://ui.perfetto.dev/ and drag the JSON file there") + except Exception as e: + print(f"Error converting trace: {e}", file=sys.stderr) + + # Clean up temporary file if not keeping raw trace + if not output_trace and self.temp_trace_file: + try: + os.unlink(self.temp_trace_file) + except: + pass + + # Generate merged folded trace if both CPU and GPU traces are available (and not disabled) + if not no_merge and cpu_profile and (chrome_trace or output_trace): + merged_output = merged_trace if merged_trace else "merged_trace.folded" + self.generate_merged_trace( + cpu_trace=cpu_profile, + gpu_trace=chrome_trace if chrome_trace else None, + gpu_raw_trace=trace_file if do_gpu_profiling else None, + output_file=merged_output + ) + + return return_code + + def generate_merged_trace(self, cpu_trace=None, gpu_trace=None, gpu_raw_trace=None, output_file=None): + """Generate merged CPU+GPU folded trace using TraceMerger""" + if not cpu_trace or not (gpu_trace or gpu_raw_trace): + return # Need both CPU and GPU traces + + if not output_file: + output_file = "merged_trace.folded" + + print(f"\nGenerating merged CPU+GPU trace: {output_file}") + + try: + merger = TraceMerger() + + # Parse CPU trace + if os.path.exists(cpu_trace): + merger.parse_cpu_trace(cpu_trace) + else: + print(f"Warning: CPU trace not found: {cpu_trace}") + return + + # Parse GPU trace (prefer JSON, fallback to raw) + if gpu_trace and os.path.exists(gpu_trace): + merger.parse_gpu_trace(gpu_trace) + elif gpu_raw_trace and os.path.exists(gpu_raw_trace): + # Convert raw trace to events first + events = self.parse_cupti_trace(gpu_raw_trace) + # Create temporary JSON for merger + import json + temp_json = tempfile.NamedTemporaryFile(mode='w', suffix='.json', delete=False) + json.dump({"traceEvents": events}, temp_json) + temp_json.close() + merger.parse_gpu_trace(temp_json.name) + os.unlink(temp_json.name) + else: + print(f"Warning: GPU trace not found") + return + + # Merge traces + merger.merge_traces() + + # Write folded output + merger.write_folded_output(output_file) + + print(f"✓ Merged trace generated: {output_file}") + print(f"\nTo generate flamegraph:") + print(f" /root/yunwei37/systemscope/cpu-tools/combined_flamegraph.pl {output_file} > merged_flamegraph.svg") + + except Exception as e: + print(f"Error generating merged trace: {e}", file=sys.stderr) + + def cleanup_temp_files(self): + """Clean up temporary files""" + if self.temp_trace_file and os.path.exists(self.temp_trace_file): + try: + os.unlink(self.temp_trace_file) + except: + pass + + def convert_trace(self, input_file, output_file): + """Convert existing CUPTI trace to Chrome format""" + + if not os.path.exists(input_file): + print(f"Error: Input file '{input_file}' not found", file=sys.stderr) + return 1 + + print(f"Converting CUPTI trace to Chrome format...") + print(f"Input: {input_file}") + print(f"Output: {output_file}") + + try: + events = self.parse_cupti_trace(input_file) + print(f"Parsed {len(events)} events") + + metadata = { + "tool": "gpuperf - GPU Performance Profiler", + "format": "Chrome Trace Format" + } + + self.parser.save_chrome_trace(events, output_file, metadata) + + print(f"\nChrome trace file written to: {output_file}") + print("\nTo visualize the trace:") + print("1. Open Chrome or Edge browser") + print("2. Navigate to chrome://tracing or edge://tracing") + print("3. Click 'Load' and select the generated JSON file") + print("\nAlternatively, visit https://ui.perfetto.dev/ and drag the JSON file there") + + return 0 + except Exception as e: + print(f"Error converting trace: {e}", file=sys.stderr) + return 1 + +def main(): + # Check if first argument is 'convert' for conversion mode + if len(sys.argv) > 1 and sys.argv[1] == 'convert': + parser = argparse.ArgumentParser( + prog='gpuperf convert', + description='Convert existing CUPTI trace to Chrome format' + ) + parser.add_argument('mode', help='Operation mode') # This will be 'convert' + parser.add_argument('-i', '--input', required=True, help='Input CUPTI trace file') + parser.add_argument('-o', '--output', default='trace.json', help='Output Chrome trace JSON file') + args = parser.parse_args() + + profiler = GPUPerf() + return profiler.convert_trace(args.input, args.output) + + # Regular run mode + parser = argparse.ArgumentParser( + description='gpuperf - GPU and CPU Performance Profiler', + usage='gpuperf [options] command [args...]\n gpuperf convert -i input.txt -o output.json' + ) + + parser.add_argument('-o', '--output', help='Save raw CUPTI trace to file (default: gpu_results.txt)') + parser.add_argument('-c', '--chrome', help='Convert trace to Chrome format and save to file (default: gpu_results.json)') + parser.add_argument('-p', '--cpu-profile', help='Also capture CPU profile and save to file (default: cpu_results.txt)') + parser.add_argument('-m', '--merged', help='Save merged CPU+GPU folded trace (default: merged_trace.folded)') + parser.add_argument('--cpu-only', action='store_true', help='Only run CPU profiler without GPU tracing') + parser.add_argument('--no-gpu', action='store_true', help='Disable GPU profiling') + parser.add_argument('--no-cpu', action='store_true', help='Disable CPU profiling') + parser.add_argument('--no-merge', action='store_true', help='Disable automatic merged trace generation') + parser.add_argument('command', nargs=argparse.REMAINDER, help='Command to run with profiling') + + args = parser.parse_args() + + profiler = GPUPerf() + + # Handle run mode + if not args.command: + parser.print_help() + return 1 + + # Use the command directly from REMAINDER + full_command = args.command + + # CPU-only mode + if args.cpu_only: + if not profiler.cpu_profiler: + print("Error: CPU profiler not available", file=sys.stderr) + return 1 + + # Start the process and immediately profile it + try: + target_proc = subprocess.Popen(full_command) + target_pid = target_proc.pid + print(f"Started target process with PID: {target_pid}") + + cpu_output = args.cpu_profile or "cpu_results.txt" + profiler.start_cpu_profiler(target_pid, cpu_output) + + return_code = target_proc.wait() + profiler.stop_cpu_profiler() + return return_code + except Exception as e: + print(f"Error: {e}", file=sys.stderr) + return 1 + + # Set up default values + gpu_output = args.output if args.output else ("gpu_results.txt" if not args.no_gpu else None) + chrome_output = args.chrome if args.chrome else ("gpu_results.json" if not args.no_gpu else None) + cpu_output = args.cpu_profile if args.cpu_profile else ("cpu_results.txt" if not args.no_cpu else None) + + # If user explicitly disabled GPU, don't run GPU profiling + if args.no_gpu: + gpu_output = None + chrome_output = None + + # If user explicitly disabled CPU, don't run CPU profiling + if args.no_cpu: + cpu_output = None + + # Combined GPU and CPU profiling (or just one based on flags) + return profiler.run_with_trace( + full_command, + output_trace=gpu_output, + chrome_trace=chrome_output, + cpu_profile=cpu_output, + merged_trace=args.merged, + no_merge=args.no_merge + ) + +if __name__ == '__main__': + sys.exit(main()) \ No newline at end of file diff --git a/src/xpu/flamegraph/merge_gpu_cpu_trace.py b/src/xpu/flamegraph/merge_gpu_cpu_trace.py new file mode 100755 index 0000000..39557d2 --- /dev/null +++ b/src/xpu/flamegraph/merge_gpu_cpu_trace.py @@ -0,0 +1,343 @@ +#!/usr/bin/env python3 +""" +Merge GPU and CPU traces into folded flamegraph format +Correlates CPU stack traces from cudaLaunchKernel uprobe with GPU kernel execution +using CUPTI correlation IDs and timestamp matching +""" + +import json +import re +import sys +import argparse +from pathlib import Path +from typing import List, Dict, Tuple, Any, Optional +from collections import defaultdict + + +class GPUKernelEvent: + """Represents a GPU kernel execution event""" + def __init__(self, name: str, start_ns: int, end_ns: int, correlation_id: int): + self.name = name + self.start_ns = start_ns + self.end_ns = end_ns + self.correlation_id = correlation_id + + def __repr__(self): + return f"GPUKernel({self.name}, {self.start_ns}-{self.end_ns}, corr={self.correlation_id})" + + +class CudaLaunchEvent: + """Represents a cudaLaunchKernel runtime API call""" + def __init__(self, start_ns: int, end_ns: int, correlation_id: int): + self.start_ns = start_ns + self.end_ns = end_ns + self.correlation_id = correlation_id + + def __repr__(self): + return f"CudaLaunch({self.start_ns}-{self.end_ns}, corr={self.correlation_id})" + + +class CPUStack: + """Represents a CPU stack trace from cudaLaunchKernel uprobe in extended folded format""" + def __init__(self, timestamp_ns: int, comm: str, pid: int, tid: int, cpu: int, stack: List[str]): + self.timestamp_ns = timestamp_ns + self.comm = comm + self.pid = pid + self.tid = tid + self.cpu = cpu + self.stack = stack # List of function names from bottom to top + + def __repr__(self): + return f"CPUStack({self.timestamp_ns}, pid={self.pid}, tid={self.tid}, depth={len(self.stack)})" + + +class TraceMerger: + """Merges GPU CUPTI traces with CPU stack traces from cudaLaunchKernel hooks""" + + def __init__(self, timestamp_tolerance_ms=10.0): + self.gpu_kernels = [] # List of GPUKernelEvent + self.cuda_launches = {} # correlation_id -> CudaLaunchEvent + self.cpu_stacks = [] # List of CPUStack from uprobe (extended folded format) + self.merged_stacks = defaultdict(int) # stack_string -> count + self.timestamp_tolerance_ns = int(timestamp_tolerance_ms * 1_000_000) + + def parse_cpu_trace(self, cpu_file: str): + """Parse CPU trace file in extended folded format from Rust profiler""" + print(f"Parsing CPU uprobe trace (extended folded format): {cpu_file}") + + with open(cpu_file, 'r') as f: + lines = f.readlines() + + stack_count = 0 + for line in lines: + line = line.strip() + if not line: + continue + + # Extended folded format: timestamp_ns comm pid tid cpu stack1;stack2;stack3 + parts = line.split(None, 5) # Split on whitespace, max 6 parts + if len(parts) < 6: + continue + + try: + timestamp_ns = int(parts[0]) + comm = parts[1] + pid = int(parts[2]) + tid = int(parts[3]) + cpu = int(parts[4]) + stack_str = parts[5] + + # Parse stack frames (separated by semicolons) + stack_frames = [] + seen_cuda_launch = False + if stack_str: + frames = stack_str.split(';') + for frame in frames: + frame = frame.strip() + if frame and frame not in ['', '_start', '__libc_start_main']: + # Clean up cudaLaunchKernel variations - keep only first occurrence + if 'cudaLaunchKernel' in frame or '__device_stub__' in frame: + if not seen_cuda_launch: + frame = 'cudaLaunchKernel' + stack_frames.append(frame) + seen_cuda_launch = True + else: + stack_frames.append(frame) + + if stack_frames: + self.cpu_stacks.append(CPUStack( + timestamp_ns, comm, pid, tid, cpu, stack_frames + )) + stack_count += 1 + + except (ValueError, IndexError) as e: + print(f"Warning: Failed to parse line: {line[:100]}... Error: {e}") + continue + + print(f"Parsed {stack_count} CPU stack traces from cudaLaunchKernel hooks") + + def parse_gpu_trace(self, gpu_json_file: str): + """Parse GPU trace JSON file and extract kernel events and launch correlations""" + print(f"Parsing GPU CUPTI trace: {gpu_json_file}") + + with open(gpu_json_file, 'r') as f: + data = json.load(f) + + events = data.get('traceEvents', []) + kernel_count = 0 + launch_count = 0 + + for event in events: + name = event.get('name', '') + category = event.get('cat', '') + correlation_id = event.get('args', {}).get('correlationId', 0) + + # Extract cudaLaunchKernel runtime events + if category == 'CUDA_Runtime' and 'LaunchKernel' in name: + start_us = event.get('ts', 0) + duration_us = event.get('dur', 0) + + if start_us > 0 and duration_us > 0 and correlation_id > 0: + start_ns = int(start_us * 1000) + end_ns = int((start_us + duration_us) * 1000) + + self.cuda_launches[correlation_id] = CudaLaunchEvent( + start_ns, end_ns, correlation_id + ) + launch_count += 1 + + # Extract actual GPU kernel executions + elif category == 'GPU_Kernel' or name.startswith('Kernel:'): + kernel_name = name.replace('Kernel: ', '') + start_us = event.get('ts', 0) + duration_us = event.get('dur', 0) + + if start_us > 0 and duration_us > 0 and correlation_id > 0: + start_ns = int(start_us * 1000) + end_ns = int((start_us + duration_us) * 1000) + + self.gpu_kernels.append(GPUKernelEvent( + kernel_name, + start_ns, + end_ns, + correlation_id + )) + kernel_count += 1 + + # Sort by correlation ID for efficient lookup + self.gpu_kernels.sort(key=lambda k: k.correlation_id) + + print(f"Parsed {kernel_count} GPU kernel events") + print(f"Parsed {launch_count} cudaLaunchKernel runtime events") + + def find_matching_kernel(self, cpu_stack: CPUStack) -> Optional[GPUKernelEvent]: + """ + Find GPU kernel that matches the CPU stack trace. + Strategy: + 1. Find cudaLaunchKernel runtime call within timestamp tolerance + 2. Use correlation ID to find actual GPU kernel execution + """ + + # Find cudaLaunchKernel runtime event that matches timestamp + best_launch = None + min_time_diff = self.timestamp_tolerance_ns + + for launch in self.cuda_launches.values(): + # Check if CPU stack timestamp is close to launch time + time_diff = abs(cpu_stack.timestamp_ns - launch.start_ns) + + if time_diff < min_time_diff: + min_time_diff = time_diff + best_launch = launch + + if not best_launch: + return None + + # Find GPU kernel with matching correlation ID + for kernel in self.gpu_kernels: + if kernel.correlation_id == best_launch.correlation_id: + return kernel + + return None + + def merge_traces(self): + """Correlate CPU stacks with GPU kernels using correlation IDs and timestamps""" + print("Correlating CPU stacks with GPU kernels...") + + matched_count = 0 + unmatched_count = 0 + + for cpu_stack in self.cpu_stacks: + # Find matching GPU kernel + gpu_kernel = self.find_matching_kernel(cpu_stack) + + # Build merged stack + merged_stack = cpu_stack.stack.copy() + + if gpu_kernel: + # Add GPU kernel to the top of the stack + merged_stack.append(f"[GPU_Kernel]{gpu_kernel.name}") + matched_count += 1 + else: + # Mark as unmatched launch (may happen if kernel hasn't executed yet) + merged_stack.append("[GPU_Launch_Pending]") + unmatched_count += 1 + + # Create folded stack string + if merged_stack: + stack_str = ';'.join(merged_stack) + self.merged_stacks[stack_str] += 1 + + print(f"Matched {matched_count} CPU stacks with GPU kernels") + print(f"Unmatched: {unmatched_count}") + print(f"Total unique stacks: {len(self.merged_stacks)}") + + def write_folded_output(self, output_file: str): + """Write folded stack format for flamegraph generation""" + print(f"Writing folded output to: {output_file}") + + with open(output_file, 'w') as f: + for stack, count in sorted(self.merged_stacks.items()): + # Folded format: stack_frame1;stack_frame2;... count + f.write(f"{stack} {count}\n") + + total_samples = sum(self.merged_stacks.values()) + print(f"Wrote {len(self.merged_stacks)} unique stacks ({total_samples} total samples)") + + def generate_summary(self): + """Generate summary statistics""" + print("\n=== Summary Statistics ===") + + # CPU statistics + if self.cpu_stacks: + cpu_start = min(s.timestamp_ns for s in self.cpu_stacks) + cpu_end = max(s.timestamp_ns for s in self.cpu_stacks) + cpu_duration_ms = (cpu_end - cpu_start) / 1_000_000 + print(f"CPU trace duration: {cpu_duration_ms:.2f} ms") + print(f"CPU stacks captured: {len(self.cpu_stacks)}") + + # GPU statistics + if self.gpu_kernels: + print(f"\nGPU kernels executed: {len(self.gpu_kernels)}") + print(f"CUDA launch events: {len(self.cuda_launches)}") + + total_kernel_time = sum(k.end_ns - k.start_ns for k in self.gpu_kernels) / 1_000_000 + print(f"Total kernel execution time: {total_kernel_time:.2f} ms") + + # Show kernel breakdown + kernel_names = defaultdict(int) + for k in self.gpu_kernels: + kernel_names[k.name] += 1 + + print("\nKernel execution counts:") + for name, count in sorted(kernel_names.items(), key=lambda x: -x[1]): + print(f" {name}: {count}") + + +def main(): + parser = argparse.ArgumentParser( + description='Merge GPU CUPTI traces with CPU cudaLaunchKernel stack traces' + ) + parser.add_argument( + '-c', '--cpu', + default='cpu_results.txt', + help='CPU uprobe trace file (extended folded format, default: cpu_results.txt)' + ) + parser.add_argument( + '-g', '--gpu', + default='gpu_results.json', + help='GPU CUPTI trace JSON file (default: gpu_results.json)' + ) + parser.add_argument( + '-o', '--output', + default='merged_trace.folded', + help='Output folded stack file (default: merged_trace.folded)' + ) + parser.add_argument( + '-t', '--tolerance', + type=float, + default=10.0, + help='Timestamp matching tolerance in milliseconds (default: 10.0)' + ) + parser.add_argument( + '-s', '--summary', + action='store_true', + help='Print summary statistics' + ) + + args = parser.parse_args() + + # Check input files exist + if not Path(args.cpu).exists(): + print(f"Error: CPU trace file not found: {args.cpu}", file=sys.stderr) + sys.exit(1) + + if not Path(args.gpu).exists(): + print(f"Error: GPU trace file not found: {args.gpu}", file=sys.stderr) + sys.exit(1) + + # Create merger and process traces + merger = TraceMerger(timestamp_tolerance_ms=args.tolerance) + + # Parse inputs + merger.parse_cpu_trace(args.cpu) + merger.parse_gpu_trace(args.gpu) + + # Merge traces + merger.merge_traces() + + # Write output + merger.write_folded_output(args.output) + + # Print summary if requested + if args.summary: + merger.generate_summary() + + print(f"\nTo generate flamegraph:") + print(f" flamegraph.pl {args.output} > merged_flamegraph.svg") + print(f"\nOr use online viewer:") + print(f" https://www.speedscope.app/ (upload {args.output})") + + +if __name__ == '__main__': + main() diff --git a/src/xpu/flamegraph/mock-test/.gitignore b/src/xpu/flamegraph/mock-test/.gitignore new file mode 100644 index 0000000..95b3447 --- /dev/null +++ b/src/xpu/flamegraph/mock-test/.gitignore @@ -0,0 +1,3 @@ +llm-inference +*.o +*.bak diff --git a/src/xpu/flamegraph/mock-test/Makefile b/src/xpu/flamegraph/mock-test/Makefile new file mode 100644 index 0000000..42ec753 --- /dev/null +++ b/src/xpu/flamegraph/mock-test/Makefile @@ -0,0 +1,53 @@ +# +# Makefile for mock LLM inference application +# +ifndef OS + OS := $(shell uname) + HOST_ARCH := $(shell uname -m) +endif + +CUDA_INSTALL_PATH ?= /usr/local/cuda-12.9 +NVCC := "$(CUDA_INSTALL_PATH)/bin/nvcc" +INCLUDES := -I"$(CUDA_INSTALL_PATH)/include" + +ifeq ($(OS),Windows_NT) + LIB_PATH ?= ..\..\lib64 +else + LIB_PATH ?= $(CUDA_INSTALL_PATH)/lib64 +endif + +# Point to the necessary cross-compiler. +NVCCFLAGS := + +ifneq ($(TARGET_ARCH), $(HOST_ARCH)) + ifeq ($(TARGET_ARCH), aarch64) + ifeq ($(TARGET_OS), linux) + HOST_COMPILER ?= aarch64-linux-gnu-g++ + else ifeq ($(TARGET_OS),qnx) + ifeq ($(QNX_HOST),) + $(error ERROR - QNX_HOST must be passed to the QNX host toolchain) + endif + ifeq ($(QNX_TARGET),) + $(error ERROR - QNX_TARGET must be passed to the QNX target toolchain) + endif + HOST_COMPILER ?= $(QNX_HOST)/usr/bin/q++ + ifndef QPP_CONFIG_VERSION + QPP_CONFIG_VERSION = 12.2.0 + endif + $(info QPP_CONFIG_VERSION = $(QPP_CONFIG_VERSION)) + NVCCFLAGS += --qpp-config $(QPP_CONFIG_VERSION),gcc_ntoaarch64le -lsocket + endif + endif + + ifdef HOST_COMPILER + NVCC_COMPILER := -ccbin $(HOST_COMPILER) + endif +endif + +all: llm-inference + +llm-inference: llm-inference.cu + $(NVCC) $(NVCC_COMPILER) $(INCLUDES) -o llm-inference llm-inference.cu -L $(LIB_PATH) -lcudart -std=c++17 -Wno-deprecated-gpu-targets --no-device-link + +clean: + rm -f llm-inference *.o *.bak diff --git a/src/xpu/flamegraph/mock-test/llm-inference.cu b/src/xpu/flamegraph/mock-test/llm-inference.cu new file mode 100644 index 0000000..685f933 --- /dev/null +++ b/src/xpu/flamegraph/mock-test/llm-inference.cu @@ -0,0 +1,702 @@ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// ============================================================================= +// Configuration using constexpr +// ============================================================================= +namespace Config { + constexpr size_t BATCH_SIZE = 16; + constexpr size_t SEQ_LENGTH = 1024; + constexpr size_t HIDDEN_DIM = 2048; + constexpr size_t NUM_HEADS = 16; + constexpr size_t HEAD_DIM = HIDDEN_DIM / NUM_HEADS; + constexpr size_t FFN_DIM = HIDDEN_DIM * 4; + constexpr size_t NUM_LAYERS = 4; + constexpr size_t VOCAB_SIZE = 4000; + constexpr int DURATION_SECONDS = 10; +} + +// ============================================================================= +// CUDA Error Checking Wrapper +// ============================================================================= +class CudaError : public std::runtime_error { +public: + explicit CudaError(const std::string& msg) : std::runtime_error(msg) {} +}; + +inline void checkCuda(cudaError_t result, const char* file, int line) { + if (result != cudaSuccess) { + throw CudaError(std::string("CUDA Error: ") + + cudaGetErrorString(result) + + " at " + file + ":" + std::to_string(line)); + } +} + +#define CUDA_CHECK(call) checkCuda((call), __FILE__, __LINE__) + +// ============================================================================= +// RAII CUDA Memory Wrapper +// ============================================================================= +template +class CudaDeviceMemory { +private: + T* data_ = nullptr; + size_t size_ = 0; + +public: + explicit CudaDeviceMemory(size_t count) : size_(count) { + if (count > 0) { + CUDA_CHECK(cudaMalloc(&data_, count * sizeof(T))); + std::cout << "[CUDA] Allocated " << (count * sizeof(T)) / (1024.0 * 1024.0) + << " MB on device" << std::endl; + } + } + + ~CudaDeviceMemory() { + if (data_) { + cudaFree(data_); + } + } + + // Delete copy operations + CudaDeviceMemory(const CudaDeviceMemory&) = delete; + CudaDeviceMemory& operator=(const CudaDeviceMemory&) = delete; + + // Allow move operations + CudaDeviceMemory(CudaDeviceMemory&& other) noexcept + : data_(other.data_), size_(other.size_) { + other.data_ = nullptr; + other.size_ = 0; + } + + CudaDeviceMemory& operator=(CudaDeviceMemory&& other) noexcept { + if (this != &other) { + if (data_) cudaFree(data_); + data_ = other.data_; + size_ = other.size_; + other.data_ = nullptr; + other.size_ = 0; + } + return *this; + } + + T* get() { return data_; } + const T* get() const { return data_; } + size_t size() const { return size_; } + + void copyFromHost(const std::vector& host_data) { + if (host_data.size() != size_) { + throw std::runtime_error("Size mismatch in copyFromHost"); + } + CUDA_CHECK(cudaMemcpy(data_, host_data.data(), + size_ * sizeof(T), cudaMemcpyHostToDevice)); + } + + void copyToHost(std::vector& host_data) const { + if (host_data.size() != size_) { + host_data.resize(size_); + } + CUDA_CHECK(cudaMemcpy(host_data.data(), data_, + size_ * sizeof(T), cudaMemcpyDeviceToHost)); + } + + void zero() { + CUDA_CHECK(cudaMemset(data_, 0, size_ * sizeof(T))); + } +}; + +// ============================================================================= +// CUDA Stream Wrapper +// ============================================================================= +class CudaStream { +private: + cudaStream_t stream_ = nullptr; + +public: + CudaStream() { + CUDA_CHECK(cudaStreamCreate(&stream_)); + } + + ~CudaStream() { + if (stream_) { + cudaStreamDestroy(stream_); + } + } + + CudaStream(const CudaStream&) = delete; + CudaStream& operator=(const CudaStream&) = delete; + + cudaStream_t get() const { return stream_; } + + void synchronize() { + CUDA_CHECK(cudaStreamSynchronize(stream_)); + } +}; + +// ============================================================================= +// GPU Kernels +// ============================================================================= +__global__ void attentionQKTKernel(const float* Q, const float* K, float* scores, + size_t batch, size_t seq_len, size_t head_dim) { + size_t b = blockIdx.z; + size_t i = blockIdx.y * blockDim.y + threadIdx.y; + size_t j = blockIdx.x * blockDim.x + threadIdx.x; + + if (b < batch && i < seq_len && j < seq_len) { + float sum = 0.0f; + for (size_t k = 0; k < head_dim; k++) { + size_t q_idx = b * seq_len * head_dim + i * head_dim + k; + size_t k_idx = b * seq_len * head_dim + j * head_dim + k; + sum += Q[q_idx] * K[k_idx]; + } + scores[b * seq_len * seq_len + i * seq_len + j] = sum / sqrtf(static_cast(head_dim)); + } +} + +__global__ void softmaxKernel(const float* input, float* output, size_t batch, size_t seq_len) { + size_t b = blockIdx.y; + size_t i = blockIdx.x * blockDim.x + threadIdx.x; + + if (b < batch && i < seq_len) { + float max_val = -INFINITY; + for (size_t j = 0; j < seq_len; j++) { + size_t idx = b * seq_len * seq_len + i * seq_len + j; + max_val = fmaxf(max_val, input[idx]); + } + + float sum = 0.0f; + for (size_t j = 0; j < seq_len; j++) { + size_t idx = b * seq_len * seq_len + i * seq_len + j; + output[idx] = expf(input[idx] - max_val); + sum += output[idx]; + } + + for (size_t j = 0; j < seq_len; j++) { + size_t idx = b * seq_len * seq_len + i * seq_len + j; + output[idx] /= sum; + } + } +} + +__global__ void layerNormKernel(const float* input, float* output, + const float* gamma, const float* beta, + size_t batch, size_t seq_len, size_t hidden_dim) { + size_t b = blockIdx.y; + size_t i = blockIdx.x * blockDim.x + threadIdx.x; + + if (b < batch && i < seq_len) { + float mean = 0.0f; + for (size_t j = 0; j < hidden_dim; j++) { + mean += input[b * seq_len * hidden_dim + i * hidden_dim + j]; + } + mean /= hidden_dim; + + float variance = 0.0f; + for (size_t j = 0; j < hidden_dim; j++) { + float diff = input[b * seq_len * hidden_dim + i * hidden_dim + j] - mean; + variance += diff * diff; + } + variance /= hidden_dim; + + float std = sqrtf(variance + 1e-5f); + for (size_t j = 0; j < hidden_dim; j++) { + size_t idx = b * seq_len * hidden_dim + i * hidden_dim + j; + output[idx] = gamma[j] * (input[idx] - mean) / std + beta[j]; + } + } +} + +__global__ void residualAddKernel(const float* input, const float* residual, + float* output, size_t n) { + size_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx < n) { + output[idx] = input[idx] + residual[idx]; + } +} + +// ============================================================================= +// Token Embedding using modern C++ +// ============================================================================= +class TokenEmbedding { +private: + std::vector embeddings_; + size_t vocab_size_; + size_t embedding_dim_; + std::mt19937 rng_; + std::uniform_real_distribution dist_; + +public: + TokenEmbedding(size_t vocab_size, size_t embedding_dim) + : vocab_size_(vocab_size) + , embedding_dim_(embedding_dim) + , rng_(std::random_device{}()) + , dist_(-1.0f, 1.0f) { + + embeddings_.resize(vocab_size * embedding_dim); + std::cout << "[Init] Creating TokenEmbedding: " + << (embeddings_.size() * sizeof(float)) / (1024.0 * 1024.0) + << " MB" << std::endl; + + // Initialize with random values + for (auto& val : embeddings_) { + val = dist_(rng_); + } + } + + void embed(const std::vector& tokens, std::vector& output) const { + // Output should be sized for full batch + size_t required_size = Config::BATCH_SIZE * Config::SEQ_LENGTH * embedding_dim_; + output.resize(required_size); + std::fill(output.begin(), output.end(), 0.0f); + + // Fill first sequence with actual embeddings + for (size_t i = 0; i < tokens.size() && i < Config::SEQ_LENGTH; ++i) { + int token_id = tokens[i] % vocab_size_; + size_t src_offset = token_id * embedding_dim_; + size_t dst_offset = i * embedding_dim_; + + std::copy_n(embeddings_.begin() + src_offset, + embedding_dim_, + output.begin() + dst_offset); + } + } + + size_t getEmbeddingDim() const { return embedding_dim_; } +}; + +// ============================================================================= +// Transformer Layer using RAII +// ============================================================================= +class TransformerLayer { +private: + CudaDeviceMemory d_Q_; + CudaDeviceMemory d_K_; + CudaDeviceMemory d_V_; + CudaDeviceMemory d_attn_scores_; + CudaDeviceMemory d_attn_probs_; + CudaDeviceMemory d_attn_output_; + CudaDeviceMemory d_ln_gamma_; + CudaDeviceMemory d_ln_beta_; + CudaDeviceMemory d_residual_; + + std::vector h_gamma_; + std::vector h_beta_; + CudaStream stream_; + +public: + TransformerLayer() + : d_Q_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HEAD_DIM) + , d_K_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HEAD_DIM) + , d_V_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HEAD_DIM) + , d_attn_scores_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::SEQ_LENGTH) + , d_attn_probs_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::SEQ_LENGTH) + , d_attn_output_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HEAD_DIM) + , d_ln_gamma_(Config::HIDDEN_DIM) + , d_ln_beta_(Config::HIDDEN_DIM) + , d_residual_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM) + , h_gamma_(Config::HIDDEN_DIM, 1.0f) + , h_beta_(Config::HIDDEN_DIM, 0.0f) { + + std::cout << "[Init] Creating TransformerLayer" << std::endl; + + d_ln_gamma_.copyFromHost(h_gamma_); + d_ln_beta_.copyFromHost(h_beta_); + } + + void forward(const CudaDeviceMemory& d_input, + CudaDeviceMemory& d_output) { + + // Do multiple passes to increase GPU compute time + // Pass 1: Layer norm + dim3 ln_grid((Config::SEQ_LENGTH + 255) / 256, Config::BATCH_SIZE); + layerNormKernel<<>>( + d_input.get(), d_residual_.get(), + d_ln_gamma_.get(), d_ln_beta_.get(), + Config::BATCH_SIZE, Config::SEQ_LENGTH, Config::HIDDEN_DIM); + + // Pass 2: Multiple softmax iterations to increase GPU compute + dim3 softmax_grid((Config::SEQ_LENGTH + 255) / 256, Config::BATCH_SIZE); + for (int i = 0; i < 22; ++i) { // Tuned to 22 iterations for ~50% GPU + softmaxKernel<<>>( + d_attn_scores_.get(), d_attn_probs_.get(), + Config::BATCH_SIZE, Config::SEQ_LENGTH); + } + + // Pass 3: Residual add + size_t total_elements = Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM; + for (int i = 0; i < 2; ++i) { + residualAddKernel<<<(total_elements + 255) / 256, 256, 0, stream_.get()>>>( + d_residual_.get(), d_input.get(), d_output.get(), total_elements); + } + + // Pass 4: Multiple layer norm passes + for (int i = 0; i < 2; ++i) { + layerNormKernel<<>>( + d_output.get(), d_residual_.get(), + d_ln_gamma_.get(), d_ln_beta_.get(), + Config::BATCH_SIZE, Config::SEQ_LENGTH, Config::HIDDEN_DIM); + } + + // Pass 5: Final residual + residualAddKernel<<<(total_elements + 255) / 256, 256, 0, stream_.get()>>>( + d_residual_.get(), d_input.get(), d_output.get(), total_elements); + + stream_.synchronize(); + } +}; + +// ============================================================================= +// File Cache Manager +// ============================================================================= +class PromptCache { +private: + std::string cache_dir_; + std::vector cached_files_; + +public: + PromptCache() { + cache_dir_ = "/tmp/llm_cache_" + std::to_string(getpid()); + std::string cmd = "mkdir -p " + cache_dir_; + std::system(cmd.c_str()); + std::cout << "[Init] Cache directory: " << cache_dir_ << std::endl; + } + + ~PromptCache() { + cleanup(); + } + + void writeCache(const std::string& key, const std::vector& data, int iteration) { + std::string filename = cache_dir_ + "/cache_" + key + "_" + std::to_string(iteration) + ".bin"; + std::ofstream file(filename, std::ios::binary); + if (file) { + file.write(reinterpret_cast(data.data()), + data.size() * sizeof(float)); + cached_files_.push_back(filename); + } + } + + bool readCache(const std::string& key, std::vector& data, int iteration) { + std::string filename = cache_dir_ + "/cache_" + key + "_" + std::to_string(iteration) + ".bin"; + std::ifstream file(filename, std::ios::binary); + if (!file) return false; + + file.seekg(0, std::ios::end); + size_t size = file.tellg() / sizeof(float); + file.seekg(0, std::ios::beg); + + data.resize(size); + file.read(reinterpret_cast(data.data()), size * sizeof(float)); + return true; + } + + void cleanup() { + for (const auto& file : cached_files_) { + std::remove(file.c_str()); + } + std::string cmd = "rm -rf " + cache_dir_; + std::system(cmd.c_str()); + } +}; + +// ============================================================================= +// Performance Timing Statistics +// ============================================================================= +struct RequestTimings { + double cpu_compute_ms = 0.0; + double gpu_compute_ms = 0.0; + double io_time_ms = 0.0; + + void add(const RequestTimings& other) { + cpu_compute_ms += other.cpu_compute_ms; + gpu_compute_ms += other.gpu_compute_ms; + io_time_ms += other.io_time_ms; + } + + double total_ms() const { + return cpu_compute_ms + gpu_compute_ms + io_time_ms; + } +}; + +// ============================================================================= +// Main Inference Pipeline +// ============================================================================= +class InferencePipeline { +private: + std::unique_ptr embedding_; + std::vector> layers_; + std::unique_ptr cache_; + + CudaDeviceMemory d_input_; + CudaDeviceMemory d_output_; + + std::vector h_input_; + std::vector h_output_; + + // Performance tracking + std::vector request_timings_; + RequestTimings accumulated_timings_; + int request_count_ = 0; + + std::array prompts_ = { + "What is artificial intelligence?", + "Explain transformer architectures", + "Describe deep learning techniques", + "What are neural networks?", + "How does machine learning work?" + }; + +public: + InferencePipeline() + : embedding_(std::make_unique(Config::VOCAB_SIZE, Config::HIDDEN_DIM)) + , cache_(std::make_unique()) + , d_input_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM) + , d_output_(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM) { + + std::cout << "[Init] Creating InferencePipeline with " + << Config::NUM_LAYERS << " layers" << std::endl; + + // Create transformer layers + for (size_t i = 0; i < Config::NUM_LAYERS; ++i) { + std::cout << "[Init] Creating layer " << (i + 1) << "/" + << Config::NUM_LAYERS << std::endl; + layers_.push_back(std::make_unique()); + } + + h_input_.resize(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM); + h_output_.resize(Config::BATCH_SIZE * Config::SEQ_LENGTH * Config::HIDDEN_DIM); + + std::cout << "[Init] Pipeline initialization complete" << std::endl; + } + + void runRequest(int request_id) { + RequestTimings timings; + auto start_time = std::chrono::high_resolution_clock::now(); + + // Select prompt + const auto& prompt = prompts_[request_id % prompts_.size()]; + + // ===== CPU COMPUTE: Tokenization ===== + auto cpu_start = std::chrono::high_resolution_clock::now(); + std::vector tokens; + tokens.reserve(Config::SEQ_LENGTH); + for (size_t i = 0; i < Config::SEQ_LENGTH && i < prompt.length(); ++i) { + tokens.push_back(static_cast(prompt[i])); + } + while (tokens.size() < Config::SEQ_LENGTH) { + tokens.push_back(0); // Padding + } + + // ===== CPU COMPUTE: Embedding lookup ===== + embedding_->embed(tokens, h_input_); + + // ===== CPU COMPUTE: Additional preprocessing (to increase CPU time) ===== + // Simulate text preprocessing, normalization, etc. + std::vector temp_buffer(Config::SEQ_LENGTH * 150); // Increased buffer + for (size_t i = 0; i < temp_buffer.size(); ++i) { + temp_buffer[i] = std::sin(static_cast(i)) * std::cos(static_cast(request_id)); + } + + // Simulate some CPU-intensive work (sorting, searching, etc.) + for (int iter = 0; iter < 12; ++iter) { // Tuned to 12 iterations for ~25% CPU + std::partial_sort(temp_buffer.begin(), temp_buffer.begin() + 1500, temp_buffer.end()); + } + + auto cpu_end = std::chrono::high_resolution_clock::now(); + timings.cpu_compute_ms = std::chrono::duration(cpu_end - cpu_start).count(); + + // ===== I/O: Transfer to GPU ===== + auto io_start = std::chrono::high_resolution_clock::now(); + d_input_.copyFromHost(h_input_); + auto io_end = std::chrono::high_resolution_clock::now(); + timings.io_time_ms += std::chrono::duration(io_end - io_start).count(); + + // ===== GPU COMPUTE: Forward pass through transformer layers ===== + auto gpu_start = std::chrono::high_resolution_clock::now(); + auto* current_input = &d_input_; + auto* current_output = &d_output_; + + for (auto& layer : layers_) { + layer->forward(*current_input, *current_output); + std::swap(current_input, current_output); + } + auto gpu_end = std::chrono::high_resolution_clock::now(); + timings.gpu_compute_ms = std::chrono::duration(gpu_end - gpu_start).count(); + + // ===== I/O: Transfer back to CPU ===== + io_start = std::chrono::high_resolution_clock::now(); + current_input->copyToHost(h_output_); + io_end = std::chrono::high_resolution_clock::now(); + timings.io_time_ms += std::chrono::duration(io_end - io_start).count(); + + // ===== I/O: Cache results (file I/O) ===== + if (request_id % 2 == 0) { + io_start = std::chrono::high_resolution_clock::now(); + cache_->writeCache("prompt_" + std::to_string(request_id % prompts_.size()), + h_output_, request_id); + io_end = std::chrono::high_resolution_clock::now(); + timings.io_time_ms += std::chrono::duration(io_end - io_start).count(); + } + + // ===== I/O: Simulate network delay ===== + io_start = std::chrono::high_resolution_clock::now(); + std::this_thread::sleep_for(std::chrono::milliseconds(10)); // Reduced from 50ms to 10ms + io_end = std::chrono::high_resolution_clock::now(); + timings.io_time_ms += std::chrono::duration(io_end - io_start).count(); + + // Track timings + request_timings_.push_back(timings); + accumulated_timings_.add(timings); + request_count_++; + + // Report every 10 requests + if (request_count_ % 10 == 0) { + reportTimings(request_count_); + } + } + + void reportTimings(int last_request_id) { + // Calculate statistics for last 10 requests + size_t start_idx = request_timings_.size() >= 10 ? request_timings_.size() - 10 : 0; + RequestTimings last_10; + + for (size_t i = start_idx; i < request_timings_.size(); ++i) { + last_10.add(request_timings_[i]); + } + + int count = request_timings_.size() - start_idx; + double avg_cpu = last_10.cpu_compute_ms / count; + double avg_gpu = last_10.gpu_compute_ms / count; + double avg_io = last_10.io_time_ms / count; + double avg_total = (avg_cpu + avg_gpu + avg_io); + + std::cout << "\n[Performance Report] Requests " << (last_request_id - count + 1) + << " - " << last_request_id << " (last " << count << " requests):" << std::endl; + std::cout << " CPU Compute: " << std::fixed << std::setprecision(2) + << avg_cpu << " ms (" << (avg_cpu / avg_total * 100) << "%)" << std::endl; + std::cout << " GPU Compute: " << avg_gpu << " ms (" + << (avg_gpu / avg_total * 100) << "%)" << std::endl; + std::cout << " I/O (+ Net): " << avg_io << " ms (" + << (avg_io / avg_total * 100) << "%)" << std::endl; + std::cout << " Total Time: " << avg_total << " ms/request" << std::endl; + } + + void printFinalReport() { + if (request_count_ == 0) return; + + std::cout << "\n=============================================================" << std::endl; + std::cout << "Final Performance Report (" << request_count_ << " total requests)" << std::endl; + std::cout << "=============================================================" << std::endl; + + double avg_cpu = accumulated_timings_.cpu_compute_ms / request_count_; + double avg_gpu = accumulated_timings_.gpu_compute_ms / request_count_; + double avg_io = accumulated_timings_.io_time_ms / request_count_; + double avg_total = (avg_cpu + avg_gpu + avg_io); + + std::cout << "Average per request:" << std::endl; + std::cout << " CPU Compute: " << std::fixed << std::setprecision(2) + << avg_cpu << " ms (" << (avg_cpu / avg_total * 100) << "%)" << std::endl; + std::cout << " GPU Compute: " << avg_gpu << " ms (" + << (avg_gpu / avg_total * 100) << "%)" << std::endl; + std::cout << " I/O (+ Net): " << avg_io << " ms (" + << (avg_io / avg_total * 100) << "%)" << std::endl; + std::cout << " Total Time: " << avg_total << " ms/request" << std::endl; + std::cout << "\nTotal time breakdown:" << std::endl; + std::cout << " CPU Compute: " << accumulated_timings_.cpu_compute_ms << " ms" << std::endl; + std::cout << " GPU Compute: " << accumulated_timings_.gpu_compute_ms << " ms" << std::endl; + std::cout << " I/O (+ Net): " << accumulated_timings_.io_time_ms << " ms" << std::endl; + std::cout << "=============================================================" << std::endl; + } +}; + +// ============================================================================= +// Global cleanup handler +// ============================================================================= +std::unique_ptr g_pipeline; +volatile sig_atomic_t g_interrupted = 0; + +void signalHandler(int signum) { + std::cout << "\n[Signal] Received signal " << signum << ", cleaning up..." << std::endl; + g_interrupted = 1; + g_pipeline.reset(); + std::cout << "[Cleanup] Complete. Exiting." << std::endl; + exit(signum); +} + +// ============================================================================= +// Main +// ============================================================================= +int main() { + try { + std::cout << "=============================================================" << std::endl; + std::cout << "Modern C++ LLM Inference Simulator" << std::endl; + std::cout << "=============================================================" << std::endl; + std::cout << "Configuration:" << std::endl; + std::cout << " - Batch Size: " << Config::BATCH_SIZE << std::endl; + std::cout << " - Sequence Length: " << Config::SEQ_LENGTH << std::endl; + std::cout << " - Hidden Dimension: " << Config::HIDDEN_DIM << std::endl; + std::cout << " - Number of Layers: " << Config::NUM_LAYERS << std::endl; + std::cout << " - Duration: " << Config::DURATION_SECONDS << " seconds" << std::endl; + std::cout << "=============================================================" << std::endl; + + // Initialize CUDA + CUDA_CHECK(cudaSetDevice(0)); + std::cout << "[Init] CUDA device initialized" << std::endl; + + // Setup signal handlers + signal(SIGINT, signalHandler); + signal(SIGTERM, signalHandler); + + // Create pipeline + g_pipeline = std::make_unique(); + + // Run request processing loop + auto start = std::chrono::steady_clock::now(); + int request_id = 0; + + std::cout << "\n[Starting] Processing requests for " << Config::DURATION_SECONDS + << " seconds..." << std::endl; + + while (!g_interrupted) { + auto now = std::chrono::steady_clock::now(); + auto elapsed = std::chrono::duration_cast(now - start).count(); + + if (elapsed >= Config::DURATION_SECONDS) { + break; + } + + g_pipeline->runRequest(request_id); + request_id++; + } + + std::cout << "\n=============================================================" << std::endl; + std::cout << "Completed " << request_id << " requests in " + << Config::DURATION_SECONDS << " seconds" << std::endl; + std::cout << "Average throughput: " + << (request_id / static_cast(Config::DURATION_SECONDS)) + << " requests/second" << std::endl; + std::cout << "=============================================================" << std::endl; + + // Print final performance report + g_pipeline->printFinalReport(); + + g_pipeline.reset(); + + return 0; + + } catch (const std::exception& e) { + std::cerr << "[ERROR] " << e.what() << std::endl; + return 1; + } +} diff --git a/src/xpu/flamegraph/profiler/.cargo/config.toml b/src/xpu/flamegraph/profiler/.cargo/config.toml new file mode 100644 index 0000000..eb9fed2 --- /dev/null +++ b/src/xpu/flamegraph/profiler/.cargo/config.toml @@ -0,0 +1,2 @@ +[build] +rustflags = ["-C", "linker=gcc"] diff --git a/src/xpu/flamegraph/profiler/.gitignore b/src/xpu/flamegraph/profiler/.gitignore new file mode 100644 index 0000000..3a1b448 --- /dev/null +++ b/src/xpu/flamegraph/profiler/.gitignore @@ -0,0 +1,2 @@ +/src/bpf/.output +/target diff --git a/src/xpu/flamegraph/profiler/Cargo.lock b/src/xpu/flamegraph/profiler/Cargo.lock new file mode 100644 index 0000000..5135b87 --- /dev/null +++ b/src/xpu/flamegraph/profiler/Cargo.lock @@ -0,0 +1,909 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "aho-corasick" +version = "1.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e60d3430d3a69478ad0993f19238d2df97c507009a52b3c10addcd7f6bcb916" +dependencies = [ + "memchr", +] + +[[package]] +name = "anstream" +version = "0.6.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ae563653d1938f79b1ab1b5e668c87c76a9930414574a6583a7b7e11a8e6192" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "862ed96ca487e809f1c8e5a8447f6ee2cf102f846893800b20cebdf541fc6bbd" + +[[package]] +name = "anstyle-parse" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e231f6134f61b71076a3eab506c379d4f36122f2af15a9ff04415ea4c3339e2" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e0633414522a32ffaac8ac6cc8f748e090c5717661fddeea04219e2344f5f2a" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys", +] + +[[package]] +name = "anyhow" +version = "1.0.99" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0674a1ddeecb70197781e945de4b3b8ffb61fa939a5597bcf48503737663100" + +[[package]] +name = "bitflags" +version = "2.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34efbcccd345379ca2868b2b2c9d3782e9cc58ba87bc7d79d5b53d9c9ae6f25d" + +[[package]] +name = "blazesym" +version = "0.2.0-rc.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29a810b7e5f883ad3c711208237841f051061bf59b6ee698ac4dc1fe12a3a5db" +dependencies = [ + "cpp_demangle", + "gimli", + "libc", + "memmap2 0.9.8", + "miniz_oxide", + "rustc-demangle", + "tracing", +] + +[[package]] +name = "camino" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d07aa9a93b00c76f71bc35d598bed923f6d4f3a9ca5c24b7737ae1a292841c0" +dependencies = [ + "serde", +] + +[[package]] +name = "cargo-platform" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e35af189006b9c0f00a064685c727031e3ed2d8020f7ba284d78cc2671bd36ea" +dependencies = [ + "serde", +] + +[[package]] +name = "cargo_metadata" +version = "0.15.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eee4243f1f26fc7a42710e7439c149e2b10b05472f88090acce52632f231a73a" +dependencies = [ + "camino", + "cargo-platform", + "semver", + "serde", + "serde_json", + "thiserror", +] + +[[package]] +name = "cc" +version = "1.2.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42bc4aea80032b7bf409b0bc7ccad88853858911b7713a8062fdc0623867bedc" +dependencies = [ + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2fd1289c04a9ea8cb22300a459a72a385d7c73d3259e2ed7dcb2af674838cfa9" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "clap" +version = "4.5.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c5e4fcf9c21d2e544ca1ee9d8552de13019a42aa7dbf32747fa7aaf1df76e57" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.5.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fecb53a0e6fcfb055f686001bc2e2592fa527efaf38dbe81a6a9563562e57d41" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.5.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14cb31bb0a7d536caef2639baa7fad459e15c3144efefa6dbd1c84562c4739f6" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b94f61472cee1439c0b966b47e3aca9ae07e45d070759512cd390ea2bebc6675" + +[[package]] +name = "colorchoice" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b05b61dc5112cbb17e4b6cd61790d9845d13888356391624cbe7e41efeac1e75" + +[[package]] +name = "cpp_demangle" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96e58d342ad113c2b878f16d5d034c03be492ae460cdbc02b7f0f2284d310c7d" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778e2ac28f6c47af28e4907f13ffd1e1ddbd400980a9abd7c8df189bf578a5ad" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + +[[package]] +name = "getrandom" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26145e563e54f2cadc477553f1ec5ee650b00862f0a58bcd12cbdc5f0ea2d2f4" +dependencies = [ + "cfg-if", + "libc", + "r-efi", + "wasi", +] + +[[package]] +name = "gimli" +version = "0.32.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc6298e594375a7fead9efd5568f0a46e6a154fb6a9bdcbe3c06946ffd81a5f6" +dependencies = [ + "fallible-iterator", + "indexmap", + "stable_deref_trait", +] + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "indexmap" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2481980430f9f78649238835720ddccc57e52df14ffce1c6f37391d61b563e9" +dependencies = [ + "equivalent", + "hashbrown", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7943c866cc5cd64cbc25b2e01621d07fa8eb2a1a23160ee81ce38704e97b8ecf" + +[[package]] +name = "itoa" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a5f13b858c8d314ee3e8f639011f7ccefe71f97f96e50151fb991f267928e2c" + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "libbpf-cargo" +version = "0.24.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704727a07f185a76c58faa7b8ed08fba3194661c212183aea1174fe2970ee185" +dependencies = [ + "anyhow", + "cargo_metadata", + "clap", + "libbpf-rs", + "memmap2 0.5.10", + "regex", + "semver", + "serde", + "serde_json", + "tempfile", +] + +[[package]] +name = "libbpf-rs" +version = "0.24.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93edd9cd673087fa7518fd63ad6c87be2cd9b4e35034b1873f3e3258c018275b" +dependencies = [ + "bitflags", + "libbpf-sys", + "libc", + "vsprintf", +] + +[[package]] +name = "libbpf-sys" +version = "1.6.1+v1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e351855cbd724ac341b2a1c163568808e72acd930c491a921331c2e5347390d3" +dependencies = [ + "cc", + "nix 0.30.1", + "pkg-config", +] + +[[package]] +name = "libc" +version = "0.2.175" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a82ae493e598baaea5209805c49bbf2ea7de956d50d7da0da1164f9c6d28543" + +[[package]] +name = "linux-raw-sys" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd945864f07fe9f5371a27ad7b52a172b4b499999f1d97574c9fa68373937e12" + +[[package]] +name = "log" +version = "0.4.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13dc2df351e3202783a1fe0d44375f7295ffb4049267b0f3018346dc122a1d94" + +[[package]] +name = "matchers" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8263075bb86c5a1b1427b5ae862e8889656f126e9f77c484496e8b47cf5c5558" +dependencies = [ + "regex-automata 0.1.10", +] + +[[package]] +name = "memchr" +version = "2.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a282da65faaf38286cf3be983213fcf1d2e2a58700e808f83f4ea9a4804bc0" + +[[package]] +name = "memmap2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83faa42c0a078c393f6b29d5db232d8be22776a891f8f56e5284faee4a20b327" +dependencies = [ + "libc", +] + +[[package]] +name = "memmap2" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843a98750cd611cc2965a8213b53b43e715f13c37a9e096c6408e69990961db7" +dependencies = [ + "libc", +] + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "nix" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[package]] +name = "nix" +version = "0.30.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "74523f3a35e05aba87a1d978330aef40f67b0304ac79c1c00b294c9830543db6" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[package]] +name = "nu-ansi-term" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77a8165726e8236064dbb45459242600304b42a5ea24ee2948e18e023bf7ba84" +dependencies = [ + "overload", + "winapi", +] + +[[package]] +name = "once_cell" +version = "1.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42f5e15c9953c5e4ccceeb2e7382a716482c34515315f7b03532b8b4e8393d2d" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4895175b425cb1f87721b59f0f286c2092bd4af812243672510e1ac53e2e0ad" + +[[package]] +name = "overload" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b15813163c1d831bf4a13c3610c05c0d03b39feb07f7e09fa234dac9b15aaf39" + +[[package]] +name = "pin-project-lite" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b3cff922bd51709b605d9ead9aa71031d81447142d828eb4a6eba76fe619f9b" + +[[package]] +name = "pkg-config" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" + +[[package]] +name = "proc-macro2" +version = "1.0.101" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89ae43fd86e4158d6db51ad8e2b80f313af9cc74f5c0e03ccb87de09998732de" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "profile" +version = "0.1.0" +dependencies = [ + "blazesym", + "clap", + "libbpf-cargo", + "libbpf-rs", + "libc", + "nix 0.29.0", + "tracing", + "tracing-subscriber", +] + +[[package]] +name = "quote" +version = "1.0.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1885c039570dc00dcb4ff087a89e185fd56bae234ddc7f056a945bf36467248d" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "regex" +version = "1.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23d7fd106d8c02486a8d64e778353d1cffe08ce79ac2e82f540c86d0facf6912" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata 0.4.10", + "regex-syntax 0.8.6", +] + +[[package]] +name = "regex-automata" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c230d73fb8d8c1b9c0b3135c5142a8acee3a0558fb8db5cf1cb65f8d7862132" +dependencies = [ + "regex-syntax 0.6.29", +] + +[[package]] +name = "regex-automata" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b9458fa0bfeeac22b5ca447c63aaf45f28439a709ccd244698632f9aa6394d6" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax 0.8.6", +] + +[[package]] +name = "regex-syntax" +version = "0.6.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f162c6dd7b008981e4d40210aca20b4bd0f9b60ca9271061b07f78537722f2e1" + +[[package]] +name = "regex-syntax" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "caf4aa5b0f434c91fe5c7f1ecb6a5ece2130b02ad2a590589dda5146df959001" + +[[package]] +name = "rustc-demangle" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56f7d92ca342cea22a06f2121d944b4fd82af56988c270852495420f961d4ace" + +[[package]] +name = "rustix" +version = "1.0.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11181fbabf243db407ef8df94a6ce0b2f9a733bd8be4ad02b4eda9602296cac8" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "ryu" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28d3b2b1366ec20994f1fd18c3c594f05c5dd4bc44d8bb0c1c632c8d6829481f" + +[[package]] +name = "semver" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56e6fa9c48d24d85fb3de5ad847117517440f6beceb7798af16b4a87d616b8d0" +dependencies = [ + "serde", +] + +[[package]] +name = "serde" +version = "1.0.219" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f0e2c6ed6606019b4e29e69dbaba95b11854410e5347d525002456dbbb786b6" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.219" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b0276cf7f2c73365f7157c8123c21cd9a50fbbd844757af28ca1f5925fc2a00" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.143" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d401abef1d108fbd9cbaebc3e46611f4b1021f714a0597a71f41ee463f5f4a5a" +dependencies = [ + "itoa", + "memchr", + "ryu", + "serde", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "simd-adler32" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d66dc143e6b11c1eddc06d5c423cfc97062865baf299914ab64caa38182078fe" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "stable_deref_trait" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a8f112729512f8e442d81f95a8a7ddf2b7c6b8a1a6f509a95864142b30cab2d3" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "syn" +version = "2.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ede7c438028d4436d71104916910f5bb611972c5cfd7f89b8300a8186e6fada6" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "tempfile" +version = "3.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15b61f8f20e3a6f7e0649d825294eaf317edce30f82cf6026e7e4cb9222a7d1e" +dependencies = [ + "fastrand", + "getrandom", + "once_cell", + "rustix", + "windows-sys", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tracing" +version = "0.1.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "784e0ac535deb450455cbfa28a6f0df145ea1bb7ae51b821cf5e7927fdcfbdd0" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81383ab64e72a7a8b8e13130c49e3dab29def6d0c7d76a03087b3cf71c5c6903" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d12581f227e93f094d3af2ae690a574abb8a2b9b7a96e7cfe9647b2b617678" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "unicode-ident" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a5f39404a5da50712a4c1eecf25e90dd62b613502b7e925fd4e4d19b5c96512" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "vsprintf" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aec2f81b75ca063294776b4f7e8da71d1d5ae81c2b1b149c8d89969230265d63" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "wasi" +version = "0.14.2+wasi-0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9683f9a5a998d873c0d21fcbe3c083009670149a8fab228644b8bd36b2c48cb3" +dependencies = [ + "wit-bindgen-rt", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-link" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6ad25900d524eaabdbbb96d20b4311e1e7ae1699af4fb28c17ae66c80d798a" + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-targets" +version = "0.53.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5fe6031c4041849d7c496a8ded650796e7b6ecc19df1a431c1a363342e5dc91" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86b8d5f90ddd19cb4a147a5fa63ca848db3df085e25fee3cc10b39b6eebae764" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7651a1f62a11b8cbd5e0d42526e55f2c99886c77e007179efff86c2b137e66c" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1dc67659d35f387f5f6c479dc4e28f1d4bb90ddd1a5d3da2e5d97b42d6272c3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ce6ccbdedbf6d6354471319e781c0dfef054c81fbc7cf83f338a4296c0cae11" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "581fee95406bb13382d2f65cd4a908ca7b1e4c2f1917f143ba16efe98a589b5d" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e55b5ac9ea33f2fc1716d1742db15574fd6fc8dadc51caab1c16a3d3b4190ba" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a6e035dd0599267ce1ee132e51c27dd29437f63325753051e71dd9e42406c57" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "271414315aff87387382ec3d271b52d7ae78726f5d44ac98b4f4030c91880486" + +[[package]] +name = "wit-bindgen-rt" +version = "0.39.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f42320e61fe2cfd34354ecb597f86f413484a798ba44a8ca1165c58d42da6c1" +dependencies = [ + "bitflags", +] diff --git a/src/xpu/flamegraph/profiler/Cargo.toml b/src/xpu/flamegraph/profiler/Cargo.toml new file mode 100644 index 0000000..78d5e1c --- /dev/null +++ b/src/xpu/flamegraph/profiler/Cargo.toml @@ -0,0 +1,19 @@ +[package] +name = "profile" +version = "0.1.0" +authors = ["Kuifeng Lee "] +license = "GPL-2.0 OR BSD-3-Clause" +edition = "2021" +rust-version = "1.71" + +[dependencies] +blazesym = { version = "0.2.0-rc.4",features = ["tracing"] } +clap = { version = "4.5", features = ["derive"] } +libbpf-rs = "0.24" +libc = "*" +nix = "0.29.0" +tracing = "0.1" +tracing-subscriber = {version = "0.3", features = ["ansi", "env-filter", "fmt"]} + +[build-dependencies] +libbpf-cargo = "0.24" diff --git a/src/xpu/flamegraph/profiler/build.rs b/src/xpu/flamegraph/profiler/build.rs new file mode 100644 index 0000000..2b8bdbc --- /dev/null +++ b/src/xpu/flamegraph/profiler/build.rs @@ -0,0 +1,30 @@ +use std::env; +use std::ffi::OsStr; +use std::path::Path; +use std::path::PathBuf; + +use libbpf_cargo::SkeletonBuilder; + +const SRC: &str = "src/bpf/profile.bpf.c"; + +fn main() { + let mut out = + PathBuf::from(env::var_os("OUT_DIR").expect("OUT_DIR must be set in build script")); + out.push("profile.skel.rs"); + + let arch = env::var("CARGO_CFG_TARGET_ARCH") + .expect("CARGO_CFG_TARGET_ARCH must be set in build script"); + + // Ensure we're building for a supported architecture + println!("cargo:warning=Building for architecture: {}", arch); + + SkeletonBuilder::new() + .source(SRC) + .clang_args([ + OsStr::new("-I"), + Path::new("../vmlinux").as_os_str() + ]) + .build_and_generate(out) + .expect("bpf compilation failed"); + println!("cargo:rerun-if-changed={}", SRC); +} diff --git a/src/xpu/flamegraph/profiler/src/bpf/profile.bpf.c b/src/xpu/flamegraph/profiler/src/bpf/profile.bpf.c new file mode 100644 index 0000000..7896e59 --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/bpf/profile.bpf.c @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: GPL-2.0 OR BSD-3-Clause +/* Copyright (c) 2022 Meta Platforms, Inc. */ +#include "vmlinux.h" +#include +#include +#include + +#include "profile.h" + +char LICENSE[] SEC("license") = "Dual BSD/GPL"; + +struct { + __uint(type, BPF_MAP_TYPE_RINGBUF); + __uint(max_entries, 256 * 1024); +} events SEC(".maps"); + +// Shared helper to collect stack trace +static __always_inline int collect_stack_trace(void *ctx, u64 cookie) +{ + int pid = bpf_get_current_pid_tgid() >> 32; + int cpu_id = bpf_get_smp_processor_id(); + struct stacktrace_event *event; + + event = bpf_ringbuf_reserve(&events, sizeof(*event), 0); + if (!event) + return 1; + + event->pid = pid; + event->cpu_id = cpu_id; + event->timestamp = bpf_ktime_get_ns(); + + if (bpf_get_current_comm(event->comm, sizeof(event->comm))) + event->comm[0] = 0; + + // Store probe_id in cpu_id field when in probe mode + // In perf mode: cpu_id is actual CPU + // In probe mode: cpu_id is probe_id, actual CPU stored in pid high bits if needed + if (cookie != 0) { + event->cpu_id = (u32)cookie; // probe_id from bpf_get_attach_cookie + } + + event->kstack_sz = bpf_get_stack(ctx, event->kstack, sizeof(event->kstack), 0); + + event->ustack_sz = + bpf_get_stack(ctx, event->ustack, sizeof(event->ustack), BPF_F_USER_STACK); + + bpf_ringbuf_submit(event, 0); + + return 0; +} + +SEC("perf_event") +int profile(void *ctx) +{ + return collect_stack_trace(ctx, 0); +} + +// Generic kprobe handler +SEC("kprobe") +int kprobe_handler(struct pt_regs *ctx) +{ + u64 probe_id = bpf_get_attach_cookie(ctx); + return collect_stack_trace(ctx, probe_id); +} + +// Generic kretprobe handler +SEC("kretprobe") +int kretprobe_handler(struct pt_regs *ctx) +{ + u64 probe_id = bpf_get_attach_cookie(ctx); + return collect_stack_trace(ctx, probe_id); +} + +// Generic uprobe handler +SEC("uprobe") +int uprobe_handler(struct pt_regs *ctx) +{ + u64 probe_id = bpf_get_attach_cookie(ctx); + return collect_stack_trace(ctx, probe_id); +} + +// Generic uretprobe handler +SEC("uretprobe") +int uretprobe_handler(struct pt_regs *ctx) +{ + u64 probe_id = bpf_get_attach_cookie(ctx); + return collect_stack_trace(ctx, probe_id); +} diff --git a/src/xpu/flamegraph/profiler/src/bpf/profile.h b/src/xpu/flamegraph/profiler/src/bpf/profile.h new file mode 100644 index 0000000..5536ffd --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/bpf/profile.h @@ -0,0 +1,27 @@ +/* SPDX-License-Identifier: (LGPL-2.1 OR BSD-2-Clause) */ +/* Copyright (c) 2022 Meta Platforms, Inc. */ +#ifndef __PROFILE_H_ +#define __PROFILE_H_ + +#ifndef TASK_COMM_LEN +#define TASK_COMM_LEN 16 +#endif + +#ifndef MAX_STACK_DEPTH +#define MAX_STACK_DEPTH 128 +#endif + +typedef __u64 stack_trace_t[MAX_STACK_DEPTH]; + +struct stacktrace_event { + __u32 pid; + __u32 cpu_id; + __u64 timestamp; + char comm[TASK_COMM_LEN]; + __s32 kstack_sz; + __s32 ustack_sz; + stack_trace_t kstack; + stack_trace_t ustack; +}; + +#endif /* __PROFILE_H_ */ diff --git a/src/xpu/flamegraph/profiler/src/event.rs b/src/xpu/flamegraph/profiler/src/event.rs new file mode 100644 index 0000000..d07892b --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/event.rs @@ -0,0 +1,309 @@ +use std::mem; +use std::time::{SystemTime, UNIX_EPOCH}; +use blazesym::symbolize; +use nix::sys::sysinfo; + +pub const MAX_STACK_DEPTH: usize = 128; +pub const TASK_COMM_LEN: usize = 16; +const ADDR_WIDTH: usize = 16; + +// A Rust version of stacktrace_event in profile.h +#[repr(C)] +pub struct StacktraceEvent { + pub pid: u32, + pub cpu_id: u32, + pub timestamp: u64, + pub comm: [u8; TASK_COMM_LEN], + pub kstack_size: i32, + pub ustack_size: i32, + pub kstack: [u64; MAX_STACK_DEPTH], + pub ustack: [u64; MAX_STACK_DEPTH], +} + +pub enum OutputFormat { + Standard, + FoldedExtended, +} + +pub struct EventHandler { + symbolizer: symbolize::Symbolizer, + format: OutputFormat, + boot_time_ns: u64, +} + +impl EventHandler { + pub fn new(format: OutputFormat) -> Self { + // Get system uptime to calculate boot time + let boot_time_ns = Self::get_boot_time_ns(); + + Self { + symbolizer: symbolize::Symbolizer::new(), + format, + boot_time_ns, + } + } + + fn get_boot_time_ns() -> u64 { + // Get current Unix timestamp in nanoseconds + let now = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("System time before Unix epoch"); + let now_ns = now.as_nanos() as u64; + + // Get system uptime in nanoseconds + let info = sysinfo::sysinfo().expect("Failed to get sysinfo"); + let uptime_ns = (info.uptime().as_secs_f64() * 1_000_000_000.0) as u64; + + // Boot time = current time - uptime + now_ns - uptime_ns + } + + pub fn handle(&self, data: &[u8]) -> ::std::os::raw::c_int { + if data.len() != mem::size_of::() { + eprintln!( + "Invalid size {} != {}", + data.len(), + mem::size_of::() + ); + return 1; + } + + let event = unsafe { &*(data.as_ptr() as *const StacktraceEvent) }; + + if event.kstack_size <= 0 && event.ustack_size <= 0 { + return 1; + } + + match self.format { + OutputFormat::Standard => self.handle_standard(event), + OutputFormat::FoldedExtended => self.handle_folded_extended(event), + } + + 0 + } + + // Helper to extract stack slice + fn get_stack_slice<'a>(stack: &'a [u64; MAX_STACK_DEPTH], size: i32) -> &'a [u64] { + if size > 0 { + &stack[0..(size as usize / mem::size_of::())] + } else { + &[] + } + } + + // Helper to get command name + fn get_comm_str(comm: &[u8; TASK_COMM_LEN]) -> &str { + std::str::from_utf8(comm) + .unwrap_or("") + .trim_end_matches('\0') + } + + fn handle_standard(&self, event: &StacktraceEvent) { + let comm = Self::get_comm_str(&event.comm); + // Convert kernel timestamp to Unix timestamp + let unix_timestamp_ns = event.timestamp + self.boot_time_ns; + let timestamp_sec = unix_timestamp_ns / 1_000_000_000; + let timestamp_nsec = unix_timestamp_ns % 1_000_000_000; + println!("[{}.{:09}] COMM: {} (pid={}) @ CPU {}", + timestamp_sec, timestamp_nsec, comm, event.pid, event.cpu_id); + + if event.kstack_size > 0 { + println!("Kernel:"); + let kstack = Self::get_stack_slice(&event.kstack, event.kstack_size); + show_stack_trace(kstack, &self.symbolizer, 0); + } else { + println!("No Kernel Stack"); + } + + if event.ustack_size > 0 { + println!("Userspace:"); + let ustack = Self::get_stack_slice(&event.ustack, event.ustack_size); + show_stack_trace(ustack, &self.symbolizer, event.pid); + } else { + println!("No Userspace Stack"); + } + + println!(); + } + + fn handle_folded_extended(&self, event: &StacktraceEvent) { + let comm = Self::get_comm_str(&event.comm); + let tid = event.pid; // For single-threaded processes, TID = PID + + let mut stack_frames = Vec::new(); + + // Process user stack (if present) + if event.ustack_size > 0 { + let ustack = Self::get_stack_slice(&event.ustack, event.ustack_size); + let user_frames = symbolize_stack_to_vec(&self.symbolizer, ustack, event.pid); + + // Add user frames in reverse order (top to bottom) + for frame in user_frames.iter().rev() { + stack_frames.push(frame.clone()); + } + } + + // Process kernel stack (if present) + if event.kstack_size > 0 { + let kstack = Self::get_stack_slice(&event.kstack, event.kstack_size); + let kernel_frames = symbolize_stack_to_vec(&self.symbolizer, kstack, 0); + + // Add kernel frames with [k] suffix in reverse order (top to bottom) + for frame in kernel_frames.iter().rev() { + stack_frames.push(format!("{}_[k]", frame)); + } + } + + // Format: timestamp_ns comm pid tid cpu stack1;stack2;stack3 + // Convert kernel timestamp to Unix timestamp + let unix_timestamp_ns = event.timestamp + self.boot_time_ns; + println!( + "{} {} {} {} {} {}", + unix_timestamp_ns, + comm, + event.pid, + tid, + event.cpu_id, + stack_frames.join(";") + ); + } +} + + +fn print_frame( + name: &str, + addr_info: Option<(blazesym::Addr, blazesym::Addr, usize)>, + code_info: &Option, +) { + let code_info = code_info.as_ref().map(|code_info| { + let path = code_info.to_path(); + let path = path.display(); + + match (code_info.line, code_info.column) { + (Some(line), Some(col)) => format!(" {path}:{line}:{col}"), + (Some(line), None) => format!(" {path}:{line}"), + (None, _) => format!(" {path}"), + } + }); + + if let Some((input_addr, addr, offset)) = addr_info { + // If we have various address information bits we have a new symbol. + println!( + "{input_addr:#0width$x}: {name} @ {addr:#x}+{offset:#x}{code_info}", + code_info = code_info.as_deref().unwrap_or(""), + width = ADDR_WIDTH + ) + } else { + // Otherwise we are dealing with an inlined call. + println!( + "{:width$} {name}{code_info} [inlined]", + " ", + code_info = code_info + .map(|info| format!(" @{info}")) + .as_deref() + .unwrap_or(""), + width = ADDR_WIDTH + ) + } +} + +// Helper function to convert stack addresses for blazesym +fn convert_stack_addresses(stack: &[u64]) -> Vec { + if mem::size_of::() != mem::size_of::() { + stack + .iter() + .copied() + .map(|addr| addr as blazesym::Addr) + .collect::>() + } else { + // For same-sized types, still need to return owned data for consistency + stack.iter().copied().map(|addr| addr as blazesym::Addr).collect() + } +} + +// Get the stack addresses as a slice (avoiding lifetime issues) +fn get_stack_slice<'a>(stack: &'a [u64], converted: &'a [blazesym::Addr]) -> &'a [blazesym::Addr] { + if mem::size_of::() != mem::size_of::() { + converted + } else { + // SAFETY: `Addr` has the same size as `u64`, so it can be trivially and + // safely converted. + unsafe { mem::transmute::<_, &[blazesym::Addr]>(stack) } + } +} + +// Get source for symbolization based on PID (0 means kernel) +fn get_symbolize_source(pid: u32) -> symbolize::source::Source<'static> { + if pid == 0 { + symbolize::source::Source::from(symbolize::source::Kernel::default()) + } else { + symbolize::source::Source::from(symbolize::source::Process::new(pid.into())) + } +} + +// Symbolize stack and return as vector of strings for folded format +fn symbolize_stack_to_vec(symbolizer: &symbolize::Symbolizer, stack: &[u64], pid: u32) -> Vec { + let converted = convert_stack_addresses(stack); + let stack_addrs = get_stack_slice(stack, &converted); + let src = get_symbolize_source(pid); + + let syms = match symbolizer.symbolize(&src, symbolize::Input::AbsAddr(stack_addrs)) { + Ok(syms) => syms, + Err(_) => { + // Return addresses if symbolization fails + return stack_addrs.iter().map(|addr| format!("{:#x}", addr)).collect(); + } + }; + + let mut result = Vec::new(); + for (addr, sym) in stack_addrs.iter().copied().zip(syms) { + match sym { + symbolize::Symbolized::Sym(symbolize::Sym { + name, + .. + }) => { + result.push(name.to_string()); + } + symbolize::Symbolized::Unknown(..) => { + result.push(format!("{:#x}", addr)); + } + } + } + result +} + +// Pid 0 means a kernel space stack. +fn show_stack_trace(stack: &[u64], symbolizer: &symbolize::Symbolizer, pid: u32) { + let converted = convert_stack_addresses(stack); + let stack_addrs = get_stack_slice(stack, &converted); + let src = get_symbolize_source(pid); + + let syms = match symbolizer.symbolize(&src, symbolize::Input::AbsAddr(stack_addrs)) { + Ok(syms) => syms, + Err(err) => { + eprintln!(" failed to symbolize addresses: {err:#}"); + return; + } + }; + + for (input_addr, sym) in stack_addrs.iter().copied().zip(syms) { + match sym { + symbolize::Symbolized::Sym(symbolize::Sym { + name, + addr, + offset, + code_info, + inlined, + .. + }) => { + print_frame(&name, Some((input_addr, addr, offset)), &code_info); + for frame in inlined.iter() { + print_frame(&frame.name, None, &frame.code_info); + } + } + symbolize::Symbolized::Unknown(..) => { + println!("{input_addr:#0width$x}: ", width = ADDR_WIDTH) + } + } + } +} \ No newline at end of file diff --git a/src/xpu/flamegraph/profiler/src/main.rs b/src/xpu/flamegraph/profiler/src/main.rs new file mode 100644 index 0000000..e196cd1 --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/main.rs @@ -0,0 +1,176 @@ +use std::mem::MaybeUninit; +use std::time::Duration; + +use clap::ArgAction; +use clap::Parser; + +use libbpf_rs::skel::OpenSkel as _; +use libbpf_rs::skel::SkelBuilder as _; +use libbpf_rs::UprobeOpts; + +use tracing::subscriber::set_global_default as set_global_subscriber; +use tracing_subscriber::filter::LevelFilter; +use tracing_subscriber::fmt::format::FmtSpan; +use tracing_subscriber::fmt::time::SystemTime; +use tracing_subscriber::FmtSubscriber; + +mod profile { + include!(concat!(env!("OUT_DIR"), "/profile.skel.rs")); +} +mod syscall; +mod event; +mod perf; + +use profile::*; + +#[derive(Parser, Debug)] +struct Args { + /// Sampling frequency (only used in perf mode) + #[arg(short, default_value_t = 50)] + freq: u64, + /// Increase verbosity (can be supplied multiple times). + #[arg(short = 'v', long = "verbose", global = true, action = ArgAction::Count)] + verbosity: u8, + /// Use software event for triggering stack trace capture. + /// + /// This can be useful for compatibility reasons if hardware event is not available + /// (which could happen in a virtual machine, for example). + #[arg(long = "sw-event")] + sw_event: bool, + /// Filter by PID (optional) + #[arg(short = 'p', long = "pid")] + pid: Option, + /// Output in extended folded format (timestamp_ns comm pid tid cpu stack1;stack2;...) + #[arg(short = 'E', long = "fold-extend")] + fold_extend: bool, + /// Attach to kprobe (format: "symbol" e.g. "tcp_v4_connect") + /// Can be specified multiple times + #[arg(long = "kprobe")] + kprobes: Vec, + /// Attach to kretprobe (format: "symbol") + #[arg(long = "kretprobe")] + kretprobes: Vec, + /// Attach to uprobe (format: "binary:symbol" e.g. "/lib/libc.so.6:malloc") + #[arg(long = "uprobe")] + uprobes: Vec, + /// Attach to uretprobe (format: "binary:symbol") + #[arg(long = "uretprobe")] + uretprobes: Vec, +} + +fn main() -> Result<(), libbpf_rs::Error> { + let args = Args::parse(); + let level = match args.verbosity { + 0 => LevelFilter::WARN, + 1 => LevelFilter::INFO, + 2 => LevelFilter::DEBUG, + _ => LevelFilter::TRACE, + }; + + let subscriber = FmtSubscriber::builder() + .with_max_level(level) + .with_span_events(FmtSpan::FULL) + .with_timer(SystemTime) + .finish(); + let () = set_global_subscriber(subscriber).expect("failed to set tracing subscriber"); + + let skel_builder = ProfileSkelBuilder::default(); + let mut open_object = MaybeUninit::uninit(); + let open_skel = skel_builder.open(&mut open_object).unwrap(); + let skel = open_skel.load().unwrap(); + + let _perf_links; + let mut pefds = Vec::new(); + let mut _probe_links = Vec::new(); + let mut probe_id: u32 = 1; + + let has_probes = !args.kprobes.is_empty() || !args.kretprobes.is_empty() + || !args.uprobes.is_empty() || !args.uretprobes.is_empty(); + + if has_probes { + // Attach kprobes + for symbol in &args.kprobes { + let link = skel.progs.kprobe_handler.attach_kprobe(false, symbol)?; + eprintln!("Attached kprobe (id={}): {}", probe_id, symbol); + _probe_links.push(link); + probe_id += 1; + } + + // Attach kretprobes + for symbol in &args.kretprobes { + let link = skel.progs.kretprobe_handler.attach_kprobe(true, symbol)?; + eprintln!("Attached kretprobe (id={}): {}", probe_id, symbol); + _probe_links.push(link); + probe_id += 1; + } + + // Attach uprobes + for spec in &args.uprobes { + let parts: Vec<&str> = spec.split(':').collect(); + if parts.len() != 2 { + eprintln!("Error: uprobe format should be 'binary:symbol'"); + std::process::exit(1); + } + let opts = UprobeOpts { + func_name: parts[1].to_string(), + cookie: probe_id as u64, + retprobe: false, + ..Default::default() + }; + let link = skel.progs.uprobe_handler.attach_uprobe_with_opts(-1, parts[0], 0, opts)?; + eprintln!("Attached uprobe (id={}): {} in {}", probe_id, parts[1], parts[0]); + _probe_links.push(link); + probe_id += 1; + } + + // Attach uretprobes + for spec in &args.uretprobes { + let parts: Vec<&str> = spec.split(':').collect(); + if parts.len() != 2 { + eprintln!("Error: uretprobe format should be 'binary:symbol'"); + std::process::exit(1); + } + let opts = UprobeOpts { + func_name: parts[1].to_string(), + cookie: probe_id as u64, + retprobe: true, + ..Default::default() + }; + let link = skel.progs.uretprobe_handler.attach_uprobe_with_opts(-1, parts[0], 0, opts)?; + eprintln!("Attached uretprobe (id={}): {} in {}", probe_id, parts[1], parts[0]); + _probe_links.push(link); + probe_id += 1; + } + } else { + // Perf mode + let freq = if args.freq < 1 { 1 } else { args.freq }; + pefds = perf::init_perf_monitor(freq, args.sw_event, args.pid)?; + _perf_links = perf::attach_perf_event(&pefds, &skel.progs.profile); + eprintln!("Perf mode: sampling at {} Hz", freq); + } + + let output_format = if args.fold_extend { + event::OutputFormat::FoldedExtended + } else { + event::OutputFormat::Standard + }; + + let event_handler = event::EventHandler::new(output_format); + + let mut builder = libbpf_rs::RingBufferBuilder::new(); + builder + .add(&skel.maps.events, move |data| { + event_handler.handle(data) + }) + .unwrap(); + + let ringbuf = builder.build().unwrap(); + while ringbuf.poll(Duration::MAX).is_ok() {} + + // Clean up perf events if in perf mode + if !pefds.is_empty() { + perf::close_perf_events(pefds)?; + } + + Ok(()) +} diff --git a/src/xpu/flamegraph/profiler/src/perf.rs b/src/xpu/flamegraph/profiler/src/perf.rs new file mode 100644 index 0000000..2644f4c --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/perf.rs @@ -0,0 +1,63 @@ +use std::io; +use std::mem; +use nix::unistd::close; +use libbpf_rs::ErrorExt as _; + +use crate::syscall; + +pub fn init_perf_monitor(freq: u64, sw_event: bool, pid_filter: Option) -> Result, libbpf_rs::Error> { + let nprocs = libbpf_rs::num_possible_cpus().unwrap(); + let pid = pid_filter.unwrap_or(-1); + let attr = syscall::perf_event_attr { + _type: if sw_event { + syscall::PERF_TYPE_SOFTWARE + } else { + syscall::PERF_TYPE_HARDWARE + }, + size: mem::size_of::() as u32, + config: if sw_event { + syscall::PERF_COUNT_SW_CPU_CLOCK + } else { + syscall::PERF_COUNT_HW_CPU_CYCLES + }, + sample: syscall::sample_un { sample_freq: freq }, + flags: 1 << 10, // freq = 1 + ..Default::default() + }; + (0..nprocs) + .map(|cpu| { + let fd = syscall::perf_event_open(&attr, pid, cpu as i32, -1, 0) as i32; + if fd == -1 { + let mut error_context = "Failed to open perf event."; + let os_error = io::Error::last_os_error(); + if !sw_event && os_error.kind() == io::ErrorKind::NotFound { + error_context = "Failed to open perf event.\n\ + Try running the profile example with the `--sw-event` option."; + } + Err(libbpf_rs::Error::from(os_error)).context(error_context) + } else { + Ok(fd) + } + }) + .collect() +} + +pub fn attach_perf_event( + pefds: &[i32], + prog: &libbpf_rs::ProgramMut, +) -> Vec> { + pefds + .iter() + .map(|pefd| prog.attach_perf_event(*pefd)) + .collect() +} + +pub fn close_perf_events(pefds: Vec) -> Result<(), libbpf_rs::Error> { + for pefd in pefds { + close(pefd) + .map_err(io::Error::from) + .map_err(libbpf_rs::Error::from) + .context("failed to close perf event")?; + } + Ok(()) +} \ No newline at end of file diff --git a/src/xpu/flamegraph/profiler/src/syscall.rs b/src/xpu/flamegraph/profiler/src/syscall.rs new file mode 100644 index 0000000..c12fc73 --- /dev/null +++ b/src/xpu/flamegraph/profiler/src/syscall.rs @@ -0,0 +1,90 @@ +use std::mem; + +extern crate libc; + +#[repr(C)] +pub union sample_un { + pub sample_period: u64, + pub sample_freq: u64, +} + +#[repr(C)] +pub union wakeup_un { + pub wakeup_events: u32, + pub wakeup_atermark: u32, +} + +#[repr(C)] +pub union bp_1_un { + pub bp_addr: u64, + pub kprobe_func: u64, + pub uprobe_path: u64, + pub config1: u64, +} + +#[repr(C)] +pub union bp_2_un { + pub bp_len: u64, + pub kprobe_addr: u64, + pub probe_offset: u64, + pub config2: u64, +} + +#[repr(C)] +pub struct perf_event_attr { + pub _type: u32, + pub size: u32, + pub config: u64, + pub sample: sample_un, + pub sample_type: u64, + pub read_format: u64, + pub flags: u64, + pub wakeup: wakeup_un, + pub bp_type: u32, + pub bp_1: bp_1_un, + pub bp_2: bp_2_un, + pub branch_sample_type: u64, + pub sample_regs_user: u64, + pub sample_stack_user: u32, + pub clockid: i32, + pub sample_regs_intr: u64, + pub aux_watermark: u32, + pub sample_max_stack: u16, + pub __reserved_2: u16, + pub aux_sample_size: u32, + pub __reserved_3: u32, +} + +impl Default for perf_event_attr { + fn default() -> Self { + unsafe { mem::zeroed() } + } +} + +pub const PERF_TYPE_HARDWARE: u32 = 0; +pub const PERF_TYPE_SOFTWARE: u32 = 1; +pub const PERF_COUNT_HW_CPU_CYCLES: u64 = 0; +pub const PERF_COUNT_SW_CPU_CLOCK: u64 = 0; + +extern "C" { + fn syscall(number: libc::c_long, ...) -> libc::c_long; +} + +pub fn perf_event_open( + hw_event: &perf_event_attr, + pid: libc::pid_t, + cpu: libc::c_int, + group_fd: libc::c_int, + flags: libc::c_ulong, +) -> libc::c_long { + unsafe { + syscall( + libc::SYS_perf_event_open, + hw_event as *const perf_event_attr, + pid, + cpu, + group_fd, + flags, + ) + } +} diff --git a/src/xpu/flamegraph/qwen3.cu/.gitattributes b/src/xpu/flamegraph/qwen3.cu/.gitattributes new file mode 100644 index 0000000..e18ec67 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/.gitattributes @@ -0,0 +1 @@ +convert_hf_to_gguf_ordered.py linguist-vendored diff --git a/src/xpu/flamegraph/qwen3.cu/.gitignore b/src/xpu/flamegraph/qwen3.cu/.gitignore new file mode 100644 index 0000000..50c3520 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/.gitignore @@ -0,0 +1,10 @@ +*.i +*.ii +*.gpu +*.ptx +*.cubin +*.fatbin +runcu +runcublas +*.gguf +Qwen3-0.6B-GGUF-FP32/ diff --git a/src/xpu/flamegraph/qwen3.cu/LICENSE b/src/xpu/flamegraph/qwen3.cu/LICENSE new file mode 100644 index 0000000..6c13926 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 William Song + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/src/xpu/flamegraph/qwen3.cu/Makefile b/src/xpu/flamegraph/qwen3.cu/Makefile new file mode 100644 index 0000000..1010cf8 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/Makefile @@ -0,0 +1,68 @@ +# choose your compiler, e.g. gcc/clang +# example override to clang: make run CC=clang +CC = gcc +# For CUDA compile +CUDA_INSTALL_PATH ?= /usr/local/cuda-12.9 +NVCC := "$(CUDA_INSTALL_PATH)/bin/nvcc" +INCLUDES := -I"$(CUDA_INSTALL_PATH)/include" +LIB_PATH ?= $(CUDA_INSTALL_PATH)/lib64 + +# compile the Cuda version (with dynamic libcudart for eBPF uprobe profiling) +.PHONY: runcu +runcu: runcu.cu + $(NVCC) $(INCLUDES) -O3 -Wno-deprecated-gpu-targets --no-device-link -o runcu runcu.cu -L $(LIB_PATH) -lcudart -lm +# compile cublas included +.PHONY: runcublas +runcublas: runcu.cu + $(NVCC) $(INCLUDES) -O3 -Wno-deprecated-gpu-targets --no-device-link -DUSE_CUBLAS -o runcublas runcu.cu -L $(LIB_PATH) -lcudart -lm -lcublas + +# download the model +.PHONY: download-model +download-model: + @if [ -f Qwen3-0.6B-FP32.gguf ] && [ $$(stat -c%s Qwen3-0.6B-FP32.gguf) -gt 1000000 ]; then \ + echo "Model already exists (size: $$(du -h Qwen3-0.6B-FP32.gguf | cut -f1))"; \ + else \ + echo "Downloading Qwen3-0.6B model (3GB - this will take a while)..."; \ + wget -c https://huggingface.co/huggit0000/Qwen3-0.6B-GGUF-FP32/resolve/main/Qwen3-0.6B-FP32.gguf -O Qwen3-0.6B-FP32.gguf || \ + curl -L -C - https://huggingface.co/huggit0000/Qwen3-0.6B-GGUF-FP32/resolve/main/Qwen3-0.6B-FP32.gguf -o Qwen3-0.6B-FP32.gguf; \ + echo "Model downloaded successfully (size: $$(du -h Qwen3-0.6B-FP32.gguf | cut -f1))"; \ + fi + +# ========================= +# The below is not used hree. + + +# the most basic way of building that is most likely to work on most systems +.PHONY: run +run: run.c + $(CC) -O3 -o run run.c -lm + +# useful for a debug build, can then e.g. analyze with valgrind, example: +# $ valgrind --leak-check=full ./run out/model.bin -n 3 +rundebug: run.c + $(CC) -g -o run run.c -lm + +# https://gcc.gnu.org/onlinedocs/gcc/Optimize-Options.html +# https://simonbyrne.github.io/notes/fastmath/ +# -Ofast enables all -O3 optimizations. +# Disregards strict standards compliance. +# It also enables optimizations that are not valid for all standard-compliant programs. +# It turns on -ffast-math, -fallow-store-data-races and the Fortran-specific +# -fstack-arrays, unless -fmax-stack-var-size is specified, and -fno-protect-parens. +# It turns off -fsemantic-interposition. +# In our specific application this is *probably* okay to use +#.PHONY: run +#runfast: run.c +# $(CC) -O3 -o run -fopenmp -march=native run.c -lm + +# additionally compiles with OpenMP, allowing multithreaded runs +# make sure to also enable multiple threads when running, e.g.: +# OMP_NUM_THREADS=4 ./run out/model.bin +.PHONY: runomp +runomp: run.c + $(CC) -O3 -fopenmp -march=native run.c -lm -o run + + +.PHONY: clean +clean: + rm -f run runcu runcublas diff --git a/src/xpu/flamegraph/qwen3.cu/README b/src/xpu/flamegraph/qwen3.cu/README new file mode 100644 index 0000000..e949c71 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/README @@ -0,0 +1,166 @@ +# qwen3.cu + +`qwen3.cu` is a **single-file, pure CUDA C implementation** for running inference on the Qwen3 model with no external libraries, no dependencies. It’s a follow-up to my earlier weekend project, [qwen3.c](https://github.com/...), inspired by Andrej Karpathy’s [`llama2.c`](https://github.com/karpathy/llama2.c). Everything’s packed into one file from tokenization all the way to CUDA kernels, staying true to the spirit of minimalism. + +This implementation runs the Qwen3 0.6B model, a small but capable model. I'm using **full-precision GGUF** here, chosen for its clarity and to help others learn its ways. Also, It’s fully self-contained, so there’s no need for any format conversion out of the box. Most GGUF models are quantized to 8-bit or lower, but for this project, you’ll want to use the FP32 version which you can download as below. Or, if you make it work from the BF16 weights, you can convert them using the included `convert_hf_to_gguf_ordered.py` script; I've made sure the layers are ordered numerically so everything aligns correctly. + +Even though GGUF files already include a binary tokenizer, this project reads vocab and merges from plain `.txt` files. It keeps things more transparent and easier to follow. Tokenization and detokenization overhead is negligible compared to the forward pass, so it doesn’t really impact TTS. + +It also supports multi-turn conversation out of the box, and native support for Qwen3’s reasoning mode. For reference, there’s also a cuBLAS version included. It’s roughly 2x faster for now, but I’ll probably try to narrow that gap in the future. I’ll add more explanation on the code later. + +### UPDATE +[Oct-27-25] Added single prompt mode (-q flag) for non-interactive usage +[Oct-27-25] Updated Makefile with --no-device-link flag and download-model target +[Aug-08-25] Remove the nonsense loop. TPS increased from ~35 to ~39. Set base for benchmarking. +[What's next] Improve kernels + +## Quick Start + +```sh +# Clone this repo +git clone https://github.com/gigit0000/qwen3.cu.git +cd qwen3.cu + +# Download FP32 model (3GB) - uses wget/curl, no Git LFS required +make download-model + +# Compile and run (interactive mode) +make runcu +./runcu Qwen3-0.6B-FP32.gguf + +# Or use single prompt mode (runs once and exits) +./runcu Qwen3-0.6B-FP32.gguf -q "What is CUDA?" +``` + +## Faster Inference +Use cuBLAS (roughly 2x faster): +```sh +# Compile and run +make runcublas +./runcublas Qwen3-0.6B-FP32.gguf + +# Single prompt with cuBLAS +./runcublas Qwen3-0.6B-FP32.gguf -q "Explain quantum computing" -r 1 +``` + +## Makefile Improvements + +The Makefile now includes: +- **`make download-model`**: Automatically downloads the 3GB FP32 model using wget/curl (no Git LFS required) +- **Fixed compilation flags**: Added `-Wno-deprecated-gpu-targets --no-device-link` to fix build issues on newer CUDA versions +- **Clean target**: `make clean` removes built binaries + +## Description + +You can enable reasoning (-k 1) or multi-turn (-m 1): +``` +./runcu Qwen3-0.6B-FP32.gguf -k 1 -m 1 +``` + +**New: Single Prompt Mode (-q)** +Run a single query and exit (useful for scripting and automation): +```sh +./runcu Qwen3-0.6B-FP32.gguf -q "What is machine learning?" +./runcu Qwen3-0.6B-FP32.gguf -q "Explain eBPF in one sentence" + +# Combine with other flags +./runcu Qwen3-0.6B-FP32.gguf -q "Why is the sky blue?" -r 1 # with TPS +./runcu Qwen3-0.6B-FP32.gguf -q "2+2=?" -t 0.3 # lower temperature +``` + +If you want to extract text files (vocab.txt, merges.txt and header.txt) on your own, you can use the scripts: +```sh +# tokenizer - vocab.txt and merges.txt +python extract_v_m.py Qwen3-0.6B-FP32.gguf + +``` + +### Inference Examples + +Multi-turn Conversation with the option m +``` +# ./runcu Qwen3-0.6B-FP32.gguf -m 1 -k 0 +Multi-turn = on, thinKing = off, Temperature = 0.60, top-P = 0.95 +Press Enter to exit the chat +Enter system prompt (or Enter to skip): Tell me in one sentence +Q: Where is the best spot in Paris? +A: The best spot in Paris is the Eiffel Tower. +Q: What about the second-best spot? +A: The second-best spot in Paris is the Louvre Museum. +``` + +Reasoning with the option k +``` +# ./runcu Qwen3-0.6B-FP32.gguf -k 1 +Multi-turn = off, thinKing = on, Temperature = 0.60, top-P = 0.95 +Press Enter to exit the chat +Enter system prompt (or Enter to skip): +Q: Why do stars shine? Give me a quick answer! +A: +Okay, the user is asking why stars shine. Let me start by recalling what I know about stars. Stars are luminous objects that emit light. So, the main reason they shine is because they produce light through nuclear fusion. + +Wait, but I should make sure. Stars form from clouds of gas and dust in space. When these clouds cool, they start fusing hydrogen into helium, which releases energy. This energy is what we see as light. So the process is nuclear fusion of hydrogen into helium, which gives off energy. + +I should also mention that the energy from stars is what we perceive as light. Maybe add that this light travels through space and we see it on Earth. But the question is why they shine, so the answer should focus on the energy production. + +I need to keep it simple and concise. The user probably wants a quick answer, so no need for too much detail. Let me check if there's any other reason, but I think that's the main one. Alright, I think that's it. + + +Stars shine because they produce light through nuclear fusion of hydrogen into helium in their cores. This energy is then released as visible light, giving them their luminous glow. +``` +You can enable and monitor TPS with the r option: +``` +./runcu Qwen3-0.6B-FP32.gguf -r 1 +Multi-turn = off, thinKing = off, tps(R) = on, Temperature = 0.60, top-P = 0.95 +Press Enter to exit the chat +Enter system prompt (or Enter to skip): You name is Tom. +Q: What is your name? +A: My name is Tom. +tok/s: 34.482759 +``` + +## Command-Line Options + +``` +Usage: ./runcu [options] +Example: ./runcu Qwen3-0.6B-FP32.gguf + ./runcu Qwen3-0.6B-FP32.gguf -q "What is CUDA?" + +Options: + -t temperature in [0,inf], default 0.6 + -p p value in top-p (nucleus) sampling in [0,1] default 0.95 + -s random seed, default time(NULL) + -m multi-turn: 0 = off (default), 1 = on + -k reasoning: 0 = off (default), 1 = on + -r TPS: 0 = off (default), 1 = on + -q single prompt mode (run once and exit) +``` + +**Usage Tips:** +- Use `-q` for automation, scripting, or quick queries +- Combine `-q` with `-r 1` to measure inference speed +- Use `-k 1` to enable Qwen3's reasoning mode (shows thinking process) +- Use `-m 1` for multi-turn conversations (maintains context) +- Lower `-t` (temperature) for more deterministic outputs +- Use `runcublas` instead of `runcu` for 2x faster inference + +## (Maybe) TODO +- [ ] Kernel optimization +- [ ] CUTLASS version +- [ ] KV cache for multi-turn conversations + +## Acknoledgement +- Inspired and baselined from Andrej Kapathy's [llama2.c](https://github.com/karpathy/llama2.c) +- Most kernels and CUDA ports were originally adopted from @rogerallen's great repo [llama2.cu](https://github.com/rogerallen/) +- Based on my qwen3.c [repo](https://github.com/gigit0000/qwen3.c/) +- GGUF [llama.cpp](https://github.com/ggml-org/llama.cpp) +- FGPF + +## License +MIT + + + + + + diff --git a/src/xpu/flamegraph/qwen3.cu/convert_hf_to_gguf_ordered.py b/src/xpu/flamegraph/qwen3.cu/convert_hf_to_gguf_ordered.py new file mode 100644 index 0000000..fd9af76 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/convert_hf_to_gguf_ordered.py @@ -0,0 +1,6843 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +from __future__ import annotations + +import ast +import logging +import argparse +import contextlib +import json +import os +import re +import sys +from enum import IntEnum +from pathlib import Path +from hashlib import sha256 +from typing import TYPE_CHECKING, Any, Callable, ContextManager, Iterable, Iterator, Literal, Sequence, TypeVar, cast +from itertools import chain +from transformers import AutoConfig + +import math +import numpy as np +import torch + +if TYPE_CHECKING: + from torch import Tensor + +if 'NO_LOCAL_GGUF' not in os.environ: + sys.path.insert(1, str(Path(__file__).parent / 'gguf-py')) +import gguf + +logger = logging.getLogger("hf-to-gguf") + + +###### MODEL DEFINITIONS ###### + +class SentencePieceTokenTypes(IntEnum): + NORMAL = 1 + UNKNOWN = 2 + CONTROL = 3 + USER_DEFINED = 4 + UNUSED = 5 + BYTE = 6 + + +class ModelType(IntEnum): + TEXT = 1 + MMPROJ = 2 + + +AnyModel = TypeVar("AnyModel", bound="type[ModelBase]") + + +class ModelBase: + _model_classes: dict[ModelType, dict[str, type[ModelBase]]] = { + ModelType.TEXT: {}, + ModelType.MMPROJ: {}, + } + + dir_model: Path + ftype: gguf.LlamaFileType + fname_out: Path + is_big_endian: bool + endianess: gguf.GGUFEndian + use_temp_file: bool + lazy: bool + part_names: list[str] + is_safetensors: bool + hparams: dict[str, Any] + tensor_names: set[str] | None + gguf_writer: gguf.GGUFWriter + model_name: str | None + metadata_override: Path | None + dir_model_card: Path + remote_hf_model_id: str | None + + # subclasses should define this! + model_arch: gguf.MODEL_ARCH + + # subclasses should initialize this! + block_count: int + tensor_map: gguf.TensorNameMap + + def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False, + use_temp_file: bool = False, eager: bool = False, + metadata_override: Path | None = None, model_name: str | None = None, + split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, + small_first_shard: bool = False, hparams: dict[str, Any] | None = None, remote_hf_model_id: str | None = None): + if type(self) is ModelBase or \ + type(self) is TextModel or \ + type(self) is MmprojModel: + raise TypeError(f"{type(self).__name__!r} should not be directly instantiated") + + self.dir_model = dir_model + self.ftype = ftype + self.fname_out = fname_out + self.is_big_endian = is_big_endian + self.endianess = gguf.GGUFEndian.BIG if is_big_endian else gguf.GGUFEndian.LITTLE + self.use_temp_file = use_temp_file + self.lazy = not eager or (remote_hf_model_id is not None) + self.remote_hf_model_id = remote_hf_model_id + if remote_hf_model_id is not None: + self.is_safetensors = True + + def get_remote_tensors() -> Iterator[tuple[str, Tensor]]: + logger.info(f"Using remote model with HuggingFace id: {remote_hf_model_id}") + remote_tensors = gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id) + self.tensor_names = set(name for name in remote_tensors.keys()) + for name, remote_tensor in gguf.utility.SafetensorRemote.get_list_tensors_hf_model(remote_hf_model_id).items(): + yield (name, LazyTorchTensor.from_remote_tensor(remote_tensor)) + + self.get_tensors = get_remote_tensors + else: + self.part_names = ModelBase.get_model_part_names(self.dir_model, "model", ".safetensors") + self.is_safetensors = len(self.part_names) > 0 + if not self.is_safetensors: + self.part_names = ModelBase.get_model_part_names(self.dir_model, "pytorch_model", ".bin") + self.hparams = ModelBase.load_hparams(self.dir_model) if hparams is None else hparams + self.tensor_names = None + self.metadata_override = metadata_override + self.model_name = model_name + self.dir_model_card = dir_model # overridden in convert_lora_to_gguf.py + + # Apply heuristics to figure out typical tensor encoding based on first layer tensor encoding type + if self.ftype == gguf.LlamaFileType.GUESSED: + # NOTE: can't use field "torch_dtype" in config.json, because some finetunes lie. + _, first_tensor = next(self.get_tensors()) + if first_tensor.dtype == torch.float16: + logger.info(f"choosing --outtype f16 from first tensor type ({first_tensor.dtype})") + self.ftype = gguf.LlamaFileType.MOSTLY_F16 + else: + logger.info(f"choosing --outtype bf16 from first tensor type ({first_tensor.dtype})") + self.ftype = gguf.LlamaFileType.MOSTLY_BF16 + + # Configure GGUF Writer + self.gguf_writer = gguf.GGUFWriter(path=None, arch=gguf.MODEL_ARCH_NAMES[self.model_arch], endianess=self.endianess, use_temp_file=self.use_temp_file, + split_max_tensors=split_max_tensors, split_max_size=split_max_size, dry_run=dry_run, small_first_shard=small_first_shard) + + @classmethod + def add_prefix_to_filename(cls, path: Path, prefix: str) -> Path: + stem, suffix = path.stem, path.suffix + new_name = f"{prefix}{stem}{suffix}" + return path.with_name(new_name) + + def find_hparam(self, keys: Iterable[str], optional: bool = False) -> Any: + key = next((k for k in keys if k in self.hparams), None) + if key is not None: + return self.hparams[key] + if optional: + return None + raise KeyError(f"could not find any of: {keys}") + + def get_tensors(self) -> Iterator[tuple[str, Tensor]]: + tensor_names_from_parts: set[str] = set() + + index_name = "model.safetensors" if self.is_safetensors else "pytorch_model.bin" + index_name += ".index.json" + index_file = self.dir_model / index_name + + if index_file.is_file(): + self.tensor_names = set() + logger.info(f"gguf: loading model weight map from '{index_name}'") + with open(index_file, "r", encoding="utf-8") as f: + index: dict[str, Any] = json.load(f) + weight_map = index.get("weight_map") + if weight_map is None or not isinstance(weight_map, dict): + raise ValueError(f"Can't load 'weight_map' from {index_name!r}") + self.tensor_names.update(weight_map.keys()) + else: + self.tensor_names = tensor_names_from_parts + weight_map = {} + + for part_name in self.part_names: + logger.info(f"gguf: loading model part '{part_name}'") + ctx: ContextManager[Any] + if self.is_safetensors: + from safetensors import safe_open + ctx = cast(ContextManager[Any], safe_open(self.dir_model / part_name, framework="pt", device="cpu")) + else: + ctx = contextlib.nullcontext(torch.load(str(self.dir_model / part_name), map_location="cpu", mmap=True, weights_only=True)) + + with ctx as model_part: + tensor_names_from_parts.update(model_part.keys()) + + for name in model_part.keys(): + if self.is_safetensors: + if self.lazy: + data = model_part.get_slice(name) + data = LazyTorchTensor.from_safetensors_slice(data) + else: + data = model_part.get_tensor(name) + else: + data = model_part[name] + if self.lazy: + data = LazyTorchTensor.from_eager(data) + yield name, data + + # verify tensor name presence and identify potentially missing files + if len(tensor_names_from_parts.symmetric_difference(self.tensor_names)) > 0: + missing = sorted(self.tensor_names.difference(tensor_names_from_parts)) + extra = sorted(tensor_names_from_parts.difference(self.tensor_names)) + missing_files = sorted(set(weight_map[n] for n in missing if n in weight_map)) + if len(extra) == 0 and len(missing_files) > 0: + raise ValueError(f"Missing or incomplete model files: {missing_files}\n" + f"Missing tensors: {missing}") + else: + raise ValueError("Mismatch between weight map and model parts for tensor names:\n" + f"Missing tensors: {missing}\n" + f"Extra tensors: {extra}") + + def format_tensor_name(self, key: gguf.MODEL_TENSOR, bid: int | None = None, suffix: str = ".weight") -> str: + if key not in gguf.MODEL_TENSORS[self.model_arch]: + raise ValueError(f"Missing {key!r} for MODEL_TENSORS of {self.model_arch!r}") + name: str = gguf.TENSOR_NAMES[key] + if "{bid}" in name: + assert bid is not None + name = name.format(bid=bid) + return name + suffix + + def match_model_tensor_name(self, name: str, key: gguf.MODEL_TENSOR, bid: int | None, suffix: str = ".weight") -> bool: + if key not in gguf.MODEL_TENSORS[self.model_arch]: + return False + key_name: str = gguf.TENSOR_NAMES[key] + if "{bid}" in key_name: + if bid is None: + return False + key_name = key_name.format(bid=bid) + else: + if bid is not None: + return False + return name == (key_name + suffix) + + def map_tensor_name(self, name: str, try_suffixes: Sequence[str] = (".weight", ".bias")) -> str: + new_name = self.tensor_map.get_name(key=name, try_suffixes=try_suffixes) + if new_name is None: + raise ValueError(f"Can not map tensor {name!r}") + return new_name + + def set_gguf_parameters(self): + raise NotImplementedError("set_gguf_parameters() must be implemented in subclasses") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + return [(self.map_tensor_name(name), data_torch)] + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + del name, new_name, bid, n_dims # unused + + return False + + # some models need extra generated tensors (like rope_freqs) + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + return () + + def prepare_tensors(self): + max_name_len = max(len(s) for _, s in self.tensor_map.mapping.values()) + len(".weight,") + + all_tensors = [] + + for name, data_torch in chain(self.generate_extra_tensors(), self.get_tensors()): + # we don't need these + if name.endswith((".attention.masked_bias", ".attention.bias", ".rotary_emb.inv_freq")): + continue + + old_dtype = data_torch.dtype + + # convert any unsupported data types to float32 + if data_torch.dtype not in (torch.float16, torch.float32): + data_torch = data_torch.to(torch.float32) + + # use the first number-like part of the tensor name as the block id + bid = None + for part in name.split("."): + if part.isdecimal(): + bid = int(part) + break + + + + for new_name, data_torch in (self.modify_tensors(data_torch, name, bid)): + all_tensors.append((new_name, data_torch, name, bid, old_dtype)) + + # Sort tensors by block ID (bid), putting None values first + def sort_key(tensor_info): + new_name, data_torch, name, bid, old_dtype = tensor_info + if bid is None: + return (-1, new_name) # Put non-block tensors first + else: + return (bid, new_name) # Sort blocks numerically, then by name + + all_tensors.sort(key=sort_key) + + + for new_name, data_torch, name, bid, old_dtype in all_tensors: + # TODO: why do we squeeze here? + # data = data_torch.squeeze().numpy() + data = data_torch.numpy() + + # if data ends up empty, it means data_torch was a scalar tensor -> restore + if len(data.shape) == 0: + data = data_torch.numpy() + + n_dims = len(data.shape) + data_qtype: gguf.GGMLQuantizationType | bool = self.tensor_force_quant(name, new_name, bid, n_dims) + + # Most of the codebase that takes in 1D tensors or norms only handles F32 tensors + if n_dims <= 1 or new_name.endswith("_norm.weight"): + data_qtype = gguf.GGMLQuantizationType.F32 + + # Conditions should closely match those in llama_model_quantize_internal in llama.cpp + # Some tensor types are always in float32 + if data_qtype is False and ( + any( + self.match_model_tensor_name(new_name, key, bid) + for key in ( + gguf.MODEL_TENSOR.FFN_GATE_INP, + gguf.MODEL_TENSOR.POS_EMBD, + gguf.MODEL_TENSOR.TOKEN_TYPES, + gguf.MODEL_TENSOR.SSM_CONV1D, + gguf.MODEL_TENSOR.TIME_MIX_FIRST, + gguf.MODEL_TENSOR.TIME_MIX_W1, + gguf.MODEL_TENSOR.TIME_MIX_W2, + gguf.MODEL_TENSOR.TIME_MIX_DECAY_W1, + gguf.MODEL_TENSOR.TIME_MIX_DECAY_W2, + gguf.MODEL_TENSOR.TIME_MIX_LERP_FUSED, + gguf.MODEL_TENSOR.POSNET_NORM1, + gguf.MODEL_TENSOR.POSNET_NORM2, + gguf.MODEL_TENSOR.V_ENC_EMBD_POS, + gguf.MODEL_TENSOR.A_ENC_EMBD_POS, + gguf.MODEL_TENSOR.ALTUP_CORRECT_COEF, + gguf.MODEL_TENSOR.ALTUP_PREDICT_COEF, + ) + ) + or not new_name.endswith(".weight") + ): + data_qtype = gguf.GGMLQuantizationType.F32 + + if data_qtype is False and any( + self.match_model_tensor_name(new_name, key, bid) + for key in ( + gguf.MODEL_TENSOR.TOKEN_EMBD, + gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + gguf.MODEL_TENSOR.OUTPUT, + gguf.MODEL_TENSOR.ALTUP_ROUTER, + gguf.MODEL_TENSOR.LAUREL_L, + gguf.MODEL_TENSOR.LAUREL_R, + ) + ): + if self.ftype in ( + gguf.LlamaFileType.MOSTLY_TQ1_0, + gguf.LlamaFileType.MOSTLY_TQ2_0, + ): + # TODO: use Q4_K and Q6_K + data_qtype = gguf.GGMLQuantizationType.F16 + + # No override (data_qtype is False), or wants to be quantized (data_qtype is True) + if isinstance(data_qtype, bool): + if self.ftype == gguf.LlamaFileType.ALL_F32: + data_qtype = gguf.GGMLQuantizationType.F32 + elif self.ftype == gguf.LlamaFileType.MOSTLY_F16: + data_qtype = gguf.GGMLQuantizationType.F16 + elif self.ftype == gguf.LlamaFileType.MOSTLY_BF16: + data_qtype = gguf.GGMLQuantizationType.BF16 + elif self.ftype == gguf.LlamaFileType.MOSTLY_Q8_0: + data_qtype = gguf.GGMLQuantizationType.Q8_0 + elif self.ftype == gguf.LlamaFileType.MOSTLY_TQ1_0: + data_qtype = gguf.GGMLQuantizationType.TQ1_0 + elif self.ftype == gguf.LlamaFileType.MOSTLY_TQ2_0: + data_qtype = gguf.GGMLQuantizationType.TQ2_0 + else: + raise ValueError(f"Unknown file type: {self.ftype.name}") + + try: + data = gguf.quants.quantize(data, data_qtype) + except gguf.QuantError as e: + logger.warning("%s, %s", e, "falling back to F16") + data_qtype = gguf.GGMLQuantizationType.F16 + data = gguf.quants.quantize(data, data_qtype) + + shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape + + # reverse shape to make it similar to the internal ggml dimension order + shape_str = f"{{{', '.join(str(n) for n in reversed(shape))}}}" + + # n_dims is implicit in the shape + logger.info(f"{f'%-{max_name_len}s' % f'{new_name},'} {old_dtype} --> {data_qtype.name}, shape = {shape_str}") + + self.gguf_writer.add_tensor(new_name, data, raw_dtype=data_qtype) + + def set_type(self): + self.gguf_writer.add_type(gguf.GGUFType.MODEL) + + def prepare_metadata(self, vocab_only: bool): + + total_params, shared_params, expert_params, expert_count = self.gguf_writer.get_total_parameter_count() + + self.metadata = gguf.Metadata.load(self.metadata_override, self.dir_model_card, self.model_name, total_params) + + # If we are using HF model id, set the metadata name to the model id + if self.remote_hf_model_id: + self.metadata.name = self.remote_hf_model_id + + # Fallback to model directory name if metadata name is still missing + if self.metadata.name is None: + self.metadata.name = self.dir_model.name + + # Generate parameter weight class (useful for leader boards) if not yet determined + if self.metadata.size_label is None and total_params > 0: + self.metadata.size_label = gguf.size_label(total_params, shared_params, expert_params, expert_count) + + self.set_type() + + logger.info("Set meta model") + self.metadata.set_gguf_meta_model(self.gguf_writer) + + logger.info("Set model parameters") + self.set_gguf_parameters() + + logger.info("Set model quantization version") + self.gguf_writer.add_quantization_version(gguf.GGML_QUANT_VERSION) + + def write_vocab(self): + raise NotImplementedError("write_vocab() must be implemented in subclasses") + + def write(self): + self.prepare_tensors() + self.prepare_metadata(vocab_only=False) + self.gguf_writer.write_header_to_file(path=self.fname_out) + self.gguf_writer.write_kv_data_to_file() + self.gguf_writer.write_tensors_to_file(progress=True) + self.gguf_writer.close() + + @staticmethod + def get_model_part_names(dir_model: Path, prefix: str, suffix: str) -> list[str]: + part_names: list[str] = [] + for filename in os.listdir(dir_model): + if filename.startswith(prefix) and filename.endswith(suffix): + part_names.append(filename) + + part_names.sort() + + return part_names + + @staticmethod + def load_hparams(dir_model: Path): + try: + # for security reason, we don't allow loading remote code by default + # if a model need remote code, we will fallback to config.json + config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict() + except Exception as e: + logger.warning(f"Failed to load model config from {dir_model}: {e}") + logger.warning("Trying to load config.json instead") + with open(dir_model / "config.json", "r", encoding="utf-8") as f: + config = json.load(f) + if "llm_config" in config: + # rename for InternVL + config["text_config"] = config["llm_config"] + if "thinker_config" in config: + # rename for Qwen2.5-Omni + config["text_config"] = config["thinker_config"]["text_config"] + return config + + @classmethod + def register(cls, *names: str) -> Callable[[AnyModel], AnyModel]: + assert names + + def func(modelcls: AnyModel) -> AnyModel: + model_type = ModelType.MMPROJ if modelcls.model_arch == gguf.MODEL_ARCH.MMPROJ else ModelType.TEXT + for name in names: + cls._model_classes[model_type][name] = modelcls + return modelcls + return func + + @classmethod + def print_registered_models(cls): + for model_type, model_classes in cls._model_classes.items(): + logger.error(f"{model_type.name} models:") + for name in sorted(model_classes.keys()): + logger.error(f" - {name}") + + @classmethod + def from_model_architecture(cls, arch: str, model_type = ModelType.TEXT) -> type[ModelBase]: + try: + return cls._model_classes[model_type][arch] + except KeyError: + raise NotImplementedError(f'Architecture {arch!r} not supported!') from None + + +class TextModel(ModelBase): + model_type = ModelType.TEXT + hf_arch: str + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.hf_arch = get_model_architecture(self.hparams, self.model_type) + + if "text_config" in self.hparams: + # move the text_config to the root level + self.hparams = {**self.hparams, **self.hparams["text_config"]} + + self.block_count = self.find_hparam(["n_layers", "num_hidden_layers", "n_layer", "num_layers"]) + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + @classmethod + def __init_subclass__(cls): + # can't use an abstract property, because overriding it without type errors + # would require using decorated functions instead of simply defining the property + if "model_arch" not in cls.__dict__: + raise TypeError(f"Missing property 'model_arch' for {cls.__name__!r}") + + def set_vocab(self): + self._set_vocab_gpt2() + + def prepare_metadata(self, vocab_only: bool): + super().prepare_metadata(vocab_only=vocab_only) + + total_params = self.gguf_writer.get_total_parameter_count()[0] + # Extract the encoding scheme from the file type name. e.g. 'gguf.LlamaFileType.MOSTLY_Q8_0' --> 'Q8_0' + output_type: str = self.ftype.name.partition("_")[2] + + # Filename Output + if self.fname_out.is_dir(): + # Generate default filename based on model specification and available metadata + if not vocab_only: + fname_default: str = gguf.naming_convention(self.metadata.name, self.metadata.basename, self.metadata.finetune, self.metadata.version, self.metadata.size_label, output_type, model_type="LoRA" if total_params < 0 else None) + else: + fname_default: str = gguf.naming_convention(self.metadata.name, self.metadata.basename, self.metadata.finetune, self.metadata.version, size_label=None, output_type=None, model_type="vocab") + + # Use the default filename + self.fname_out = self.fname_out / f"{fname_default}.gguf" + else: + # Output path is a custom defined templated filename + # Note: `not is_dir()` is used because `.is_file()` will not detect + # file template strings as it doesn't actually exist as a file + + # Process templated file name with the output ftype, useful with the "auto" ftype + self.fname_out = self.fname_out.parent / gguf.fill_templated_filename(self.fname_out.name, output_type) + + logger.info("Set model tokenizer") + self.set_vocab() + + def set_gguf_parameters(self): + self.gguf_writer.add_block_count(self.block_count) + + if (n_ctx := self.find_hparam(["max_position_embeddings", "n_ctx", "n_positions", "max_length"], optional=True)) is not None: + self.gguf_writer.add_context_length(n_ctx) + logger.info(f"gguf: context length = {n_ctx}") + + if (n_embd := self.find_hparam(["hidden_size", "n_embd", "dim"], optional=True)) is not None: + self.gguf_writer.add_embedding_length(n_embd) + logger.info(f"gguf: embedding length = {n_embd}") + + if (n_ff := self.find_hparam(["intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None: + self.gguf_writer.add_feed_forward_length(n_ff) + logger.info(f"gguf: feed forward length = {n_ff}") + + if (n_head := self.find_hparam(["num_attention_heads", "n_head", "n_heads"], optional=True)) is not None: + self.gguf_writer.add_head_count(n_head) + logger.info(f"gguf: head count = {n_head}") + + if (n_head_kv := self.hparams.get("num_key_value_heads")) is not None: + self.gguf_writer.add_head_count_kv(n_head_kv) + logger.info(f"gguf: key-value head count = {n_head_kv}") + + if (rope_theta := self.hparams.get("rope_theta")) is not None: + self.gguf_writer.add_rope_freq_base(rope_theta) + logger.info(f"gguf: rope theta = {rope_theta}") + if (f_rms_eps := self.hparams.get("rms_norm_eps")) is not None: + self.gguf_writer.add_layer_norm_rms_eps(f_rms_eps) + logger.info(f"gguf: rms norm epsilon = {f_rms_eps}") + if (f_norm_eps := self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon"], optional=True)) is not None: + self.gguf_writer.add_layer_norm_eps(f_norm_eps) + logger.info(f"gguf: layer norm epsilon = {f_norm_eps}") + if (n_experts := self.hparams.get("num_local_experts")) is not None: + self.gguf_writer.add_expert_count(n_experts) + logger.info(f"gguf: expert count = {n_experts}") + if (n_experts_used := self.hparams.get("num_experts_per_tok")) is not None: + self.gguf_writer.add_expert_used_count(n_experts_used) + logger.info(f"gguf: experts used count = {n_experts_used}") + + if (head_dim := self.hparams.get("head_dim")) is not None: + self.gguf_writer.add_key_length(head_dim) + self.gguf_writer.add_value_length(head_dim) + + self.gguf_writer.add_file_type(self.ftype) + logger.info(f"gguf: file type = {self.ftype}") + + def write_vocab(self): + if len(self.gguf_writer.tensors) != 1: + raise ValueError('Splitting the vocabulary is not supported') + + self.prepare_metadata(vocab_only=True) + self.gguf_writer.write_header_to_file(path=self.fname_out) + self.gguf_writer.write_kv_data_to_file() + self.gguf_writer.close() + + def does_token_look_special(self, token: str | bytes) -> bool: + if isinstance(token, (bytes, bytearray)): + token_text = token.decode(encoding="utf-8") + elif isinstance(token, memoryview): + token_text = token.tobytes().decode(encoding="utf-8") + else: + token_text = token + + # Some models mark some added tokens which ought to be control tokens as not special. + # (e.g. command-r, command-r-plus, deepseek-coder, gemma{,-2}) + seems_special = token_text in ( + "", # deepseek-coder + "", "<2mass>", "[@BOS@]", # gemma{,-2} + ) + + seems_special = seems_special or (token_text.startswith("<|") and token_text.endswith("|>")) + seems_special = seems_special or (token_text.startswith("<｜") and token_text.endswith("｜>")) # deepseek-coder + + # TODO: should these be marked as UNUSED instead? (maybe not) + seems_special = seems_special or (token_text.startswith("")) # gemma{,-2} + + return seems_special + + # used for GPT-2 BPE and WordPiece vocabs + def get_vocab_base(self) -> tuple[list[str], list[int], str]: + tokens: list[str] = [] + toktypes: list[int] = [] + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) + assert max(tokenizer.vocab.values()) < vocab_size + + tokpre = self.get_vocab_base_pre(tokenizer) + + reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() + + added_tokens_decoder = tokenizer.added_tokens_decoder + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + else: + token: str = reverse_vocab[i] + if token in added_vocab: + # The tokenizer in llama.cpp assumes the CONTROL and USER_DEFINED tokens are pre-normalized. + # To avoid unexpected issues - we make sure to normalize non-normalized tokens + if not added_tokens_decoder[i].normalized: + previous_token = token + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) + if previous_token != token: + logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer") + + if added_tokens_decoder[i].special or self.does_token_look_special(token): + toktypes.append(gguf.TokenType.CONTROL) + else: + # NOTE: this was added for Gemma. + # Encoding and decoding the tokens above isn't sufficient for this case. + token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") # pre-normalize user-defined spaces + toktypes.append(gguf.TokenType.USER_DEFINED) + else: + toktypes.append(gguf.TokenType.NORMAL) + tokens.append(token) + + return tokens, toktypes, tokpre + + # NOTE: this function is generated by convert_hf_to_gguf_update.py + # do not modify it manually! + # ref: https://github.com/ggml-org/llama.cpp/pull/6920 + # Marker: Start get_vocab_base_pre + def get_vocab_base_pre(self, tokenizer) -> str: + # encoding this string and hashing the resulting tokens would (hopefully) give us a unique identifier that + # is specific for the BPE pre-tokenizer used by the model + # we will use this unique identifier to write a "tokenizer.ggml.pre" entry in the GGUF file which we can + # use in llama.cpp to implement the same pre-tokenizer + + chktxt = '\n \n\n \n\n\n \t \t\t \t\n \n \n \n \n🚀 (normal) 😶\u200d🌫️ (multiple emojis concatenated) ✅ 🦙🦙 3 33 333 3333 33333 333333 3333333 33333333 3.3 3..3 3...3 កាន់តែពិសេសអាច😁 ?我想在apple工作1314151天～ ------======= нещо на Български \'\'\'\'\'\'```````""""......!!!!!!?????? I\'ve been \'told he\'s there, \'RE you sure? \'M not sure I\'ll make it, \'D you like some tea? We\'Ve a\'lL' + + chktok = tokenizer.encode(chktxt) + chkhsh = sha256(str(chktok).encode()).hexdigest() + + logger.debug(f"chktok: {chktok}") + logger.debug(f"chkhsh: {chkhsh}") + + res = None + + # NOTE: if you get an error here, you need to update the convert_hf_to_gguf_update.py script + # or pull the latest version of the model from Huggingface + # don't edit the hashes manually! + if chkhsh == "0ef9807a4087ebef797fc749390439009c3b9eda9ad1a097abbe738f486c01e5": + # ref: https://huggingface.co/meta-llama/Meta-Llama-3-8B + res = "llama-bpe" + if chkhsh == "049ecf7629871e3041641907f3de7c733e4dbfdc736f57d882ba0b0845599754": + # ref: https://huggingface.co/deepseek-ai/deepseek-llm-7b-base + res = "deepseek-llm" + if chkhsh == "347715f544604f9118bb75ed199f68779f423cabb20db6de6f31b908d04d7821": + # ref: https://huggingface.co/deepseek-ai/deepseek-coder-6.7b-base + res = "deepseek-coder" + if chkhsh == "8aeee3860c56296a157a1fe2fad249ec40aa59b1bb5709f4ade11c4e6fe652ed": + # ref: https://huggingface.co/tiiuae/falcon-7b + res = "falcon" + if chkhsh == "0876d13b50744004aa9aeae05e7b0647eac9d801b5ba4668afc01e709c15e19f": + # ref: https://huggingface.co/BAAI/bge-small-en-v1.5 + res = "bert-bge" + if chkhsh == "9d032fcbd5501f4a38150912590928bfb36091efb5df11b8e2124b0390e3fb1e": + # ref: https://huggingface.co/tiiuae/Falcon3-7B-Base + res = "falcon3" + if chkhsh == "8e62295832751ca1e8f92f2226f403dea30dc5165e448b5bfa05af5340c64ec7": + # ref: https://huggingface.co/BAAI/bge-large-zh-v1.5 + res = "bert-bge-large" + if chkhsh == "b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166": + # ref: https://huggingface.co/mosaicml/mpt-7b + res = "mpt" + if chkhsh == "35d91631860c815f952d711435f48d356ebac988362536bed955d43bfa436e34": + # ref: https://huggingface.co/bigcode/starcoder2-3b + res = "starcoder" + if chkhsh == "3ce83efda5659b07b1ad37ca97ca5797ea4285d9b9ab0dc679e4a720c9da7454": + # ref: https://huggingface.co/openai-community/gpt2 + res = "gpt-2" + if chkhsh == "32d85c31273f8019248f2559fed492d929ea28b17e51d81d3bb36fff23ca72b3": + # ref: https://huggingface.co/stabilityai/stablelm-2-zephyr-1_6b + res = "stablelm2" + if chkhsh == "6221ad2852e85ce96f791f476e0b390cf9b474c9e3d1362f53a24a06dc8220ff": + # ref: https://huggingface.co/smallcloudai/Refact-1_6-base + res = "refact" + if chkhsh == "9c2227e4dd922002fb81bde4fc02b0483ca4f12911410dee2255e4987644e3f8": + # ref: https://huggingface.co/CohereForAI/c4ai-command-r-v01 + res = "command-r" + if chkhsh == "e636dc30a262dcc0d8c323492e32ae2b70728f4df7dfe9737d9f920a282b8aea": + # ref: https://huggingface.co/Qwen/Qwen1.5-7B + res = "qwen2" + if chkhsh == "b6dc8df998e1cfbdc4eac8243701a65afe638679230920b50d6f17d81c098166": + # ref: https://huggingface.co/allenai/OLMo-1.7-7B-hf + res = "olmo" + if chkhsh == "a8594e3edff7c29c003940395316294b2c623e09894deebbc65f33f1515df79e": + # ref: https://huggingface.co/databricks/dbrx-base + res = "dbrx" + if chkhsh == "c7699093ba4255a91e702aa38a596aa81669f3525dae06c2953267dde580f448": + # ref: https://huggingface.co/jinaai/jina-reranker-v1-tiny-en + res = "jina-v1-en" + if chkhsh == "0876d13b50744004aa9aeae05e7b0647eac9d801b5ba4668afc01e709c15e19f": + # ref: https://huggingface.co/jinaai/jina-embeddings-v2-base-en + res = "jina-v2-en" + if chkhsh == "171aeeedd6fb548d418a7461d053f11b6f1f1fc9b387bd66640d28a4b9f5c643": + # ref: https://huggingface.co/jinaai/jina-embeddings-v2-base-es + res = "jina-v2-es" + if chkhsh == "27949a2493fc4a9f53f5b9b029c82689cfbe5d3a1929bb25e043089e28466de6": + # ref: https://huggingface.co/jinaai/jina-embeddings-v2-base-de + res = "jina-v2-de" + if chkhsh == "c136ed14d01c2745d4f60a9596ae66800e2b61fa45643e72436041855ad4089d": + # ref: https://huggingface.co/abacusai/Smaug-Llama-3-70B-Instruct + res = "smaug-bpe" + if chkhsh == "c7ea5862a53e4272c035c8238367063e2b270d51faa48c0f09e9d5b54746c360": + # ref: https://huggingface.co/LumiOpen/Poro-34B-chat + res = "poro-chat" + if chkhsh == "7967bfa498ade6b757b064f31e964dddbb80f8f9a4d68d4ba7998fcf281c531a": + # ref: https://huggingface.co/jinaai/jina-embeddings-v2-base-code + res = "jina-v2-code" + if chkhsh == "7fc505bd3104ca1083b150b17d088b59534ede9bde81f0dd2090967d7fe52cee": + # ref: https://huggingface.co/LumiOpen/Viking-7B + res = "viking" + if chkhsh == "b53802fb28e26d645c3a310b34bfe07da813026ec7c7716883404d5e0f8b1901": + # ref: https://huggingface.co/core42/jais-13b + res = "jais" + if chkhsh == "7b3e7548e4308f52a76e8229e4e6cc831195d0d1df43aed21ac6c93da05fec5f": + # ref: https://huggingface.co/WisdomShell/CodeShell-7B + res = "codeshell" + if chkhsh == "63b97e4253352e6f357cc59ea5b583e3a680eaeaf2632188c2b952de2588485e": + # ref: https://huggingface.co/mistralai/Mistral-Nemo-Base-2407 + res = "tekken" + if chkhsh == "855059429035d75a914d1eda9f10a876752e281a054a7a3d421ef0533e5b6249": + # ref: https://huggingface.co/HuggingFaceTB/SmolLM-135M + res = "smollm" + if chkhsh == "3c30d3ad1d6b64202cd222813e7736c2db6e1bd6d67197090fc1211fbc612ae7": + # ref: https://huggingface.co/bigscience/bloom + res = "bloom" + if chkhsh == "bc01ce58980e1db43859146dc51b1758b3b88729b217a74792e9f8d43e479d21": + # ref: https://huggingface.co/TurkuNLP/gpt3-finnish-small + res = "gpt3-finnish" + if chkhsh == "4e2b24cc4770243d65a2c9ec19770a72f08cffc161adbb73fcbb6b7dd45a0aae": + # ref: https://huggingface.co/LGAI-EXAONE/EXAONE-3.0-7.8B-Instruct + res = "exaone" + if chkhsh == "fcace8b9cac38ce847670c970cd5892031a753a1ef381abd1d9af00f713da085": + # ref: https://huggingface.co/microsoft/phi-2 + res = "phi-2" + if chkhsh == "60824e3c0d9401f89943cbb2fff727f0e2d4c545ba4df2d6e4f09a6db0f5b450": + # ref: https://huggingface.co/facebook/chameleon-7b + res = "chameleon" + if chkhsh == "8b5a93ed704057481f240da0be7e7dca721d7f8f4755263b6807227a2cbeae65": + # ref: https://huggingface.co/sentence-transformers/stsb-roberta-base + res = "roberta-bpe" + if chkhsh == "ad851be1dba641f2e3711822f816db2c265f788b37c63b4e1aeacb9ee92de8eb": + # ref: https://huggingface.co/ai-sage/GigaChat-20B-A3B-instruct + res = "gigachat" + if chkhsh == "d4c8f286ea6b520b3d495c4455483cfa2302c0cfcd4be05d781b6a8a0a7cdaf1": + # ref: https://huggingface.co/Infinigence/Megrez-3B-Instruct + res = "megrez" + if chkhsh == "877081d19cf6996e2c4ff0e1236341e9b7bde288f5311a56a937f0afbbb3aeb5": + # ref: https://huggingface.co/deepseek-ai/DeepSeek-V3 + res = "deepseek-v3" + if chkhsh == "b3f499bb4255f8ca19fccd664443283318f2fd2414d5e0b040fbdd0cc195d6c5": + # ref: https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + res = "deepseek-r1-qwen" + if chkhsh == "ccc2ef013c104be7bae2965776d611e1d7a8a2a9c547dd93a682c9a9fc80352e": + # ref: https://huggingface.co/Xenova/gpt-4o + res = "gpt-4o" + if chkhsh == "7dec86086fcc38b66b7bc1575a160ae21cf705be7718b9d5598190d7c12db76f": + # ref: https://huggingface.co/UW/OLMo2-8B-SuperBPE-t180k + res = "superbpe" + if chkhsh == "1994ffd01900cfb37395608534236ecd63f2bd5995d6cb1004dda1af50240f15": + # ref: https://huggingface.co/trillionlabs/Trillion-7B-preview + res = "trillion" + if chkhsh == "96a5f08be6259352137b512d4157e333e21df7edd3fcd152990608735a65b224": + # ref: https://huggingface.co/inclusionAI/Ling-lite + res = "bailingmoe" + if chkhsh == "d353350c764d8c3b39c763113960e4fb4919bea5fbf208a0e3b22e8469dc7406": + # ref: https://huggingface.co/meta-llama/Llama-4-Scout-17B-16E-Instruct + res = "llama4" + if chkhsh == "0e9433cbbb161f89e264eb32e8e64bfe69e834973ffca5d41d3948a604a3e2a3": + # ref: https://huggingface.co/mistral-community/pixtral-12b + res = "pixtral" + if chkhsh == "d5f1dd6f980fec569fb218a81a7658ac45fc56b38c5a0adeb1c232fbe04ef5ec": + # ref: https://huggingface.co/ByteDance-Seed/Seed-Coder-8B-Base + res = "seed-coder" + if chkhsh == "b6e8e1518dc4305be2fe39c313ed643381c4da5db34a98f6a04c093f8afbe99b": + # ref: https://huggingface.co/THUDM/glm-4-9b-chat + res = "chatglm-bpe" + if chkhsh == "81d72c7348a9f0ebe86f23298d37debe0a5e71149e29bd283904c02262b27516": + # ref: https://huggingface.co/THUDM/glm-4-9b-chat + res = "chatglm-bpe" + if chkhsh == "a1336059768a55c99a734006ffb02203cd450fed003e9a71886c88acf24fdbc2": + # ref: https://huggingface.co/THUDM/glm-4-9b-hf + res = "glm4" + if chkhsh == "1431a23e583c97432bc230bff598d103ddb5a1f89960c8f1d1051aaa944d0b35": + # ref: https://huggingface.co/sapienzanlp/Minerva-7B-base-v1.0 + res = "minerva-7b" + + if res is None: + logger.warning("\n") + logger.warning("**************************************************************************************") + logger.warning("** WARNING: The BPE pre-tokenizer was not recognized!") + logger.warning("** There are 2 possible reasons for this:") + logger.warning("** - the model has not been added to convert_hf_to_gguf_update.py yet") + logger.warning("** - the pre-tokenization config has changed upstream") + logger.warning("** Check your model files and convert_hf_to_gguf_update.py and update them accordingly.") + logger.warning("** ref: https://github.com/ggml-org/llama.cpp/pull/6920") + logger.warning("**") + logger.warning(f"** chkhsh: {chkhsh}") + logger.warning("**************************************************************************************") + logger.warning("\n") + raise NotImplementedError("BPE pre-tokenizer was not recognized - update get_vocab_base_pre()") + + logger.debug(f"tokenizer.ggml.pre: {repr(res)}") + logger.debug(f"chkhsh: {chkhsh}") + + return res + # Marker: End get_vocab_base_pre + + def _set_vocab_none(self) -> None: + self.gguf_writer.add_tokenizer_model("none") + + def _set_vocab_gpt2(self) -> None: + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) + + def _set_vocab_qwen(self): + dir_model = self.dir_model + hparams = self.hparams + tokens: list[str] = [] + toktypes: list[int] = [] + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True) + vocab_size = hparams["vocab_size"] + assert max(tokenizer.get_vocab().values()) < vocab_size + + tokpre = self.get_vocab_base_pre(tokenizer) + + merges = [] + vocab = {} + mergeable_ranks = tokenizer.mergeable_ranks + for token, rank in mergeable_ranks.items(): + vocab[QwenModel.token_bytes_to_string(token)] = rank + if len(token) == 1: + continue + merged = QwenModel.bpe(mergeable_ranks, token, max_rank=rank) + assert len(merged) == 2 + merges.append(' '.join(map(QwenModel.token_bytes_to_string, merged))) + + # for this kind of tokenizer, added_vocab is not a subset of vocab, so they need to be combined + added_vocab = tokenizer.special_tokens + reverse_vocab = {id_ : encoded_tok for encoded_tok, id_ in {**vocab, **added_vocab}.items()} + + for i in range(vocab_size): + if i not in reverse_vocab: + tokens.append(f"[PAD{i}]") + toktypes.append(gguf.TokenType.UNUSED) + elif reverse_vocab[i] in added_vocab: + tokens.append(reverse_vocab[i]) + toktypes.append(gguf.TokenType.CONTROL) + else: + tokens.append(reverse_vocab[i]) + toktypes.append(gguf.TokenType.NORMAL) + + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(dir_model, load_merges=False) + special_vocab.merges = merges + # only add special tokens when they were not already loaded from config.json + if len(special_vocab.special_token_ids) == 0: + special_vocab._set_special_token("bos", tokenizer.special_tokens["<|endoftext|>"]) + special_vocab._set_special_token("eos", tokenizer.special_tokens["<|endoftext|>"]) + # this one is usually not in config.json anyway + special_vocab._set_special_token("unk", tokenizer.special_tokens["<|endoftext|>"]) + special_vocab.add_to_gguf(self.gguf_writer) + + def _set_vocab_sentencepiece(self, add_to_gguf=True): + tokens, scores, toktypes = self._create_vocab_sentencepiece() + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def _create_vocab_sentencepiece(self): + from sentencepiece import SentencePieceProcessor + + tokenizer_path = self.dir_model / 'tokenizer.model' + + if not tokenizer_path.is_file(): + raise FileNotFoundError(f"File not found: {tokenizer_path}") + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.find_hparam([ + "vocab_size_per_layer_input", # gemma3n + "vocab_size", + ], optional=True) or tokenizer.vocab_size() + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + for token_id in range(tokenizer.vocab_size()): + if token_id >= vocab_size: + logger.warning(f'ignore tokens from {token_id}: id is out of range, max={vocab_size - 1}') + break + + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + added_tokens_file = self.dir_model / 'added_tokens.json' + if added_tokens_file.is_file(): + with open(added_tokens_file, "r", encoding="utf-8") as f: + added_tokens_json = json.load(f) + for key in added_tokens_json: + token_id = added_tokens_json[key] + if token_id >= vocab_size: + logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + + tokens[token_id] = key.encode("utf-8") + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + added_tokens_decoder = tokenizer_config_json.get("added_tokens_decoder", {}) + for token_id, token_data in added_tokens_decoder.items(): + token_id = int(token_id) + token: str = token_data["content"] + if token_id >= vocab_size: + logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + if toktypes[token_id] != SentencePieceTokenTypes.UNUSED: + if tokens[token_id] != token.encode("utf-8"): + logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token!r}') + if token_data.get("special") or self.does_token_look_special(token): + toktypes[token_id] = SentencePieceTokenTypes.CONTROL + else: + token = token.replace(b"\xe2\x96\x81".decode("utf-8"), " ") # pre-normalize user-defined spaces + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + + scores[token_id] = -1000.0 + tokens[token_id] = token.encode("utf-8") + + if vocab_size > len(tokens): + pad_count = vocab_size - len(tokens) + logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]") + for i in range(1, pad_count + 1): + tokens.append(bytes(f"[PAD{i}]", encoding="utf-8")) + scores.append(-1000.0) + toktypes.append(SentencePieceTokenTypes.UNUSED) + + return tokens, scores, toktypes + + def _set_vocab_llama_hf(self): + vocab = gguf.LlamaHfVocab(self.dir_model) + tokens = [] + scores = [] + toktypes = [] + + for text, score, toktype in vocab.all_tokens(): + tokens.append(text) + scores.append(score) + toktypes.append(toktype) + + assert len(tokens) == vocab.vocab_size + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def _set_vocab_rwkv_world(self): + assert (self.dir_model / "rwkv_vocab_v20230424.txt").is_file() + vocab_size = self.hparams.get("vocab_size", 65536) + + tokens: list[bytes] = [''.encode("utf-8")] + toktypes: list[int] = [gguf.TokenType.CONTROL] + + with open(self.dir_model / "rwkv_vocab_v20230424.txt", "r", encoding="utf-8") as f: + lines = f.readlines() + for line in lines: + parts = line.split(' ') + assert len(parts) >= 3 + token, token_len = ast.literal_eval(' '.join(parts[1:-1])), int(parts[-1]) + token = token.encode("utf-8") if isinstance(token, str) else token + assert isinstance(token, bytes) + assert len(token) == token_len + token_text: str = repr(token)[2:-1] # "b'\xff'" -> "\xff" + tokens.append(token_text.encode("utf-8")) + toktypes.append(gguf.TokenType.NORMAL) + remainder = vocab_size - len(tokens) + assert remainder >= 0 + for i in range(len(tokens), vocab_size): + tokens.append(f"[PAD{i}]".encode("utf-8")) + toktypes.append(gguf.TokenType.UNUSED) + + self.gguf_writer.add_tokenizer_model("rwkv") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=False) + special_vocab.chat_template = "rwkv-world" + # hack: Add '\n\n' as the EOT token to make it chat normally + special_vocab._set_special_token("eot", 261) + # hack: Override these as they have already been set (incorrectly) + special_vocab.special_token_ids["bos"] = 0 + special_vocab.special_token_ids["eos"] = 0 + + special_vocab.add_to_gguf(self.gguf_writer) + + def _set_vocab_builtin(self, model_name: Literal["gpt-neox", "llama-spm"], vocab_size: int): + tokenizer_path = Path(sys.path[0]) / "models" / f"ggml-vocab-{model_name}.gguf" + logger.warning(f"Using tokenizer from '{os.path.relpath(tokenizer_path, os.getcwd())}'") + vocab_reader = gguf.GGUFReader(tokenizer_path, "r") + + default_pre = "mpt" if model_name == "gpt-neox" else "default" + + field = vocab_reader.get_field(gguf.Keys.Tokenizer.MODEL) + assert field # tokenizer model + self.gguf_writer.add_tokenizer_model(bytes(field.parts[-1]).decode("utf-8")) + + field = vocab_reader.get_field(gguf.Keys.Tokenizer.PRE) + self.gguf_writer.add_tokenizer_pre(bytes(field.parts[-1]).decode("utf-8") if field else default_pre) + + field = vocab_reader.get_field(gguf.Keys.Tokenizer.LIST) + assert field # token list + self.gguf_writer.add_token_list([bytes(field.parts[i]) for i in field.data][:vocab_size]) + + if model_name == "llama-spm": + field = vocab_reader.get_field(gguf.Keys.Tokenizer.SCORES) + assert field # token scores + self.gguf_writer.add_token_scores([field.parts[i].tolist()[0] for i in field.data][:vocab_size]) + + field = vocab_reader.get_field(gguf.Keys.Tokenizer.TOKEN_TYPE) + assert field # token types + self.gguf_writer.add_token_types([field.parts[i].tolist()[0] for i in field.data][:vocab_size]) + + if model_name != "llama-spm": + field = vocab_reader.get_field(gguf.Keys.Tokenizer.MERGES) + assert field # token merges + self.gguf_writer.add_token_merges([bytes(field.parts[i]) for i in field.data]) + + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.BOS_ID)) is not None: + self.gguf_writer.add_bos_token_id(field.parts[-1].tolist()[0]) + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.EOS_ID)) is not None: + self.gguf_writer.add_eos_token_id(field.parts[-1].tolist()[0]) + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.UNK_ID)) is not None: + self.gguf_writer.add_unk_token_id(field.parts[-1].tolist()[0]) + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.PAD_ID)) is not None: + self.gguf_writer.add_pad_token_id(field.parts[-1].tolist()[0]) + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.ADD_BOS)) is not None: + self.gguf_writer.add_add_bos_token(field.parts[-1].tolist()[0]) + if (field := vocab_reader.get_field(gguf.Keys.Tokenizer.ADD_EOS)) is not None: + self.gguf_writer.add_add_eos_token(field.parts[-1].tolist()[0]) + + def _try_set_pooling_type(self) -> None: + # get pooling path + pooling_path = None + module_path = self.dir_model / "modules.json" + if module_path.is_file(): + with open(module_path, encoding="utf-8") as f: + modules = json.load(f) + for mod in modules: + if mod["type"] == "sentence_transformers.models.Pooling": + pooling_path = mod["path"] + break + + # get pooling type + if pooling_path is not None: + with open(self.dir_model / pooling_path / "config.json", encoding="utf-8") as f: + pooling = json.load(f) + if pooling["pooling_mode_mean_tokens"]: + pooling_type = gguf.PoolingType.MEAN + elif pooling["pooling_mode_cls_token"]: + pooling_type = gguf.PoolingType.CLS + elif pooling["pooling_mode_lasttoken"]: + pooling_type = gguf.PoolingType.LAST + else: + raise NotImplementedError("Only MEAN, CLS, and LAST pooling types supported") + self.gguf_writer.add_pooling_type(pooling_type) + + +class MmprojModel(ModelBase): + model_type = ModelType.MMPROJ + model_arch = gguf.MODEL_ARCH.MMPROJ + preprocessor_config: dict[str, Any] + global_config: dict[str, Any] + + n_block_keys = ["n_layers", "num_hidden_layers", "n_layer", "num_layers", "depth"] + + has_vision_encoder: bool = True # by default + has_audio_encoder: bool = False + + # for models having multiple encoders, we need to separate their hparams + hparams_vision: dict[str, Any] | None = None + hparams_audio: dict[str, Any] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if self.model_arch != gguf.MODEL_ARCH.MMPROJ: + raise TypeError("MmprojModel must be subclassed with model_arch = gguf.MODEL_ARCH.MMPROJ") + + # get n_embd of the text model + if "text_config" not in self.hparams: + self.hparams["text_config"] = {} + if "audio_config" not in self.hparams: + self.hparams["audio_config"] = {} + text_config = {**self.hparams, **self.hparams["text_config"]} + self.n_embd_text = text_config.get("hidden_size", text_config.get("n_embd", 0)) + assert self.n_embd_text > 0, "n_embd not found in hparams" + + # move vision config to the top level, while preserving the original hparams in global_config + import copy + self.global_config = copy.deepcopy(self.hparams) + self.hparams_vision = self.get_vision_config() + self.hparams_audio = self.get_audio_config() + + if self.hparams_vision is None and self.hparams_audio is None: + raise ValueError("vision_config / audio_config not found in hparams") + + # for compat with vision-only models + self.hparams = self.hparams_vision or self.hparams_audio or self.hparams + + # TODO @ngxson : this is a hack to support both vision and audio encoders + have_multiple_encoders = self.has_audio_encoder and self.has_vision_encoder + self.block_count = 128 if have_multiple_encoders else self.find_hparam(self.n_block_keys, True) + self.tensor_map = gguf.get_tensor_name_map(gguf.MODEL_ARCH.MMPROJ, self.block_count) + + # load preprocessor config + with open(self.dir_model / "preprocessor_config.json", "r", encoding="utf-8") as f: + self.preprocessor_config = json.load(f) + + def get_vision_config(self) -> dict[str, Any] | None: + return self.global_config.get("vision_config") + + def get_audio_config(self) -> dict[str, Any] | None: + return self.global_config.get("audio_config") + + def set_type(self): + self.gguf_writer.add_type(gguf.GGUFType.MMPROJ) + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + + if self.has_vision_encoder: + self.gguf_writer.add_clip_has_vision_encoder(True) + self.gguf_writer.add_vision_projection_dim(self.n_embd_text) + + # vision config + self.gguf_writer.add_vision_image_size(self.find_vparam(["image_size"])) + self.gguf_writer.add_vision_patch_size(self.find_vparam(["patch_size"])) + self.gguf_writer.add_vision_embedding_length(self.find_vparam(["hidden_size"])) + self.gguf_writer.add_vision_feed_forward_length(self.find_vparam(["intermediate_size"])) + self.gguf_writer.add_vision_block_count(self.find_vparam(self.n_block_keys)) + self.gguf_writer.add_vision_head_count(self.find_vparam(["num_attention_heads"])) + + # preprocessor config + self.gguf_writer.add_vision_image_mean(self.preprocessor_config["image_mean"]) + self.gguf_writer.add_vision_image_std(self.preprocessor_config["image_std"]) + + if self.has_audio_encoder: + self.gguf_writer.add_clip_has_audio_encoder(True) + self.gguf_writer.add_audio_projection_dim(self.n_embd_text) + + # audio config + self.gguf_writer.add_audio_embedding_length(self.find_aparam(["hidden_size"])) + self.gguf_writer.add_audio_feed_forward_length(self.find_aparam(["intermediate_size"])) + self.gguf_writer.add_audio_block_count(self.find_aparam(self.n_block_keys)) + self.gguf_writer.add_audio_head_count(self.find_aparam(["num_attention_heads"])) + + if not self.has_vision_encoder and not self.has_audio_encoder: + raise ValueError("MmprojModel must have either vision or audio encoder") + + def write_vocab(self): + raise ValueError("MmprojModel does not support vocab writing") + + def find_vparam(self, keys: Iterable[str], optional: bool = False) -> Any: + assert self.hparams_vision is not None + return self._find_param(self.hparams_vision, keys, optional) + + def find_aparam(self, keys: Iterable[str], optional: bool = False) -> Any: + assert self.hparams_audio is not None + return self._find_param(self.hparams_audio, keys, optional) + + def _find_param(self, obj: dict[str, Any], keys: Iterable[str], optional: bool = False) -> Any: + key = next((k for k in keys if k in obj), None) + if key is not None: + return obj[key] + if optional: + return None + raise KeyError(f"could not find any of: {keys}") + + +@ModelBase.register("GPTNeoXForCausalLM") +class GPTNeoXModel(TextModel): + model_arch = gguf.MODEL_ARCH.GPTNEOX + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + + self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_rope_dimension_count( + int(self.hparams["rotary_pct"] * (self.hparams["hidden_size"] // self.hparams["num_attention_heads"])), + ) + self.gguf_writer.add_head_count(self.hparams["num_attention_heads"]) + self.gguf_writer.add_parallel_residual(self.hparams.get("use_parallel_residual", True)) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_eps"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads")) + n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed")) + + tensors: list[tuple[str, Tensor]] = [] + + if re.match(r"gpt_neox\.layers\.\d+\.attention\.query_key_value\.weight", name): + # Map bloom-style qkv_linear to gpt-style qkv_linear + # bloom: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py#L238-L252 # noqa + # gpt-2: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L312 # noqa + qkv_weights = data_torch.reshape((n_head, 3, n_embed // n_head, n_embed)) + data_torch = torch.cat( + ( + qkv_weights[:, 0, :, :].reshape((-1, n_embed)), + qkv_weights[:, 1, :, :].reshape((-1, n_embed)), + qkv_weights[:, 2, :, :].reshape((-1, n_embed)), + ), + dim=0, + ) + logger.info("re-format attention.linear_qkv.weight") + elif re.match(r"gpt_neox\.layers\.\d+\.attention\.query_key_value\.bias", name): + qkv_bias = data_torch.reshape((n_head, 3, n_embed // n_head)) + data_torch = torch.cat( + ( + qkv_bias[:, 0, :].reshape((n_embed,)), + qkv_bias[:, 1, :].reshape((n_embed,)), + qkv_bias[:, 2, :].reshape((n_embed,)), + ), + dim=0, + ) + logger.info("re-format attention.linear_qkv.bias") + + tensors.append((self.map_tensor_name(name), data_torch)) + + return tensors + + +@ModelBase.register("BloomForCausalLM", "BloomModel") +class BloomModel(TextModel): + model_arch = gguf.MODEL_ARCH.BLOOM + + def set_gguf_parameters(self): + n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed")) + n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads")) + self.gguf_writer.add_context_length(self.hparams.get("seq_length", n_embed)) + self.gguf_writer.add_embedding_length(n_embed) + self.gguf_writer.add_feed_forward_length(4 * n_embed) + self.gguf_writer.add_block_count(self.hparams["n_layer"]) + self.gguf_writer.add_head_count(n_head) + self.gguf_writer.add_head_count_kv(n_head) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads")) + n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed")) + + name = re.sub(r'transformer\.', '', name) + + tensors: list[tuple[str, Tensor]] = [] + + if re.match(r"h\.\d+\.self_attention\.query_key_value\.weight", name): + # Map bloom-style qkv_linear to gpt-style qkv_linear + # bloom: https://github.com/huggingface/transformers/blob/main/src/transformers/models/bloom/modeling_bloom.py#L238-L252 # noqa + # gpt-2: https://github.com/huggingface/transformers/blob/main/src/transformers/models/gpt2/modeling_gpt2.py#L312 # noqa + qkv_weights = data_torch.reshape((n_head, 3, n_embed // n_head, n_embed)) + data_torch = torch.cat( + ( + qkv_weights[:, 0, :, :].reshape((-1, n_embed)), + qkv_weights[:, 1, :, :].reshape((-1, n_embed)), + qkv_weights[:, 2, :, :].reshape((-1, n_embed)), + ), + dim=0, + ) + logger.info("re-format attention.linear_qkv.weight") + elif re.match(r"h\.\d+\.self_attention\.query_key_value\.bias", name): + qkv_bias = data_torch.reshape((n_head, 3, n_embed // n_head)) + data_torch = torch.cat( + ( + qkv_bias[:, 0, :].reshape((n_embed,)), + qkv_bias[:, 1, :].reshape((n_embed,)), + qkv_bias[:, 2, :].reshape((n_embed,)), + ), + dim=0, + ) + logger.info("re-format attention.linear_qkv.bias") + + tensors.append((self.map_tensor_name(name), data_torch)) + + return tensors + + +@ModelBase.register("MPTForCausalLM") +class MPTModel(TextModel): + model_arch = gguf.MODEL_ARCH.MPT + + def set_vocab(self): + try: + self._set_vocab_gpt2() + except Exception: + # Fallback for SEA-LION model + self._set_vocab_sentencepiece() + self.gguf_writer.add_add_bos_token(False) + self.gguf_writer.add_pad_token_id(3) + self.gguf_writer.add_eos_token_id(1) + self.gguf_writer.add_unk_token_id(0) + + def set_gguf_parameters(self): + block_count = self.hparams["n_layers"] + self.gguf_writer.add_context_length(self.hparams["max_seq_len"]) + self.gguf_writer.add_embedding_length(self.hparams["d_model"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["d_model"]) + self.gguf_writer.add_head_count(self.hparams["n_heads"]) + if kv_n_heads := self.hparams["attn_config"].get("kv_n_heads"): + self.gguf_writer.add_head_count_kv(kv_n_heads) + self.gguf_writer.add_layer_norm_eps(1e-5) + if self.hparams["attn_config"]["clip_qkv"] is not None: + self.gguf_writer.add_clamp_kqv(self.hparams["attn_config"]["clip_qkv"]) + if self.hparams["attn_config"]["alibi"]: + self.gguf_writer.add_max_alibi_bias(self.hparams["attn_config"]["alibi_bias_max"]) + else: + self.gguf_writer.add_max_alibi_bias(0.0) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if "scales" in name: + new_name = self.map_tensor_name(name, try_suffixes=(".weight", ".bias", ".scales")) + new_name = new_name.replace("scales", "act.scales") + else: + new_name = self.map_tensor_name(name, try_suffixes=(".weight", ".bias")) + + return [(new_name, data_torch)] + + +@ModelBase.register("OrionForCausalLM") +class OrionModel(TextModel): + model_arch = gguf.MODEL_ARCH.ORION + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + head_count = self.hparams["num_attention_heads"] + head_count_kv = self.hparams.get("num_key_value_heads", head_count) + + ctx_length = 0 + if "max_sequence_length" in self.hparams: + ctx_length = self.hparams["max_sequence_length"] + elif "max_position_embeddings" in self.hparams: + ctx_length = self.hparams["max_position_embeddings"] + elif "model_max_length" in self.hparams: + ctx_length = self.hparams["model_max_length"] + else: + raise ValueError("gguf: can not find ctx length parameter.") + + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_tensor_data_layout("Meta AI original pth") + self.gguf_writer.add_context_length(ctx_length) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_head_count(head_count) + self.gguf_writer.add_head_count_kv(head_count_kv) + # note: config provides rms norm but it is actually layer norm + # ref: https://huggingface.co/OrionStarAI/Orion-14B-Chat/blob/276a17221ce42beb45f66fac657a41540e71f4f5/modeling_orion.py#L570-L571 + self.gguf_writer.add_layer_norm_eps(self.hparams["rms_norm_eps"]) + + +@ModelBase.register("BaichuanForCausalLM", "BaiChuanForCausalLM") +class BaichuanModel(TextModel): + model_arch = gguf.MODEL_ARCH.BAICHUAN + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + head_count = self.hparams["num_attention_heads"] + head_count_kv = self.hparams.get("num_key_value_heads", head_count) + + ctx_length = 0 + if "max_sequence_length" in self.hparams: + ctx_length = self.hparams["max_sequence_length"] + elif "max_position_embeddings" in self.hparams: + ctx_length = self.hparams["max_position_embeddings"] + elif "model_max_length" in self.hparams: + ctx_length = self.hparams["model_max_length"] + else: + raise ValueError("gguf: can not find ctx length parameter.") + + self.gguf_writer.add_tensor_data_layout("Meta AI original pth") + self.gguf_writer.add_context_length(ctx_length) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"]) + self.gguf_writer.add_head_count(head_count) + self.gguf_writer.add_head_count_kv(head_count_kv) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_file_type(self.ftype) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + head_count = self.hparams["num_attention_heads"] + head_count_kv = self.hparams.get("num_key_value_heads", head_count) + + tensors: list[tuple[str, Tensor]] = [] + + if bid is not None and name == f"model.layers.{bid}.self_attn.W_pack.weight": + logger.info(f"Unpacking and permuting layer {bid}") + tensors = [ + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), + self._reverse_hf_permute_part(data_torch, 0, head_count, head_count)), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), + self._reverse_hf_permute_part(data_torch, 1, head_count, head_count_kv)), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), + self._reverse_hf_part(data_torch, 2)), + ] + else: + tensors = [(self.map_tensor_name(name), data_torch)] + + return tensors + + def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor: + if n_kv_head is not None and n_head != n_kv_head: + n_head //= n_kv_head + + return ( + weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape) + ) + + def _reverse_hf_permute_part( + self, weights: Tensor, n_part: int, n_head: int, n_head_kv: int | None = None, + ) -> Tensor: + r = weights.shape[0] // 3 + return self._reverse_hf_permute(weights[r * n_part:r * n_part + r, ...], n_head, n_head_kv) + + def _reverse_hf_part(self, weights: Tensor, n_part: int) -> Tensor: + r = weights.shape[0] // 3 + return weights[r * n_part:r * n_part + r, ...] + + +@ModelBase.register("XverseForCausalLM") +class XverseModel(TextModel): + model_arch = gguf.MODEL_ARCH.XVERSE + + def set_vocab(self): + assert (self.dir_model / "tokenizer.json").is_file() + dir_model = self.dir_model + hparams = self.hparams + + tokens: list[bytes] = [] + toktypes: list[int] = [] + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(dir_model) + vocab_size = hparams.get("vocab_size", len(tokenizer.vocab)) + # Since we are checking the maximum index, we need to ensure it's strictly less than vocab_size, + # because vocab_size is the count of items, and indexes start at 0. + max_vocab_index = max(tokenizer.get_vocab().values()) + if max_vocab_index >= vocab_size: + raise ValueError("Vocabulary size exceeds expected maximum size.") + + reverse_vocab: dict[int, str] = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() + + for token_id in range(vocab_size): + token_text = reverse_vocab[token_id].encode('utf-8') + # replace "\x00" to string with length > 0 + if token_text == b"\x00": + toktype = gguf.TokenType.BYTE # special + token_text = f"<{token_text}>".encode('utf-8') + elif re.fullmatch(br"<0x[0-9A-Fa-f]{2}>", token_text): + toktype = gguf.TokenType.BYTE # special + elif reverse_vocab[token_id] in added_vocab: + if tokenizer.added_tokens_decoder[token_id].special: + toktype = gguf.TokenType.CONTROL + else: + toktype = gguf.TokenType.USER_DEFINED + else: + toktype = gguf.TokenType.NORMAL + + tokens.append(token_text) + toktypes.append(toktype) + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + head_count = self.hparams["num_attention_heads"] + head_count_kv = self.hparams.get("num_key_value_heads", head_count) + + ctx_length = 0 + if "max_sequence_length" in self.hparams: + ctx_length = self.hparams["max_sequence_length"] + elif "max_position_embeddings" in self.hparams: + ctx_length = self.hparams["max_position_embeddings"] + elif "model_max_length" in self.hparams: + ctx_length = self.hparams["model_max_length"] + else: + raise ValueError("gguf: can not find ctx length parameter.") + + self.gguf_writer.add_tensor_data_layout("Meta AI original pth") + self.gguf_writer.add_context_length(ctx_length) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"]) + self.gguf_writer.add_head_count(head_count) + self.gguf_writer.add_head_count_kv(head_count_kv) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_file_type(self.ftype) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + head_count = self.hparams["num_attention_heads"] + head_count_kv = self.hparams.get("num_key_value_heads", head_count) + + # HF models permute some of the tensors, so we need to undo that + if name.endswith("q_proj.weight"): + data_torch = self._reverse_hf_permute(data_torch, head_count, head_count) + if name.endswith("k_proj.weight"): + data_torch = self._reverse_hf_permute(data_torch, head_count, head_count_kv) + + return [(self.map_tensor_name(name), data_torch)] + + def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor: + if n_kv_head is not None and n_head != n_kv_head: + n_head //= n_kv_head + + return ( + weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape) + ) + + +@ModelBase.register("FalconForCausalLM", "RWForCausalLM") +class FalconModel(TextModel): + model_arch = gguf.MODEL_ARCH.FALCON + + def set_gguf_parameters(self): + block_count = self.hparams.get("num_hidden_layers") + if block_count is None: + block_count = self.hparams["n_layer"] # old name + + n_head = self.hparams.get("num_attention_heads") + if n_head is None: + n_head = self.hparams["n_head"] # old name + + n_head_kv = self.hparams.get("num_kv_heads") + if n_head_kv is None: + n_head_kv = self.hparams.get("n_head_kv", 1) # old name + + self.gguf_writer.add_context_length(2048) # not in config.json + self.gguf_writer.add_tensor_data_layout("jploski") # qkv tensor transform + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(n_head) + self.gguf_writer.add_head_count_kv(n_head_kv) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + # QKV tensor transform + # The original query_key_value tensor contains n_head_kv "kv groups", + # each consisting of n_head/n_head_kv query weights followed by one key + # and one value weight (shared by all query heads in the kv group). + # This layout makes it a big pain to work with in GGML. + # So we rearrange them here,, so that we have n_head query weights + # followed by n_head_kv key weights followed by n_head_kv value weights, + # in contiguous fashion. + # ref: https://github.com/jploski/ggml/blob/falcon40b/examples/falcon/convert-hf-to-ggml.py + + if "query_key_value" in name: + n_head = self.find_hparam(["num_attention_heads", "n_head"]) + n_head_kv = self.find_hparam(["num_kv_heads", "n_head_kv"], optional=True) or 1 + head_dim = self.hparams["hidden_size"] // n_head + + qkv = data_torch.view(n_head_kv, n_head // n_head_kv + 2, head_dim, head_dim * n_head) + q = qkv[:, :-2].reshape(n_head * head_dim, head_dim * n_head) + k = qkv[:, [-2]].reshape(n_head_kv * head_dim, head_dim * n_head) + v = qkv[:, [-1]].reshape(n_head_kv * head_dim, head_dim * n_head) + data_torch = torch.cat((q, k, v)).reshape_as(data_torch) + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("GPTBigCodeForCausalLM") +class StarCoderModel(TextModel): + model_arch = gguf.MODEL_ARCH.STARCODER + + def set_gguf_parameters(self): + block_count = self.hparams["n_layer"] + + self.gguf_writer.add_context_length(self.hparams["n_positions"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["n_embd"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_head_count_kv(1) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + +@ModelBase.register("GPTRefactForCausalLM") +class RefactModel(TextModel): + model_arch = gguf.MODEL_ARCH.REFACT + + def set_vocab(self): + super().set_vocab() + + # TODO: how to determine special FIM tokens automatically? + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=False, + special_token_types = ['prefix', 'suffix', 'middle', 'eot']) + special_vocab._set_special_token("prefix", 1) + special_vocab._set_special_token("suffix", 3) + special_vocab._set_special_token("middle", 2) + special_vocab.chat_template = None # do not add it twice + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + hidden_dim = self.hparams["n_embd"] + inner_dim = 4 * hidden_dim + hidden_dim = int(2 * inner_dim / 3) + multiple_of = 256 + ff_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + + block_count = self.hparams["n_layer"] + + # refact uses Alibi. So this is from config.json which might be used by training. + self.gguf_writer.add_context_length(self.hparams["n_positions"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + + self.gguf_writer.add_feed_forward_length(ff_dim) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_head_count_kv(1) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + hidden_dim = self.hparams["n_embd"] + inner_dim = 4 * hidden_dim + hidden_dim = int(2 * inner_dim / 3) + multiple_of = 256 + ff_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) + n_head = self.hparams["n_head"] + n_head_kv = 1 + head_dim = self.hparams["n_embd"] // n_head + + tensors: list[tuple[str, Tensor]] = [] + + if bid is not None: + if name == f"transformer.h.{bid}.attn.kv.weight": + tensors.append((self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), data_torch[:n_head_kv * head_dim])) + tensors.append((self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), data_torch[n_head_kv * head_dim:])) + elif name == f"transformer.h.{bid}.attn.q.weight": + tensors.append((self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), data_torch)) + elif name == f"transformer.h.{bid}.mlp.gate_up_proj.weight": + tensors.append((self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), data_torch[:ff_dim])) + tensors.append((self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), data_torch[ff_dim:])) + + if len(tensors) == 0: + tensors.append((self.map_tensor_name(name), data_torch)) + + return tensors + + +@ModelBase.register("StableLmForCausalLM", "StableLMEpochForCausalLM", "LlavaStableLMEpochForCausalLM") +class StableLMModel(TextModel): + model_arch = gguf.MODEL_ARCH.STABLELM + + def set_vocab(self): + if (self.dir_model / "tokenizer.json").is_file(): + self._set_vocab_gpt2() + else: + # StableLM 2 1.6B used to have a vocab in a similar format to Qwen's vocab + self._set_vocab_qwen() + + def set_gguf_parameters(self): + hparams = self.hparams + block_count = hparams["num_hidden_layers"] + + self.gguf_writer.add_context_length(hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + rotary_factor = self.find_hparam(["partial_rotary_factor", "rope_pct"]) + self.gguf_writer.add_rope_dimension_count(int(rotary_factor * (hparams["hidden_size"] // hparams["num_attention_heads"]))) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(hparams["num_key_value_heads"]) + self.gguf_writer.add_parallel_residual(hparams["use_parallel_residual"] if "use_parallel_residual" in hparams else True) + self.gguf_writer.add_layer_norm_eps(self.find_hparam(["layer_norm_eps", "norm_eps"])) + self.gguf_writer.add_file_type(self.ftype) + + _q_norms: list[dict[str, Tensor]] | None = None + _k_norms: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams["num_key_value_heads"] + + if name.find("q_layernorm.norms") != -1: + assert bid is not None + + if self._q_norms is None: + self._q_norms = [{} for _ in range(self.block_count)] + + self._q_norms[bid][name] = data_torch + + if len(self._q_norms[bid]) >= n_head: + return self._stack_qk_norm(bid, n_head, self._q_norms[bid], "q_layernorm") + else: + return [] + + if name.find("k_layernorm.norms") != -1: + assert bid is not None + + if self._k_norms is None: + self._k_norms = [{} for _ in range(self.block_count)] + + self._k_norms[bid][name] = data_torch + + if len(self._k_norms[bid]) >= n_kv_head: + return self._stack_qk_norm(bid, n_kv_head, self._k_norms[bid], "k_layernorm") + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def _stack_qk_norm(self, bid: int, n_head: int, norms: dict[str, Tensor], layer_name: str = "q_layernorm"): + datas: list[Tensor] = [] + # extract the norms in order + for xid in range(n_head): + ename = f"model.layers.{bid}.self_attn.{layer_name}.norms.{xid}.weight" + datas.append(norms[ename]) + del norms[ename] + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.self_attn.{layer_name}.weight" + new_name = self.map_tensor_name(merged_name) + + return [(new_name, data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._q_norms is not None or self._k_norms is not None: + # flatten two `list[dict[str, Tensor]]` into a single `list[str]` + norms = ( + [k for d in self._q_norms for k in d.keys()] if self._q_norms is not None else [] + ) + ( + [k for d in self._k_norms for k in d.keys()] if self._k_norms is not None else [] + ) + if len(norms) > 0: + raise ValueError(f"Unprocessed norms: {norms}") + + +@ModelBase.register( + "LLaMAForCausalLM", + "LlamaForCausalLM", + "MistralForCausalLM", + "MixtralForCausalLM", + "VLlama3ForCausalLM", + "LlavaForConditionalGeneration", + "LlamaModel") +class LlamaModel(TextModel): + model_arch = gguf.MODEL_ARCH.LLAMA + undo_permute = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # fix for SmolVLM2, missing `num_attention_heads` in config.json + if self.hf_arch == "VLlama3ForCausalLM": + self.hparams["num_attention_heads"] = self.hparams.get("num_attention_heads", 32) + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + try: + self._set_vocab_llama_hf() + except (FileNotFoundError, TypeError): + # Llama 3 + self._set_vocab_gpt2() + + # Apply to CodeLlama only (and ignore for Llama 3 with a vocab size of 128256) + if self.hparams.get("vocab_size", 32000) == 32016: + special_vocab = gguf.SpecialVocab( + self.dir_model, load_merges=False, + special_token_types = ['prefix', 'suffix', 'middle', 'eot'] + ) + special_vocab._set_special_token("prefix", 32007) + special_vocab._set_special_token("suffix", 32008) + special_vocab._set_special_token("middle", 32009) + special_vocab._set_special_token("eot", 32010) + special_vocab.add_to_gguf(self.gguf_writer) + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + if "add_prefix_space" in tokenizer_config_json: + self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"]) + + # Apply to granite small models only + if self.hparams.get("vocab_size", 32000) == 49152: + self.gguf_writer.add_add_bos_token(False) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + if (rope_dim := hparams.get("head_dim")) is None: + rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(rope_dim) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + @staticmethod + def permute(weights: Tensor, n_head: int, n_head_kv: int | None): + if n_head_kv is not None and n_head != n_head_kv: + n_head = n_head_kv + return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape)) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + is_vision_tensor = "vision_tower" in name \ + or "vision_model" in name \ + or "model.connector" in name \ + or "multi_modal_projector" in name + + if is_vision_tensor: + return [] # skip vision tensors + elif self.hf_arch == "LlamaModel": + name = "model." + name + elif name.startswith("model.text_model"): + name = name.replace("text_model.", "") # for SmolVLM + elif name.startswith("language_model."): + name = name.replace("language_model.", "") # for the rest + + if self.undo_permute: + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + + # process the experts separately + if name.find("block_sparse_moe.experts") != -1: + n_experts = self.hparams["num_local_experts"] + + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for wid in ["w1", "w2", "w3"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + if rope_scaling := self.find_hparam(["rope_scaling"], optional=True): + if rope_scaling.get("rope_type", '').lower() == "llama3": + base = self.hparams.get("rope_theta", 10000.0) + if (dim := self.hparams.get("head_dim")) is None: + dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + + factor = rope_scaling.get("factor", 8.0) + low_freq_factor = rope_scaling.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling.get("high_freq_factor", 4.0) + old_context_len = self.hparams.get("original_max_position_embeddings", 8192) + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + # assert low_freq_wavelen != high_freq_wavelen # Errors for Llama4 + + rope_factors = [] + for freq in freqs: + wavelen = 2 * math.pi / freq + if wavelen < high_freq_wavelen: + rope_factors.append(1) + elif wavelen > low_freq_wavelen: + rope_factors.append(factor) + else: + smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) + rope_factors.append(1 / ((1 - smooth) / factor + smooth)) + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32)) + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("ArceeForCausalLM") +class ArceeModel(LlamaModel): + model_arch = gguf.MODEL_ARCH.ARCEE + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self._try_set_pooling_type() + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + + +@ModelBase.register( + "LlavaForConditionalGeneration", # pixtral + "Mistral3ForConditionalGeneration", # mistral small 3.1 +) +class LlavaVisionModel(MmprojModel): + img_break_tok_id = -1 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.hparams["model_type"] == "pixtral": + # layer_norm_eps is not in config.json, it is hard-coded in modeling_pixtral.py + self.hparams["layer_norm_eps"] = self.hparams.get("layer_norm_eps", 1e-5) + self.img_break_tok_id = self.get_token_id("[IMG_BREAK]") + logger.info(f"Image break token id: {self.img_break_tok_id}") + else: + raise ValueError(f"Unsupported model type: {self.hparams['model_type']}") + + def get_token_id(self, token: str) -> int: + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + added_tokens_decoder = json.load(f)['added_tokens_decoder'] + for id_, token_data in added_tokens_decoder.items(): + if token_data["content"] == token: + return int(id_) + raise ValueError(f"Token '{token}' not found in tokenizer config.") + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + if hparams["model_type"] == "pixtral": + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.PIXTRAL) + self.gguf_writer.add_vision_attention_layernorm_eps(hparams["layer_norm_eps"]) + + # hidden_act + if hparams["hidden_act"] == "silu": + self.gguf_writer.add_vision_use_silu(True) + elif hparams["hidden_act"] == "gelu": + self.gguf_writer.add_vision_use_gelu(True) + else: + raise ValueError(f"Unsupported hidden_act: {hparams['hidden_act']}") + + # spatial_merge_size + if "spatial_merge_size" in self.global_config: + self.gguf_writer.add_vision_spatial_merge_size(self.global_config["spatial_merge_size"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + n_head = self.hparams["num_attention_heads"] + n_kv_head = n_head + + if name.startswith("multi_modal_projector.") or name.startswith("vision_tower."): + # process vision tensors + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + return [(self.map_tensor_name(name), data_torch)] + + if self.img_break_tok_id > 0 and "embed_tokens.weight" in name: + logger.info(f"Extracting [IMG_BREAK] token embedding from {name}") + # for pixtral model, we need to extract the [IMG_BREAK] token embedding + img_break_embd = data_torch[self.img_break_tok_id] + name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK] + return [(self.map_tensor_name(name), img_break_embd)] + + return [] # skip other tensors + + +@ModelBase.register("Idefics3ForConditionalGeneration", "SmolVLMForConditionalGeneration") +class SmolVLMModel(MmprojModel): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + if self.hparams["model_type"] == "smolvlm_vision": + # fix for SmolVLM2, missing some keys in config.json + # default values are taken from transformers code + self.hparams["hidden_size"] = self.hparams.get("hidden_size", 1152) + self.hparams["num_attention_heads"] = self.hparams.get("num_attention_heads", 16) + self.hparams["intermediate_size"] = self.hparams.get("intermediate_size", 3072) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.IDEFICS3) + self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams.get("layer_norm_eps", 1e-5)) + self.gguf_writer.add_vision_projector_scale_factor(self.global_config.get("scale_factor", 2)) + self.gguf_writer.add_vision_use_gelu(True) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, new_name, n_dims # unused + if ".embeddings." in name: + return gguf.GGMLQuantizationType.F32 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + is_vision_tensor = "vision_tower" in name or "vision_model" in name or "model.connector" in name + + if is_vision_tensor: + return [(self.map_tensor_name(name), data_torch)] + + return [] # skip other tensors + + +@ModelBase.register("Llama4ForConditionalGeneration") +class Llama4Model(LlamaModel): + model_arch = gguf.MODEL_ARCH.LLAMA4 + undo_permute = False + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # IMPORTANT: the normal "intermediate_size" is renamed to "intermediate_size_mlp", we need to undo this + self.hparams["intermediate_size_moe"] = self.hparams["intermediate_size"] + self.hparams["intermediate_size"] = self.hparams["intermediate_size_mlp"] + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_interleave_moe_layer_step(self.hparams["interleave_moe_layer_step"]) + self.gguf_writer.add_expert_feed_forward_length(self.hparams["intermediate_size_moe"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + if name.startswith("language_model."): + name = name.replace("language_model.", "") + + # split the gate_up into gate and up + if "gate_up_proj" in name: + name_up = name.replace("gate_up_proj", "up_proj.weight") + name_gate = name.replace("gate_up_proj", "gate_proj.weight") + dim_half = data_torch.shape[-1] // 2 + gate_proj_weight, up_proj_weight = data_torch.transpose(-1, -2).split(dim_half, dim=-2) + return [ + (self.map_tensor_name(name_gate), gate_proj_weight), + (self.map_tensor_name(name_up), up_proj_weight) + ] + + if name.endswith("down_proj"): + name += ".weight" + data_torch = data_torch.transpose(-1, -2) + + if "multi_modal_projector" in name or "vision_model" in name: + return [] + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("Llama4ForConditionalGeneration") +class Llama4VisionModel(MmprojModel): + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.LLAMA4) + self.gguf_writer.add_vision_attention_layernorm_eps(self.hparams["norm_eps"]) + self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / self.hparams["pixel_shuffle_ratio"])) + assert self.hparams["hidden_act"] == "gelu" + self.gguf_writer.add_vision_use_gelu(True) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + if "multi_modal_projector" in name or "vision_model" in name: + # process vision tensors + if "positional_embedding_vlm" in name and ".weight" not in name: + name += ".weight" + if "multi_modal_projector.linear_1" in name: + # despite the name with number postfix, this is a single fully connected layer + return [(gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_MMPROJ_FC] + '.weight', data_torch)] + return [(self.map_tensor_name(name), data_torch)] + return [] + + +@ModelBase.register("Mistral3ForConditionalGeneration") +class Mistral3Model(LlamaModel): + model_arch = gguf.MODEL_ARCH.LLAMA + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + name = name.replace("language_model.", "") + if "multi_modal_projector" in name or "vision_tower" in name: + return [] + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("DeciLMForCausalLM") +class DeciModel(TextModel): + model_arch = gguf.MODEL_ARCH.DECI + + @staticmethod + def _ffn_mult_to_intermediate_size(ffn_mult: float, n_embd: int) -> int: + # DeciLM-specific code + intermediate_size = int(2 * ffn_mult * n_embd / 3) + return DeciModel._find_multiple(intermediate_size, 256) + + @staticmethod + def _find_multiple(n: int, k: int) -> int: + # DeciLM-specific code + if n % k == 0: + return n + return n + k - (n % k) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + if "block_configs" in self.hparams: # Llama-3_1-Nemotron-51B + _block_configs: list[dict[str,Any]] = self.hparams["block_configs"] + assert self.block_count == len(_block_configs) + self._num_kv_heads = list() + self._num_heads = list() + _ffn_multipliers = list() + # ***linear attention layer*** + # if n_heads_in_group is None and replace_with_linear is True + # then _num_kv_heads[il] is 0 and _num_heads[il] is num_attention_heads + # ***attention-free layer*** + # if n_heads_in_group is None and replace_with_linear is False + # then _num_kv_heads[il] is 0 and _num_heads[il] is 0 + # ***normal attention-layer*** + # if n_heads_in_group is not None, then + # _num_kv_heads[il] is num_attention_head // n_heads_in_group and + # _num_heads[il] is num_attention_head + # ***dummy layer*** for nemotron 253B + # if n_heads_in_group is None and ffn_mult is None + # then _num_kv_heads[il] is 0 and _num_heads[il] is 0 and _ffn_dims is 0 + for il in range(len(_block_configs)): + if _block_configs[il]["attention"]["n_heads_in_group"] is None: + if _block_configs[il]["attention"]["replace_with_linear"] is True: + self._num_kv_heads.append(0) + self._num_heads.append(self.hparams["num_attention_heads"]) + else: + self._num_kv_heads.append(0) + self._num_heads.append(0) + else: + self._num_kv_heads.append(self.hparams["num_attention_heads"] // _block_configs[il]["attention"]["n_heads_in_group"]) + self._num_heads.append(self.hparams["num_attention_heads"]) + if _block_configs[il]["ffn"]["ffn_mult"] is None: # dummy layer + _ffn_multipliers.append(0.0) + else: + _ffn_multipliers.append(_block_configs[il]["ffn"]["ffn_mult"]) + assert self.block_count == len(self._num_kv_heads) + assert self.block_count == len(self._num_heads) + assert self.block_count == len(_ffn_multipliers) + assert isinstance(self._num_kv_heads, list) and isinstance(self._num_kv_heads[0], int) + assert isinstance(self._num_heads, list) and isinstance(self._num_heads[0], int) + assert isinstance(_ffn_multipliers, list) and isinstance(_ffn_multipliers[0], float) + self._ffn_dims: list[int] = [ + DeciModel._ffn_mult_to_intermediate_size(multiplier, self.hparams["hidden_size"]) + for multiplier in _ffn_multipliers + ] + + def set_vocab(self): + # Please change tokenizer_config.json of Llama-3_1-Nemotron-51B's + # eos_token from '|eot_id|' to '|end_of_text|' + if self.hparams.get("vocab_size", 128256) == 128256: + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab.add_to_gguf(self.gguf_writer) + else: + # DeciLM-7B + self._set_vocab_llama_hf() + + def set_gguf_parameters(self): + if "block_configs" in self.hparams: # Llama-3_1-Nemotron-51B + assert self.block_count == len(self._num_kv_heads) + assert self.block_count == len(self._num_heads) + assert self.block_count == len(self._ffn_dims) + if (rope_theta := self.hparams.get("rope_theta")) is not None: + self.gguf_writer.add_rope_freq_base(rope_theta) + self.gguf_writer.add_head_count_kv(self._num_kv_heads) + self.gguf_writer.add_head_count(self._num_heads) + self.gguf_writer.add_feed_forward_length(self._ffn_dims) + self.gguf_writer.add_block_count(self.block_count) + self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_key_length(self.hparams["hidden_size"] // self.hparams["num_attention_heads"]) + self.gguf_writer.add_value_length(self.hparams["hidden_size"] // self.hparams["num_attention_heads"]) + self.gguf_writer.add_file_type(self.ftype) + else: # DeciLM-7B + super().set_gguf_parameters() + if "num_key_value_heads_per_layer" in self.hparams: # DeciLM-7B + self._num_kv_heads: list[int] = self.hparams["num_key_value_heads_per_layer"] + assert self.block_count == len(self._num_kv_heads) + self.gguf_writer.add_head_count_kv(self._num_kv_heads) + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + if (rope_dim := hparams.get("head_dim")) is None: + rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(rope_dim) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + @staticmethod + def permute(weights: Tensor, n_head: int, n_head_kv: int | None): + if n_head_kv is not None and n_head != n_head_kv: + n_head = n_head_kv + return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + if bid is not None: + if "num_key_value_heads_per_layer" in self.hparams: + n_kv_head = self.hparams["num_key_value_heads_per_layer"][bid] + elif "block_configs" in self.hparams: + n_kv_head = self._num_kv_heads[bid] + n_head = self._num_heads[bid] + else: + n_kv_head = self.hparams.get("num_key_value_heads") + else: + n_kv_head = self.hparams.get("num_key_value_heads") + + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = DeciModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = DeciModel.permute(data_torch, n_head, n_kv_head) + return [(self.map_tensor_name(name), data_torch)] + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + if rope_scaling := self.find_hparam(["rope_scaling"], optional=True): + if rope_scaling.get("rope_type", '').lower() == "llama3": + base = self.hparams.get("rope_theta", 10000.0) + if (dim := self.hparams.get("head_dim")) is None: + dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + + factor = rope_scaling.get("factor", 8.0) + low_freq_factor = rope_scaling.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling.get("high_freq_factor", 4.0) + old_context_len = self.hparams.get("original_max_position_embeddings", 8192) + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + assert low_freq_wavelen != high_freq_wavelen + + rope_factors = [] + for freq in freqs: + wavelen = 2 * math.pi / freq + if wavelen < high_freq_wavelen: + rope_factors.append(1) + elif wavelen > low_freq_wavelen: + rope_factors.append(factor) + else: + smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) + rope_factors.append(1 / ((1 - smooth) / factor + smooth)) + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32)) + + def prepare_tensors(self): + super().prepare_tensors() + + +@ModelBase.register("BitnetForCausalLM") +class BitnetModel(TextModel): + model_arch = gguf.MODEL_ARCH.BITNET + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(1.0) + + def weight_quant(self, weight: Tensor) -> Tensor: + dtype = weight.dtype + weight = weight.float() + scale = weight.abs().mean().clamp(min=1e-5) + iscale = 1 / scale + # TODO: multiply by the scale directly instead of inverting it twice + # (this is also unnecessarily doubly inverted upstream) + # ref: https://huggingface.co/1bitLLM/bitnet_b1_58-3B/blob/af89e318d78a70802061246bf037199d2fb97020/utils_quant.py#L10 + result = (weight * iscale).round().clamp(-1, 1) / iscale + return result.type(dtype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + new_name = self.map_tensor_name(name) + + if any(self.match_model_tensor_name(new_name, key, bid) for key in [ + gguf.MODEL_TENSOR.ATTN_Q, + gguf.MODEL_TENSOR.ATTN_K, + gguf.MODEL_TENSOR.ATTN_V, + gguf.MODEL_TENSOR.ATTN_OUT, + gguf.MODEL_TENSOR.FFN_UP, + gguf.MODEL_TENSOR.FFN_DOWN, + gguf.MODEL_TENSOR.FFN_GATE, + ]): + # transform weight into 1/0/-1 (in fp32) + data_torch = self.weight_quant(data_torch) + + yield (new_name, data_torch) + + +@ModelBase.register("GrokForCausalLM") +class GrokModel(TextModel): + model_arch = gguf.MODEL_ARCH.GROK + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # process the experts separately + if name.find(".moe.") != -1: + n_experts = self.hparams["num_local_experts"] + + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for wid in ["linear", "linear_1", "linear_v"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"transformer.decoder_layer.{bid}.moe.{xid}.{wid}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"transformer.decoder_layer.{bid}.moe.{wid}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("DbrxForCausalLM") +class DbrxModel(TextModel): + model_arch = gguf.MODEL_ARCH.DBRX + + def set_gguf_parameters(self): + ffn_config = self.hparams["ffn_config"] + attn_config = self.hparams["attn_config"] + self.gguf_writer.add_block_count(self.hparams["n_layers"]) + + self.gguf_writer.add_context_length(self.hparams["max_seq_len"]) + self.gguf_writer.add_embedding_length(self.hparams["d_model"]) + self.gguf_writer.add_feed_forward_length(ffn_config["ffn_hidden_size"]) + + self.gguf_writer.add_head_count(self.hparams["n_heads"]) + self.gguf_writer.add_head_count_kv(attn_config["kv_n_heads"]) + + self.gguf_writer.add_rope_freq_base(attn_config["rope_theta"]) + + self.gguf_writer.add_clamp_kqv(attn_config["clip_qkv"]) + + self.gguf_writer.add_expert_count(ffn_config["moe_num_experts"]) + self.gguf_writer.add_expert_used_count(ffn_config["moe_top_k"]) + + self.gguf_writer.add_layer_norm_eps(1e-5) + + self.gguf_writer.add_file_type(self.ftype) + logger.info(f"gguf: file type = {self.ftype}") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + n_expert = self.hparams["ffn_config"]["moe_num_experts"] + n_ff = self.hparams["ffn_config"]["ffn_hidden_size"] + n_embd = self.hparams["d_model"] + + # Specific behavior for experts tensors: suffix .weight, view as 3D and transpose + # original implementation expects (n_expert, n_ff, n_embd) for all experts weights + # But llama.cpp moe graph works differently + # AND the dimensions in ggml are typically in the reverse order of the pytorch dimensions + # so (n_expert, n_ff, n_embd) in pytorch is {n_embd, n_ff, n_expert} in ggml_tensor + exp_tensor_names = {"ffn.experts.mlp.w1": None, # LLM_TENSOR_FFN_GATE_EXPS ggml_tensor->ne{n_embd, n_ff, n_expert} + "ffn.experts.mlp.w2": (0, 2, 1), # LLM_TENSOR_FFN_DOWN_EXPS ggml_tensor->ne{n_ff, n_embd, n_expert} + "ffn.experts.mlp.v1": None} # LLM_TENSOR_FFN_UP_EXPS ggml_tensor->ne{n_embd, n_ff, n_expert} + experts = False + + for exp_tensor_name in exp_tensor_names.keys(): + if name.find(exp_tensor_name) != -1 and name.find(".weight") == -1: + experts = True + data_torch = data_torch.view(n_expert, n_ff, n_embd) + if (permute_tensor := exp_tensor_names[exp_tensor_name]) is not None: + data_torch = data_torch.permute(*permute_tensor) + break + + # map tensor names + # In MoE models the ffn tensors are typically most of the model weights, + # and need to be quantizable. Quantize expects tensor names to be suffixed by .weight. + # Every other model has the weight names ending in .weight, + # let's assume that is the convention which is not the case for dbrx: + # https://huggingface.co/databricks/dbrx-instruct/blob/main/model.safetensors.index.json#L15 + new_name = self.map_tensor_name(name if not experts else name + ".weight", try_suffixes=(".weight",)) + + return [(new_name, data_torch)] + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + del name, new_name, bid # unused + + return n_dims > 1 + + +@ModelBase.register("MiniCPMForCausalLM") +class MiniCPMModel(TextModel): + model_arch = gguf.MODEL_ARCH.MINICPM + + def set_gguf_parameters(self): + super().set_gguf_parameters() + embedding_scale = float(self.hparams["scale_emb"]) + self.gguf_writer.add_embedding_scale(embedding_scale) + logger.info(f"gguf: (minicpm) embedding_scale = {embedding_scale}") + residual_scale = self.hparams["scale_depth"] / self.hparams["num_hidden_layers"] ** 0.5 + self.gguf_writer.add_residual_scale(residual_scale) + logger.info(f"gguf: (minicpm) residual_scale = {residual_scale}") + logit_scale = self.hparams["hidden_size"] / self.hparams["dim_model_base"] + self.gguf_writer.add_logit_scale(logit_scale) + logger.info(f"gguf: (minicpm) logit_scale = {logit_scale}") + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "longrope": + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LONGROPE) + logger.info(f"gguf: (minicpm) rope_scaling_type = {gguf.RopeScalingType.LONGROPE}") + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + rope_dims = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + + rope_scaling = self.find_hparam(['rope_scaling'], True) + if rope_scaling is not None: + long_factors = rope_scaling.get('long_factor', None) + short_factors = rope_scaling.get('short_factor', None) + + if long_factors is None or short_factors is None: + raise KeyError('Missing the required key rope_scaling.long_factor or rope_scaling_short_factor') + + if len(long_factors) != len(short_factors) or len(long_factors) != rope_dims / 2: + raise ValueError(f'The length of rope long and short factors must be {rope_dims / 2}') + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_LONG), torch.tensor(long_factors, dtype=torch.float32)) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_SHORT), torch.tensor(short_factors, dtype=torch.float32)) + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + + # HF models permute some of the tensors, so we need to undo that + if name.endswith(("q_proj.weight")): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight")): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("MiniCPM3ForCausalLM") +class MiniCPM3Model(TextModel): + model_arch = gguf.MODEL_ARCH.MINICPM3 + + def set_gguf_parameters(self): + hparams = self.hparams + + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_context_length(hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_block_count(self.block_count) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(hparams["num_key_value_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"]) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + if "q_lora_rank" in hparams and hparams["q_lora_rank"] is not None: + self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) + self.gguf_writer.add_key_length(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]) + self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + rope_scaling = self.find_hparam(['rope_scaling'], True) + if rope_scaling is not None: + rope_dims = self.hparams["qk_rope_head_dim"] + + long_factors = rope_scaling.get('long_factor', None) + short_factors = rope_scaling.get('short_factor', None) + + if long_factors is None or short_factors is None: + raise KeyError('Missing the required key rope_scaling.long_factor or rope_scaling_short_factor') + + if len(long_factors) != len(short_factors) or len(long_factors) != rope_dims / 2: + raise ValueError(f'The length of rope long and short factors must be {rope_dims / 2}') + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_LONG), torch.tensor(long_factors, dtype=torch.float32)) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_SHORT), torch.tensor(short_factors, dtype=torch.float32)) + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def _reverse_hf_permute(self, weights: Tensor, n_head: int, n_kv_head: int | None = None) -> Tensor: + if n_kv_head is not None and n_head != n_kv_head: + n_head //= n_kv_head + + return ( + weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape) + ) + + +@ModelBase.register("QWenLMHeadModel") +class QwenModel(TextModel): + model_arch = gguf.MODEL_ARCH.QWEN + + @staticmethod + def token_bytes_to_string(b): + from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode + byte_encoder = bytes_to_unicode() + return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')]) + + @staticmethod + def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]: + parts = [bytes([b]) for b in token] + while True: + min_idx = None + min_rank = None + for i, pair in enumerate(zip(parts[:-1], parts[1:])): + rank = mergeable_ranks.get(pair[0] + pair[1]) + if rank is not None and (min_rank is None or rank < min_rank): + min_idx = i + min_rank = rank + if min_rank is None or (max_rank is not None and min_rank >= max_rank): + break + assert min_idx is not None + parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:] + return parts + + def set_vocab(self): + self._set_vocab_qwen() + + def set_gguf_parameters(self): + self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"]) + self.gguf_writer.add_block_count(self.hparams["num_hidden_layers"]) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_rope_freq_base(self.hparams["rotary_emb_base"]) + self.gguf_writer.add_rope_dimension_count(self.hparams["hidden_size"] // self.hparams["num_attention_heads"]) + self.gguf_writer.add_head_count(self.hparams["num_attention_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + +@ModelBase.register("Qwen2Model", "Qwen2ForCausalLM", "Qwen2AudioForConditionalGeneration") +class Qwen2Model(TextModel): + model_arch = gguf.MODEL_ARCH.QWEN2 + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self._try_set_pooling_type() + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if self.hf_arch == "Qwen2Model": + name = f"model.{name}" # map to Qwen2ForCausalLM tensors + if "language_model." in name: + name = name.replace("language_model.", "") # for InternVL + if name.startswith("mlp") or name.startswith("multi_modal_projector") \ + or name.startswith("vision_model") or name.startswith("audio_tower"): + # skip vision and audio tensors + return [] + yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("Ernie4_5_ForCausalLM") +class Ernie4_5Model(TextModel): + model_arch = gguf.MODEL_ARCH.ERNIE4_5 + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + num_heads = self.hparams["num_attention_heads"] + num_kv_heads = self.hparams["num_key_value_heads"] + head_dim = self.hparams["head_dim"] + + if "ernie." in name: + name = name.replace("ernie.", "model.") + # split the qkv weights + # qkv_proj shape: [(num_heads + 2 * num_kv_heads) * head_dim, hidden_size] + if "qkv_proj" in name: + name_q = name.replace("qkv_proj.weight", "q_proj.weight") + name_k = name.replace("qkv_proj.weight", "k_proj.weight") + name_v = name.replace("qkv_proj.weight", "v_proj.weight") + total_q_dim = num_heads * head_dim + total_k_dim = num_kv_heads * head_dim + total_v_dim = num_kv_heads * head_dim + q_proj_weight, k_proj_weight, v_proj_weight = data_torch.split([total_q_dim, total_k_dim, total_v_dim], dim=0) + return [ + (self.map_tensor_name(name_q), q_proj_weight), + (self.map_tensor_name(name_k), k_proj_weight), + (self.map_tensor_name(name_v), v_proj_weight) + ] + # split the up_gate_proj into gate and up + # up_gate_proj shape: [2 * intermediate_size, hidden_size] + if "up_gate_proj" in name: + name_up = name.replace("up_gate_proj.weight", "up_proj.weight") + name_gate = name.replace("up_gate_proj.weight", "gate_proj.weight") + dim_half = data_torch.shape[0] // 2 + gate_proj_weight, up_proj_weight = data_torch.split(dim_half, dim=0) + return [ + (self.map_tensor_name(name_gate), gate_proj_weight), + (self.map_tensor_name(name_up), up_proj_weight) + ] + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register( + "Qwen2VLModel", + "Qwen2VLForConditionalGeneration", + "Qwen2_5_VLForConditionalGeneration", + "Qwen2_5OmniModel", +) +class Qwen2VLModel(TextModel): + model_arch = gguf.MODEL_ARCH.QWEN2VL + + def set_gguf_parameters(self): + super().set_gguf_parameters() + mrope_section = self.hparams["rope_scaling"]["mrope_section"] + mrope_section += [0] * max(0, 4 - len(mrope_section)) + self.gguf_writer.add_rope_dimension_sections(mrope_section) + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_gpt2() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + if name.startswith("thinker."): + name = name.replace("thinker.", "") + if name.startswith("visual") or name.startswith("audio") or \ + name.startswith("talker") or name.startswith("token2wav"): + # skip multimodal tensors + return [] + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("Qwen2VLModel", "Qwen2VLForConditionalGeneration", "Qwen2_5_VLForConditionalGeneration") +class Qwen2VLVisionModel(MmprojModel): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams_vision is not None + self.hparams_vision["image_size"] = self.hparams_vision.get("image_size", 560) + # rename config.json values + self.hparams_vision["num_attention_heads"] = self.hparams_vision.get("num_heads") + self.hparams_vision["num_hidden_layers"] = self.hparams_vision.get("depth") + if "embed_dim" in self.hparams_vision: # qwen2vl + self.hparams_vision["intermediate_size"] = self.hparams_vision.get("hidden_size") + self.hparams_vision["hidden_size"] = self.hparams_vision.get("embed_dim") + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + hparams = self.hparams_vision + model_type = self.global_config['model_type'] + if model_type == 'qwen2_vl': + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.QWEN2VL) + elif model_type == 'qwen2_5_vl' or model_type == 'qwen2_5_omni': + if model_type == 'qwen2_5_omni': + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.QWEN25O) + else: + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.QWEN25VL) + self.gguf_writer.add_vision_use_silu(True) + # find n_wa_pattern (window attention pattern) + fullatt_block_indexes = hparams.get("fullatt_block_indexes") + assert fullatt_block_indexes is not None, "fullatt_block_indexes is required for qwen2_5_vl" + n_wa_pattern = fullatt_block_indexes[0] + 1 + # validate n_wa_pattern + for i in range(1, len(fullatt_block_indexes)): + if fullatt_block_indexes[i] - fullatt_block_indexes[i - 1] != n_wa_pattern: + raise ValueError(f"Invalid fullatt_block_indexes: {fullatt_block_indexes}") + self.gguf_writer.add_vision_n_wa_pattern(n_wa_pattern) + else: + raise ValueError(f"Unknown QwenVL model type: {self.global_config['model_type']}") + # default values below are taken from HF tranformers code + self.gguf_writer.add_vision_attention_layernorm_eps(self.global_config.get("rms_norm_eps", 1e-6)) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, name, n_dims # unused + if ".patch_embd." in new_name: + return gguf.GGMLQuantizationType.F16 + if ".position_embd." in new_name: + return gguf.GGMLQuantizationType.F32 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + if name.startswith("visual."): + # process visual tensors + # split QKV tensors if needed + if ".qkv." in name: + if data_torch.ndim == 2: # weight + c3, _ = data_torch.shape + else: # bias + c3 = data_torch.shape[0] + assert c3 % 3 == 0 + c = c3 // 3 + wq = data_torch[:c] + wk = data_torch[c: c * 2] + wv = data_torch[c * 2:] + return [ + (self.map_tensor_name(name.replace("qkv", "q")), wq), + (self.map_tensor_name(name.replace("qkv", "k")), wk), + (self.map_tensor_name(name.replace("qkv", "v")), wv), + ] + elif 'patch_embed.proj.weight' in name: + # split Conv3D into Conv2Ds + c1, c2, kt, kh, kw = data_torch.shape + del c1, c2, kh, kw # unused + assert kt == 2, "Current implmentation only support temporal_patch_size of 2" + return [ + (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight" , data_torch[:, :, 0, ...]), + (gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + ".weight.1", data_torch[:, :, 1, ...]), + ] + else: + return [(self.map_tensor_name(name), data_torch)] + return [] # skip other tensors + + +@ModelBase.register("Qwen2_5OmniModel") +class Qwen25OmniModel(Qwen2VLVisionModel): + has_vision_encoder = True + has_audio_encoder = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams_audio is not None + self.hparams_audio["hidden_size"] = self.hparams_audio["d_model"] + self.hparams_audio["intermediate_size"] = self.hparams_audio["encoder_ffn_dim"] + self.hparams_audio["num_attention_heads"] = self.hparams_audio["encoder_attention_heads"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_audio is not None + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("layer_norm_eps", 1e-5)) + + def get_vision_config(self) -> dict[str, Any] | None: + return self.global_config["thinker_config"].get("vision_config") + + def get_audio_config(self) -> dict[str, Any] | None: + return self.global_config["thinker_config"].get("audio_config") + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + # SinusoidsPositionEmbedding + assert self.hparams_audio is not None + max_timescale = 10000 + length = 1500 + channels = self.hparams_audio["hidden_size"] + log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1) + inv_timescales = torch.exp(-log_timescale_increment * torch.arange(channels // 2).float()) + scaled_time = torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :] + pos_embd = torch.cat([torch.sin(scaled_time), torch.cos(scaled_time)], dim=1).to(dtype=torch.float32) + yield ("audio_tower.embed_positions.weight", pos_embd) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, new_name, n_dims # unused + if ".conv" in name and ".weight" in name: + return gguf.GGMLQuantizationType.F16 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("thinker."): + name = name.replace("thinker.", "") + + if name.startswith("audio_tower"): + # process audio tensors + if "conv1.bias" in name or "conv2.bias" in name: + # transpose conv1 and conv2 bias + data_torch = data_torch.unsqueeze(-1) + if "audio_bos_eos_token" in name: + # this tensor is left unused in transformers code + # https://github.com/huggingface/transformers/blob/6e3063422c4b1c014aa60c32b9254fd2902f0f28/src/transformers/models/qwen2_5_omni/modular_qwen2_5_omni.py#L1809 + return [] + return [(self.map_tensor_name(name), data_torch)] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("InternVisionModel") +class InternVisionModel(MmprojModel): + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.INTERNVL) + self.gguf_writer.add_vision_attention_layernorm_eps(hparams["layer_norm_eps"]) + # hidden_act + if hparams["hidden_act"] == "silu": + self.gguf_writer.add_vision_use_silu(True) + elif hparams["hidden_act"] == "gelu": + self.gguf_writer.add_vision_use_gelu(True) + else: + raise ValueError(f"Unsupported hidden_act: {hparams['hidden_act']}") + # downsample_ratio + downsample_ratio = self.global_config.get("downsample_ratio") + assert downsample_ratio is not None + self.gguf_writer.add_vision_projector_scale_factor(int(1.0 / downsample_ratio)) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, name, n_dims # unused + if ".patch_embd." in new_name: + return gguf.GGMLQuantizationType.F16 + if ".position_embd." in new_name: + return gguf.GGMLQuantizationType.F32 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + if name.startswith("vision_model") or name.startswith("mlp"): + # process visual tensors + # correct name + if name.startswith("vision_model"): + name = "vision_tower." + name + if (".ls" in name or "position_embedding" in name) and not name.endswith(".weight"): + name += ".weight" + # split QKV tensors if needed + if ".qkv." in name: + if data_torch.ndim == 2: # weight + c3, _ = data_torch.shape + else: # bias + c3 = data_torch.shape[0] + assert c3 % 3 == 0 + c = c3 // 3 + wq = data_torch[:c] + wk = data_torch[c: c * 2] + wv = data_torch[c * 2:] + return [ + (self.map_tensor_name(name.replace("attn.qkv", "self_attn.q_proj")), wq), + (self.map_tensor_name(name.replace("attn.qkv", "self_attn.k_proj")), wk), + (self.map_tensor_name(name.replace("attn.qkv", "self_attn.v_proj")), wv), + ] + return [(self.map_tensor_name(name), data_torch)] + return [] # skip other tensors + + +@ModelBase.register("WavTokenizerDec") +class WavTokenizerDecModel(TextModel): + model_arch = gguf.MODEL_ARCH.WAVTOKENIZER_DEC + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if \ + name.endswith("codebook.cluster_size") or \ + name.endswith("codebook.embed_avg") or \ + name.endswith("codebook.inited"): + logger.debug(f"Skipping {name!r}") + return [] + + logger.info(f"{self.map_tensor_name(name)} -> {data_torch.shape}") + + return [(self.map_tensor_name(name), data_torch)] + + def set_vocab(self): + self._set_vocab_none() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_vocab_size (self.hparams["vocab_size"]) + self.gguf_writer.add_features_length (self.hparams["n_embd_features"]) + self.gguf_writer.add_feed_forward_length(self.hparams["n_ff"]) + self.gguf_writer.add_group_norm_eps (self.hparams["group_norm_epsilon"]) + self.gguf_writer.add_group_norm_groups (self.hparams["group_norm_groups"]) + + self.gguf_writer.add_posnet_embedding_length(self.hparams["posnet"]["n_embd"]) + self.gguf_writer.add_posnet_block_count (self.hparams["posnet"]["n_layer"]) + + self.gguf_writer.add_convnext_embedding_length(self.hparams["convnext"]["n_embd"]) + self.gguf_writer.add_convnext_block_count (self.hparams["convnext"]["n_layer"]) + + self.gguf_writer.add_causal_attention(False) + + +@ModelBase.register("Qwen2MoeForCausalLM") +class Qwen2MoeModel(TextModel): + model_arch = gguf.MODEL_ARCH.QWEN2MOE + + def set_gguf_parameters(self): + super().set_gguf_parameters() + if (n_experts := self.hparams.get("num_experts")) is not None: + self.gguf_writer.add_expert_count(n_experts) + if (moe_intermediate_size := self.hparams.get("moe_intermediate_size")) is not None: + self.gguf_writer.add_expert_feed_forward_length(moe_intermediate_size) + logger.info(f"gguf: expert feed forward length = {moe_intermediate_size}") + if (shared_expert_intermediate_size := self.hparams.get('shared_expert_intermediate_size')) is not None: + self.gguf_writer.add_expert_shared_feed_forward_length(shared_expert_intermediate_size) + logger.info(f"gguf: expert shared feed forward length = {shared_expert_intermediate_size}") + # YaRN is not enabled by default + # To enable it, please refer to this guide: https://huggingface.co/Qwen/Qwen3-30B-A3B#processing-long-texts + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # process the experts separately + if name.find("experts") != -1: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("Qwen3ForCausalLM") +class Qwen3Model(Qwen2Model): + model_arch = gguf.MODEL_ARCH.QWEN3 + + +@ModelBase.register("Qwen3MoeForCausalLM") +class Qwen3MoeModel(Qwen2MoeModel): + model_arch = gguf.MODEL_ARCH.QWEN3MOE + + +@ModelBase.register("GPT2LMHeadModel") +class GPT2Model(TextModel): + model_arch = gguf.MODEL_ARCH.GPT2 + + def set_gguf_parameters(self): + self.gguf_writer.add_block_count(self.hparams["n_layer"]) + self.gguf_writer.add_context_length(self.hparams["n_ctx"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["n_embd"]) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + tensors: list[tuple[str, Tensor]] = [] + + # we don't need these + if name.endswith((".attn.bias", ".attn.masked_bias")): + return tensors + + if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight", ".c_proj.weight")): + data_torch = data_torch.transpose(1, 0) + + new_name = self.map_tensor_name(name) + + tensors.append((new_name, data_torch)) + + return tensors + + +@ModelBase.register("PhiForCausalLM") +class Phi2Model(TextModel): + model_arch = gguf.MODEL_ARCH.PHI2 + + def set_gguf_parameters(self): + block_count = self.find_hparam(["num_hidden_layers", "n_layer"]) + + rot_pct = self.find_hparam(["partial_rotary_factor"]) + n_embd = self.find_hparam(["hidden_size", "n_embd"]) + n_head = self.find_hparam(["num_attention_heads", "n_head"]) + + self.gguf_writer.add_context_length(self.find_hparam(["n_positions", "max_position_embeddings"])) + + self.gguf_writer.add_embedding_length(n_embd) + self.gguf_writer.add_feed_forward_length(4 * n_embd) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(n_head) + self.gguf_writer.add_head_count_kv(n_head) + self.gguf_writer.add_layer_norm_eps(self.find_hparam(["layer_norm_epsilon", "layer_norm_eps"])) + self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head) + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_add_bos_token(False) + + +@ModelBase.register("Phi3ForCausalLM") +class Phi3MiniModel(TextModel): + model_arch = gguf.MODEL_ARCH.PHI3 + + def set_vocab(self): + # Phi-4 model uses GPT2Tokenizer + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + tokenizer_class = tokenizer_config_json['tokenizer_class'] + if tokenizer_class == 'GPT2Tokenizer': + return self._set_vocab_gpt2() + + from sentencepiece import SentencePieceProcessor + + tokenizer_path = self.dir_model / 'tokenizer.model' + + if not tokenizer_path.is_file(): + raise ValueError(f'Error: Missing {tokenizer_path}') + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + for token_id in range(tokenizer.vocab_size()): + + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + added_tokens_file = self.dir_model / 'added_tokens.json' + if added_tokens_file.is_file(): + with open(added_tokens_file, "r", encoding="utf-8") as f: + added_tokens_json = json.load(f) + + for key in added_tokens_json: + token_id = added_tokens_json[key] + if token_id >= vocab_size: + logger.debug(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + + tokens[token_id] = key.encode("utf-8") + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + added_tokens_decoder = tokenizer_config_json.get("added_tokens_decoder", {}) + for token_id, foken_data in added_tokens_decoder.items(): + token_id = int(token_id) + token = foken_data["content"].encode("utf-8") + if toktypes[token_id] != SentencePieceTokenTypes.UNUSED: + if tokens[token_id] != token: + logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}') + tokens[token_id] = token + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + if foken_data.get("special"): + toktypes[token_id] = SentencePieceTokenTypes.CONTROL + + tokenizer_file = self.dir_model / 'tokenizer.json' + if tokenizer_file.is_file(): + with open(tokenizer_file, "r", encoding="utf-8") as f: + tokenizer_json = json.load(f) + added_tokens = tokenizer_json.get("added_tokens", []) + for foken_data in added_tokens: + token_id = int(foken_data["id"]) + token = foken_data["content"].encode("utf-8") + if toktypes[token_id] != SentencePieceTokenTypes.UNUSED: + if tokens[token_id] != token: + logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}') + tokens[token_id] = token + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + if foken_data.get("special"): + toktypes[token_id] = SentencePieceTokenTypes.CONTROL + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + block_count = self.find_hparam(["num_hidden_layers", "n_layer"]) + + n_embd = self.find_hparam(["hidden_size", "n_embd"]) + n_head = self.find_hparam(["num_attention_heads", "n_head"]) + n_head_kv = self.find_hparam(["num_key_value_heads", "n_head_kv"]) + rms_eps = self.find_hparam(["rms_norm_eps"]) + max_pos_embds = self.find_hparam(["n_positions", "max_position_embeddings"]) + orig_max_pos_embds = self.find_hparam(["original_max_position_embeddings"]) + rot_pct = self.hparams.get("partial_rotary_factor", 1.0) + rope_dims = int(rot_pct * n_embd) // n_head + + self.gguf_writer.add_context_length(max_pos_embds) + self.gguf_writer.add_rope_scaling_orig_ctx_len(orig_max_pos_embds) + self.gguf_writer.add_embedding_length(n_embd) + self.gguf_writer.add_feed_forward_length(self.find_hparam(["intermediate_size"])) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(n_head) + self.gguf_writer.add_head_count_kv(n_head_kv) + self.gguf_writer.add_layer_norm_rms_eps(rms_eps) + self.gguf_writer.add_rope_dimension_count(rope_dims) + self.gguf_writer.add_rope_freq_base(self.find_hparam(["rope_theta"])) + self.gguf_writer.add_file_type(self.ftype) + sliding_window = self.hparams.get("sliding_window") + # use zero value of sliding_window to distinguish Phi-4 from other PHI3 models + if sliding_window is None: + sliding_window = 0 + self.gguf_writer.add_sliding_window(sliding_window) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + n_embd = self.find_hparam(["hidden_size", "n_embd"]) + n_head = self.find_hparam(["num_attention_heads", "n_head"]) + max_pos_embds = self.find_hparam(["n_positions", "max_position_embeddings"]) + orig_max_pos_embds = self.find_hparam(["original_max_position_embeddings"]) + rot_pct = self.hparams.get("partial_rotary_factor", 1.0) + rope_dims = int(rot_pct * n_embd) // n_head + + # write rope scaling for long context (128k) model + rope_scaling = self.find_hparam(['rope_scaling'], True) + if rope_scaling is None: + return + + scale = max_pos_embds / orig_max_pos_embds + + rope_scaling_type = rope_scaling.get('rope_type', rope_scaling.get('type', '')).lower() + if len(rope_scaling_type) == 0: + raise KeyError('Missing the required key rope_scaling.type') + + if rope_scaling_type == 'su' or rope_scaling_type == 'longrope': + attn_factor = math.sqrt(1 + math.log(scale) / math.log(orig_max_pos_embds)) if scale > 1.0 else 1.0 + elif rope_scaling_type == 'yarn': + attn_factor = 0.1 * math.log(scale) + 1.0 if scale > 1.0 else 1.0 + else: + raise NotImplementedError(f'The rope scaling type {rope_scaling_type} is not supported yet') + + self.gguf_writer.add_rope_scaling_attn_factors(attn_factor) + + long_factors = rope_scaling.get('long_factor', None) + short_factors = rope_scaling.get('short_factor', None) + + if long_factors is None or short_factors is None: + raise KeyError('Missing the required key rope_scaling.long_factor or rope_scaling_short_factor') + + if len(long_factors) != len(short_factors) or len(long_factors) != rope_dims / 2: + raise ValueError(f'The length of rope long and short factors must be {rope_dims / 2}. long_factors = {len(long_factors)}, short_factors = {len(short_factors)}.') + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_LONG), torch.tensor(long_factors, dtype=torch.float32)) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FACTORS_SHORT), torch.tensor(short_factors, dtype=torch.float32)) + + +@ModelBase.register("PhiMoEForCausalLM") +class PhiMoeModel(Phi3MiniModel): + model_arch = gguf.MODEL_ARCH.PHIMOE + + _experts: list[dict[str, Tensor]] | None = None + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_expert_used_count(self.hparams["num_experts_per_tok"]) + self.gguf_writer.add_expert_count(self.hparams["num_local_experts"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # process the experts separately + if name.find("block_sparse_moe.experts") != -1: + n_experts = self.hparams["num_local_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for w_name in ["w1", "w2", "w3"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.block_sparse_moe.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("PlamoForCausalLM") +class PlamoModel(TextModel): + model_arch = gguf.MODEL_ARCH.PLAMO + + def set_vocab(self): + self._set_vocab_sentencepiece() + + def set_gguf_parameters(self): + hparams = self.hparams + block_count = hparams["num_hidden_layers"] + + self.gguf_writer.add_context_length(4096) # not in config.json + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(5) # hparams["num_key_value_heads"]) is wrong + self.gguf_writer.add_layer_norm_rms_eps(hparams["rms_norm_eps"]) + self.gguf_writer.add_file_type(self.ftype) + + def shuffle_attn_q_weight(self, data_torch): + assert data_torch.size() == (5120, 5120) + data_torch = data_torch.reshape(8, 5, 128, 5120) + data_torch = torch.permute(data_torch, (1, 0, 2, 3)) + data_torch = torch.reshape(data_torch, (5120, 5120)) + return data_torch + + def shuffle_attn_output_weight(self, data_torch): + assert data_torch.size() == (5120, 5120) + data_torch = data_torch.reshape(5120, 8, 5, 128) + data_torch = torch.permute(data_torch, (0, 2, 1, 3)) + data_torch = torch.reshape(data_torch, (5120, 5120)) + return data_torch + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + new_name = self.map_tensor_name(name) + + # shuffle for broadcasting of gqa in ggml_mul_mat + if new_name.endswith("attn_q.weight"): + data_torch = self.shuffle_attn_q_weight(data_torch) + elif new_name.endswith("attn_output.weight"): + data_torch = self.shuffle_attn_output_weight(data_torch) + + return [(new_name, data_torch)] + + +@ModelBase.register("CodeShellForCausalLM") +class CodeShellModel(TextModel): + model_arch = gguf.MODEL_ARCH.CODESHELL + + def set_gguf_parameters(self): + block_count = self.hparams["n_layer"] + + self.gguf_writer.add_context_length(self.hparams["n_positions"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + self.gguf_writer.add_feed_forward_length(4 * self.hparams["n_embd"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_head_count_kv(self.hparams["num_query_groups"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_rope_freq_base(10000.0) + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(1.0) + + _has_tok_embd = False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + output_name = self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT) + tok_embd_name = self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD) + + new_name = self.map_tensor_name(name) + + # assuming token_embd.weight is seen before output.weight + if not self._has_tok_embd and new_name == self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT): + # even though the tensor file(s) does not contain the word embeddings they are still in the weight map + if self.tensor_names and "transformer.wte.weight" in self.tensor_names: + logger.debug(f"{tok_embd_name} not found before {output_name}, assuming they are tied") + self.tensor_names.remove("transformer.wte.weight") + elif new_name == tok_embd_name: + self._has_tok_embd = True + + return [(new_name, data_torch)] + + +@ModelBase.register("InternLM2ForCausalLM") +class InternLM2Model(TextModel): + model_arch = gguf.MODEL_ARCH.INTERNLM2 + + def set_vocab(self): + # (TODO): Is there a better way? + # Copy from _set_vocab_sentencepiece, The only difference is that we will treat the character + # \x00 specially and convert it into an emoji character to prevent it from being mistakenly + # recognized as an empty string in C++. + from sentencepiece import SentencePieceProcessor + from sentencepiece import sentencepiece_model_pb2 as model + + tokenizer_path = self.dir_model / 'tokenizer.model' + + tokens: list[bytes] = [] + scores: list[float] = [] + toktypes: list[int] = [] + + if not tokenizer_path.is_file(): + logger.error(f'Error: Missing {tokenizer_path}') + sys.exit(1) + + sentencepiece_model = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] + sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read()) + add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) + + for token_id in range(vocab_size): + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + if text == b"\x00": + # (TODO): fixme + # Hack here and replace the \x00 characters. + logger.warning(f"InternLM2 convert token '{text}' to '🐉'!") + text = "🐉".encode("utf-8") + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + # take care of ununsed raw token + if piece.startswith('[UNUSED'): + toktype = SentencePieceTokenTypes.UNUSED + + tokens.append(text) + scores.append(score) + toktypes.append(toktype) + + added_tokens_file = self.dir_model / 'added_tokens.json' + if added_tokens_file.is_file(): + with open(added_tokens_file, "r", encoding="utf-8") as f: + added_tokens_json = json.load(f) + + for key in added_tokens_json: + tokens.append(key.encode("utf-8")) + scores.append(-1000.0) + toktypes.append(SentencePieceTokenTypes.USER_DEFINED) + + chat_eos_token = '<|im_end|>' + chat_eos_token_id = None + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + added_tokens_decoder = tokenizer_config_json.get("added_tokens_decoder", {}) + for token_id, foken_data in added_tokens_decoder.items(): + token_id = int(token_id) + token = foken_data["content"] + if token == chat_eos_token: + chat_eos_token_id = token_id + token = token.encode("utf-8") + if toktypes[token_id] != SentencePieceTokenTypes.UNUSED: + if tokens[token_id] != token: + logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}') + tokens[token_id] = token + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + if foken_data.get("special"): + toktypes[token_id] = SentencePieceTokenTypes.CONTROL + + tokenizer_file = self.dir_model / 'tokenizer.json' + if tokenizer_file.is_file(): + with open(tokenizer_file, "r", encoding="utf-8") as f: + tokenizer_json = json.load(f) + added_tokens = tokenizer_json.get("added_tokens", []) + for foken_data in added_tokens: + token_id = int(foken_data["id"]) + token = foken_data["content"] + if token == chat_eos_token: + chat_eos_token_id = token_id + token = token.encode("utf-8") + if toktypes[token_id] != SentencePieceTokenTypes.UNUSED: + if tokens[token_id] != token: + logger.warning(f'replacing token {token_id}: {tokens[token_id].decode("utf-8")!r} -> {token.decode("utf-8")!r}') + tokens[token_id] = token + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + if foken_data.get("special"): + toktypes[token_id] = SentencePieceTokenTypes.CONTROL + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(add_prefix) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + old_eos = special_vocab.special_token_ids["eos"] + if chat_eos_token_id is not None: + # For the chat model, we replace the eos with '<|im_end|>'. + # TODO: this is a hack, should be fixed + # https://github.com/ggml-org/llama.cpp/pull/6745#issuecomment-2067687048 + special_vocab.special_token_ids["eos"] = chat_eos_token_id + logger.warning(f"Replace eos:{old_eos} with a special token:{chat_eos_token_id}" + " in chat mode so that the conversation can end normally.") + + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + self.gguf_writer.add_context_length(self.hparams["max_position_embeddings"]) + self.gguf_writer.add_block_count(self.hparams["num_hidden_layers"]) + self.gguf_writer.add_embedding_length(self.hparams["hidden_size"]) + self.gguf_writer.add_feed_forward_length(self.hparams["intermediate_size"]) + self.gguf_writer.add_rope_freq_base(self.hparams["rope_theta"]) + self.gguf_writer.add_head_count(self.hparams["num_attention_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_head_count_kv(self.hparams["num_key_value_heads"]) + self.gguf_writer.add_file_type(self.ftype) + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + num_heads = self.hparams["num_attention_heads"] + num_kv_heads = self.hparams["num_key_value_heads"] + n_embd = self.hparams["hidden_size"] + q_per_kv = num_heads // num_kv_heads + head_dim = n_embd // num_heads + num_groups = num_heads // q_per_kv + + name = name.replace("language_model.", "") # InternVL + if name.startswith("mlp") or name.startswith("vision_model"): + # skip visual tensors + return [] + + if bid is not None and f"model.layers.{bid}.attention.wqkv" in name: + qkv = data_torch + + qkv = qkv.reshape((num_groups, q_per_kv + 2, head_dim, n_embd)) + q, k, v = qkv[:, : q_per_kv], qkv[:, -2], qkv[:, -1] + + # The model weights of q and k equire additional reshape. + q = LlamaModel.permute(q.reshape((-1, q.shape[-1])), num_heads, num_heads) + k = LlamaModel.permute(k.reshape((-1, k.shape[-1])), num_heads, num_kv_heads) + v = v.reshape((-1, v.shape[-1])) + + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v), + ] + else: + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("InternLM3ForCausalLM") +class InternLM3Model(TextModel): + model_arch = gguf.MODEL_ARCH.LLAMA + + def set_vocab(self): + tokens, scores, toktypes = self._create_vocab_sentencepiece() + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + if "add_prefix_space" in tokenizer_config_json: + self.gguf_writer.add_add_space_prefix(tokenizer_config_json["add_prefix_space"]) + + if "added_tokens_decoder" in tokenizer_config_json: + for token_id, token_data in tokenizer_config_json["added_tokens_decoder"].items(): + if token_data.get("special"): + token_id = int(token_id) + token = token_data["content"] + special_vocab._set_special_token(token, token_id) + # update eos token + if token == '<|im_end|>' and "eos" in special_vocab.special_token_ids: + special_vocab.special_token_ids["eos"] = token_id + + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + if (rope_dim := hparams.get("head_dim")) is None: + rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(rope_dim) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + name = name.replace("language_model.", "") # InternVL + if name.startswith("mlp") or name.startswith("vision_model"): + # skip visual tensors + return [] + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("BertModel", "BertForMaskedLM", "CamembertModel", "BertForSequenceClassification") +class BertModel(TextModel): + model_arch = gguf.MODEL_ARCH.BERT + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.vocab_size = None + + if cls_out_labels := self.hparams.get("id2label"): + if len(cls_out_labels) == 2 and cls_out_labels[0] == "LABEL_0": + # Remove dummy labels added by AutoConfig + cls_out_labels = None + self.cls_out_labels = cls_out_labels + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_causal_attention(False) + self._try_set_pooling_type() + + if self.cls_out_labels: + self.gguf_writer.add_classifier_output_labels([v for k, v in sorted(self.cls_out_labels.items())]) + + def set_vocab(self): + tokens, toktypes, tokpre = self.get_vocab_base() + self.vocab_size = len(tokens) + + # we need this to validate the size of the token_type embeddings + # though currently we are passing all zeros to the token_type embeddings + # "Sequence A" or "Sequence B" + self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1)) + + # convert to phantom space vocab + def phantom(tok): + if tok.startswith("[") and tok.endswith("]"): + return tok + if tok.startswith("##"): + return tok[2:] + return "\u2581" + tok + tokens = list(map(phantom, tokens)) + + # add vocab to gguf + self.gguf_writer.add_tokenizer_model("bert") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + + # handle special tokens + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if name.startswith("bert."): + name = name[5:] + + if name.endswith(".gamma"): + name = name[:-6] + ".weight" + + if name.endswith(".beta"): + name = name[:-5] + ".bias" + + # we are only using BERT for embeddings so we don't need the pooling layer + if name in ("embeddings.position_ids", "pooler.dense.weight", "pooler.dense.bias"): + return [] # we don't need these + + if name.startswith("cls.predictions"): + return [] + + if name.startswith("cls.seq_relationship"): + return [] + + if self.cls_out_labels: + # For BertForSequenceClassification (direct projection layer) + if name == "classifier.weight": + name = "classifier.out_proj.weight" + + if name == "classifier.bias": + name = "classifier.out_proj.bias" + + return [(self.map_tensor_name(name), data_torch)] + + def _xlmroberta_tokenizer_init(self) -> None: + # we need the pad_token_id to know how to chop down position_embd matrix + if (pad_token_id := self.hparams.get("pad_token_id")) is not None: + self._position_offset = 1 + pad_token_id + if "max_position_embeddings" in self.hparams: + self.hparams["max_position_embeddings"] -= self._position_offset + else: + self._position_offset = None + + def _xlmroberta_set_vocab(self) -> None: + # to avoid TypeError: Descriptors cannot be created directly + # exception when importing sentencepiece_model_pb2 + os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" + from sentencepiece import SentencePieceProcessor + from sentencepiece import sentencepiece_model_pb2 as model + + tokenizer_path = self.dir_model / 'sentencepiece.bpe.model' + + tokenizer_json = {} + tokenizer_config_json = {} + if not tokenizer_path.is_file(): + tokenizer_path = self.dir_model / 'tokenizer.json' + tokenizer_config_path = self.dir_model / 'tokenizer_config.json' + + if not tokenizer_path.is_file(): + raise FileNotFoundError(f"File not found: {tokenizer_path}") + + from base64 import b64decode + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + + with open(tokenizer_path, "r", encoding="utf-8") as fp: + tokenizer_json = json.load(fp) + + if tokenizer_config_path.is_file(): + with open(tokenizer_config_path, "r", encoding="utf-8") as fp: + tokenizer_config_json = json.load(fp) + + add_prefix = tokenizer.add_prefix_space + remove_whitespaces = tokenizer.clean_up_tokenization_spaces + precompiled_charsmap = b64decode(tokenizer_json["normalizer"]["precompiled_charsmap"]) + + vocab_size = max(self.hparams.get("vocab_size", 0), tokenizer.vocab_size) + else: + sentencepiece_model = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] + sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read()) + assert sentencepiece_model.trainer_spec.model_type == 1 # UNIGRAM + + add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix + remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces + precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = max(self.hparams.get("vocab_size", 0), tokenizer.vocab_size()) + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + if isinstance(tokenizer, SentencePieceProcessor): + for token_id in range(tokenizer.vocab_size()): + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + else: + added_vocab = tokenizer.get_added_vocab() + unk_token = tokenizer_config_json.get("unk_token") + unk_token_id = added_vocab.get(unk_token, tokenizer_json["model"].get("unk_id", 3)) + + for token_id in range(tokenizer.vocab_size): + piece = tokenizer._convert_id_to_token(token_id) + if (piece := tokenizer._convert_id_to_token(token_id)) is not None: + text = piece.encode("utf-8") + score = tokenizer_json["model"]["vocab"][token_id][1] + + toktype = SentencePieceTokenTypes.NORMAL + if token_id == unk_token_id: + toktype = SentencePieceTokenTypes.UNKNOWN + elif token_id in tokenizer.all_special_ids: + toktype = SentencePieceTokenTypes.CONTROL + elif token_id in added_vocab.values(): + toktype = SentencePieceTokenTypes.USER_DEFINED + # No reliable way to detect this, but jina doesn't have any + # elif tokenizer.IsByte(token_id): + # toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + if isinstance(tokenizer, SentencePieceProcessor): + # realign tokens (see HF tokenizer code) + tokens = [b'', b'', b'', b''] + tokens[3:-1] + scores = [0.0, 0.0, 0.0, 0.0] + scores[3:-1] + toktypes = [ + SentencePieceTokenTypes.CONTROL, + SentencePieceTokenTypes.CONTROL, + SentencePieceTokenTypes.CONTROL, + SentencePieceTokenTypes.UNKNOWN, + ] + toktypes[3:-1] + + if self.model_arch == gguf.MODEL_ARCH.NOMIC_BERT_MOE: + # Add mask token missing from sentencepiece.bpe.model + tokens[250001] = b'' + scores[250001] = 0.0 + toktypes[250001] = SentencePieceTokenTypes.CONTROL + + self.gguf_writer.add_tokenizer_model("t5") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(add_prefix) + self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1)) + self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces) + if precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + +@ModelBase.register("DistilBertModel", "DistilBertForMaskedLM", "DistilBertForSequenceClassification") +class DistilBertModel(BertModel): + model_arch = gguf.MODEL_ARCH.BERT + + def set_gguf_parameters(self): + self.gguf_writer.add_layer_norm_eps(1e-12) + logger.info("gguf: layer norm epsilon = 1e-12") + super().set_gguf_parameters() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("distilbert."): + name = name[11:] + + # These layers act as MLM head, so we don't need them + if name.startswith("vocab_"): + return [] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("RobertaModel", "RobertaForSequenceClassification") +class RobertaModel(BertModel): + model_arch = gguf.MODEL_ARCH.BERT + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # we need the pad_token_id to know how to chop down position_embd matrix + if (pad_token_id := self.hparams.get("pad_token_id")) is not None: + self._position_offset = 1 + pad_token_id + if "max_position_embeddings" in self.hparams: + self.hparams["max_position_embeddings"] -= self._position_offset + else: + self._position_offset = None + + def set_vocab(self): + """Support BPE tokenizers for roberta models""" + bpe_tok_path = self.dir_model / "tokenizer.json" + if bpe_tok_path.exists(): + self._set_vocab_gpt2() + + # we need this to validate the size of the token_type embeddings + # though currently we are passing all zeros to the token_type embeddings + # "Sequence A" or "Sequence B" + self.gguf_writer.add_token_type_count(self.hparams.get("type_vocab_size", 1)) + + else: + return super().set_vocab() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # if name starts with "roberta.", remove the prefix + # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main + if name.startswith("roberta."): + name = name[8:] + + # position embeddings start at pad_token_id + 1, so just chop down the weight tensor + if name == "embeddings.position_embeddings.weight": + if self._position_offset is not None: + data_torch = data_torch[self._position_offset:,:] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("NomicBertModel") +class NomicBertModel(BertModel): + model_arch = gguf.MODEL_ARCH.BERT + + def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, **kwargs: Any): + hparams = kwargs.pop("hparams", None) + if hparams is None: + hparams = ModelBase.load_hparams(dir_model) + + self.is_moe = bool(hparams.get("moe_every_n_layers")) + self.model_arch = gguf.MODEL_ARCH.NOMIC_BERT_MOE if self.is_moe else gguf.MODEL_ARCH.NOMIC_BERT + + super().__init__(dir_model, ftype, fname_out, hparams=hparams, **kwargs) + + self._tokenizer_is_xlmroberta = self._is_tokenizer_xlmroberta() + if self._tokenizer_is_xlmroberta: + self._xlmroberta_tokenizer_init() + + npos, mtp = self.hparams["n_positions"], self.hparams.get("max_trained_positions", 2048) + if npos == 8192 and mtp == 2048: + self.hparams["n_positions"] = 2048 # nomic-embed-text v1 and v1.5 are trained for 2048 tokens. + elif npos == 2048 and mtp == 2048: + self.hparams["n_positions"] = 512 # nomic-embed-text-v2-moe is trained for 512 tokens. + else: + raise ValueError(f"unrecognized parameters: n_positions={npos}, max_trained_positions={mtp}") + + assert self.hparams["activation_function"] == "gelu" if self.is_moe else "swiglu" + + # this doesn't do anything in the HF version + assert self.hparams["causal"] is False + # no bias tensors unless MoE + assert self.hparams["qkv_proj_bias"] == self.is_moe + assert self.hparams["mlp_fc1_bias"] == self.is_moe + assert self.hparams["mlp_fc2_bias"] == self.is_moe + + # norm at end of layer + assert self.hparams["prenorm"] is False + # standard RoPE + assert self.hparams["rotary_emb_fraction"] == 1.0 + assert self.hparams["rotary_emb_interleaved"] is False + assert self.hparams["rotary_emb_scale_base"] is None + + def set_vocab(self) -> None: + if self._tokenizer_is_xlmroberta: + return self._xlmroberta_set_vocab() + return super().set_vocab() + + def modify_tensors(self, data_torch: torch.Tensor, name: str, bid: int | None) -> Iterable[tuple[str, torch.Tensor]]: + # If the tensor is an experts bias tensor, skip it by returning an empty list. + if "mlp.experts.bias" in name: + return [] # Explicitly return an empty list. + + if "mlp.experts.mlp.w1" in name: + data_torch = data_torch.view(self.hparams["num_experts"], self.hparams["n_inner"], self.hparams["n_embd"]) + name += ".weight" + + if "mlp.experts.mlp.w2" in name: + data_torch = data_torch.view(self.hparams["num_experts"], self.hparams["n_inner"], self.hparams["n_embd"]) + data_torch = data_torch.transpose(1, 2) + name += ".weight" + + return [(self.map_tensor_name(name), data_torch)] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_rope_freq_base(self.hparams["rotary_emb_base"]) + if self.is_moe: + self.gguf_writer.add_moe_every_n_layers(self.hparams["moe_every_n_layers"]) + self.gguf_writer.add_expert_count(self.hparams["num_experts"]) + self.gguf_writer.add_expert_used_count(self.hparams["moe_top_k"]) + + def _is_tokenizer_xlmroberta(self) -> bool: + with open(self.dir_model / "tokenizer.json") as f: + tokenizer_json = json.load(f) + toktyp = tokenizer_json["model"]["type"] + if toktyp == "Unigram": + return True + if toktyp == "WordPiece": + return False + raise ValueError(f"unknown tokenizer: {toktyp}") + + +@ModelBase.register("NeoBERT", "NeoBERTLMHead", "NeoBERTForSequenceClassification") +class NeoBert(BertModel): + model_arch = gguf.MODEL_ARCH.NEO_BERT + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + # NeoBERT uses 2/3 of the intermediate size as feed forward length + self.gguf_writer.add_feed_forward_length(int(2 * self.hparams["intermediate_size"] / 3)) + self.gguf_writer.add_rope_freq_base(10000.0) # default value for NeoBERT + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + + f_rms_eps = self.hparams.get("norm_eps", 1e-6) # default value for NeoBERT + self.gguf_writer.add_layer_norm_rms_eps(f_rms_eps) + logger.info(f"gguf: rms norm epsilon = {f_rms_eps}") + + self.gguf_writer.add_pooling_type(gguf.PoolingType.CLS) # https://huggingface.co/chandar-lab/NeoBERT#how-to-use + + def modify_tensors(self, data_torch, name, bid): + if name.startswith("decoder."): + return [] + + if name.startswith("model."): + name = name[6:] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("XLMRobertaModel", "XLMRobertaForSequenceClassification") +class XLMRobertaModel(BertModel): + model_arch = gguf.MODEL_ARCH.BERT + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._xlmroberta_tokenizer_init() + + def set_vocab(self): + self._xlmroberta_set_vocab() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # if name starts with "roberta.", remove the prefix + # e.g. https://huggingface.co/BAAI/bge-reranker-v2-m3/tree/main + if name.startswith("roberta."): + name = name[8:] + + # position embeddings start at pad_token_id + 1, so just chop down the weight tensor + if name == "embeddings.position_embeddings.weight": + if self._position_offset is not None: + data_torch = data_torch[self._position_offset:,:] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("GemmaForCausalLM") +class GemmaModel(TextModel): + model_arch = gguf.MODEL_ARCH.GEMMA + + def set_vocab(self): + self._set_vocab_sentencepiece() + + # TODO: these special tokens should be exported only for the CodeGemma family + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=False, + special_token_types = ['prefix', 'suffix', 'middle', 'fsep', 'eot']) + special_vocab._set_special_token("prefix", 67) + special_vocab._set_special_token("suffix", 69) + special_vocab._set_special_token("middle", 68) + special_vocab._set_special_token("fsep", 70) + special_vocab._set_special_token("eot", 107) + special_vocab.chat_template = None # do not add it twice + special_vocab.add_to_gguf(self.gguf_writer) + + self.gguf_writer.add_add_space_prefix(False) + + def set_gguf_parameters(self): + hparams = self.hparams + block_count = hparams["num_hidden_layers"] + + self.gguf_writer.add_context_length(hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(self.hparams["num_key_value_heads"] if "num_key_value_heads" in hparams else hparams["num_attention_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_key_length(hparams["head_dim"]) + self.gguf_writer.add_value_length(hparams["head_dim"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + # lm_head is not used in llama.cpp, while autoawq will include this tensor in model + # To prevent errors, skip loading lm_head.weight. + if name == "lm_head.weight": + logger.debug(f"Skipping get tensor {name!r} in safetensors so that convert can end normally.") + return [] + + # ref: https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/gemma/modeling_gemma.py#L89 + if name.endswith("norm.weight"): + data_torch = data_torch + 1 + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("Gemma2ForCausalLM") +class Gemma2Model(TextModel): + model_arch = gguf.MODEL_ARCH.GEMMA2 + + def set_vocab(self): + self._set_vocab_sentencepiece() + + self.gguf_writer.add_add_space_prefix(False) + + def set_gguf_parameters(self): + hparams = self.hparams + block_count = hparams["num_hidden_layers"] + + self.gguf_writer.add_context_length(hparams["max_position_embeddings"]) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_head_count(hparams["num_attention_heads"]) + self.gguf_writer.add_head_count_kv(self.hparams["num_key_value_heads"] if "num_key_value_heads" in hparams else hparams["num_attention_heads"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["rms_norm_eps"]) + self.gguf_writer.add_key_length(hparams["head_dim"]) + self.gguf_writer.add_value_length(hparams["head_dim"]) + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_attn_logit_softcapping( + self.hparams["attn_logit_softcapping"] + ) + self.gguf_writer.add_final_logit_softcapping( + self.hparams["final_logit_softcapping"] + ) + self.gguf_writer.add_sliding_window(self.hparams["sliding_window"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + # lm_head is not used in llama.cpp, while autoawq will include this tensor in model + # To prevent errors, skip loading lm_head.weight. + if name == "lm_head.weight": + logger.debug(f"Skipping get tensor {name!r} in safetensors so that convert can end normally.") + return [] + + # ref: https://github.com/huggingface/transformers/blob/fc37f38915372c15992b540dfcbbe00a916d4fc6/src/transformers/models/gemma/modeling_gemma.py#L89 + if name.endswith("norm.weight"): + data_torch = data_torch + 1 + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("Gemma3ForCausalLM", "Gemma3ForConditionalGeneration") +class Gemma3Model(TextModel): + model_arch = gguf.MODEL_ARCH.GEMMA3 + norm_shift = 1.0 # Gemma3RMSNorm adds 1.0 to the norm value + + def set_vocab(self): + self._set_vocab_sentencepiece() + + self.gguf_writer.add_add_space_prefix(False) + + def set_gguf_parameters(self): + hparams = self.hparams + block_count = hparams["num_hidden_layers"] + + # some default values are not specified in the hparams + self.gguf_writer.add_context_length(hparams.get("max_position_embeddings", 131072)) + self.gguf_writer.add_embedding_length(hparams["hidden_size"]) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_feed_forward_length(hparams["intermediate_size"]) + self.gguf_writer.add_head_count(hparams.get("num_attention_heads", 8)) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("rms_norm_eps", 1e-6)) + self.gguf_writer.add_key_length(hparams.get("head_dim", 256)) + self.gguf_writer.add_value_length(hparams.get("head_dim", 256)) + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_rope_freq_base(hparams.get("rope_theta", 1_000_000.0)) # for global layers + # attn_logit_softcapping is removed in Gemma3 + assert hparams.get("attn_logit_softcapping") is None + self.gguf_writer.add_sliding_window(hparams["sliding_window"]) + self.gguf_writer.add_head_count_kv(hparams.get("num_key_value_heads", 4)) + if hparams.get("rope_scaling") is not None: + assert hparams["rope_scaling"]["rope_type"] == "linear" + # important: this rope_scaling is only applied for global layers, and not used by 1B model + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(hparams["rope_scaling"]["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if "language_model." in name: + name = name.replace("language_model.", "") + + elif name.startswith("multi_modal_projector.") or name.startswith("vision_tower.") \ + or name.startswith("multimodal_projector.") or name.startswith("vision_model."): + return [] # skip vision tensors + + # remove OOV (out-of-vocabulary) rows in token_embd + if "embed_tokens.weight" in name: + vocab = self._create_vocab_sentencepiece() + tokens = vocab[0] + data_torch = data_torch[:len(tokens)] + + # ref code in Gemma3RMSNorm + # output = output * (1.0 + self.weight.float()) + # note: this is not the case on gemma3n + if name.endswith("norm.weight"): + data_torch = data_torch + self.norm_shift + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("Gemma3ForConditionalGeneration") +class Gemma3VisionModel(MmprojModel): + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GEMMA3) + # default values below are taken from HF tranformers code + self.gguf_writer.add_vision_attention_layernorm_eps(hparams.get("layer_norm_eps", 1e-6)) + self.gguf_writer.add_vision_use_gelu(True) + # calculate proj_scale_factor (used by tinygemma3 test model) + image_seq_length = self.preprocessor_config.get("image_seq_length", 256) + n_per_side = int(image_seq_length ** 0.5) + image_size = self.hparams["image_size"] + patch_size = self.hparams["patch_size"] + proj_scale_factor = (image_size // patch_size) // n_per_side + if proj_scale_factor > 0 and proj_scale_factor != 4: + # we only need to write this if it's not the default value + # in this case, we are converting a test model + self.gguf_writer.add_vision_projector_scale_factor(proj_scale_factor) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, new_name, n_dims # unused + # related to https://github.com/ggml-org/llama.cpp/issues/13025 + if "input_projection" in name: + return gguf.GGMLQuantizationType.F16 + if ".embeddings." in name: + return gguf.GGMLQuantizationType.F32 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if "vision_model.head." in name: + return [] # skip redundant tensors for tinygemma3 + + if name.startswith("multi_modal_projector.") or name.startswith("vision_tower.") \ + or name.startswith("multimodal_projector.") or name.startswith("vision_model."): + # process vision tensors + name = name.replace("_weight", ".weight") + + # correct norm value ; only this "soft_emb_norm" need to be corrected as it's part of Gemma projector + # the other norm values are part of SigLIP model, and they are already correct + # ref code: Gemma3RMSNorm + if "soft_emb_norm.weight" in name: + logger.info(f"Correcting norm value for '{name}'") + data_torch = data_torch + 1 + + return [(self.map_tensor_name(name), data_torch)] + + return [] # skip other tensors + + +@ModelBase.register("Gemma3nForConditionalGeneration") +class Gemma3NModel(Gemma3Model): + model_arch = gguf.MODEL_ARCH.GEMMA3N + norm_shift = 0.0 # same value with Gemma3p5RMSNorm scale_shift on python code + + _altup_proj: list[Tensor] = [] + _altup_unembd: list[Tensor] = [] + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams["altup_num_inputs"] == 4, "Current conversion only supports 4 altup inputs" + self._altup_proj = [ + torch.Tensor(), # to be replaced + torch.Tensor(), # to be replaced + torch.Tensor(), # to be replaced + ] + self._altup_unembd = [ + torch.Tensor(), # to be replaced + torch.Tensor(), # to be replaced + torch.Tensor(), # to be replaced + ] + + def set_vocab(self): + super().set_vocab() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_altup_active_idx(self.hparams["altup_active_idx"]) + self.gguf_writer.add_altup_num_inputs(self.hparams["altup_num_inputs"]) + self.gguf_writer.add_embedding_length_per_layer_input(self.hparams["hidden_size_per_layer_input"]) + self.gguf_writer.add_shared_kv_layers(self.hparams["num_kv_shared_layers"]) + + activation_sparsity_scale = [] + for s in self.hparams["activation_sparsity_pattern"]: + normal_dist = torch.distributions.normal.Normal(0, 1) + std_multiplier = normal_dist.icdf(torch.tensor(s, dtype=torch.float32)) + activation_sparsity_scale.append(std_multiplier.item()) + self.gguf_writer.add_activation_sparsity_scale(activation_sparsity_scale) + + sliding_window_pattern = [] + for t in self.hparams["layer_types"]: + sliding_window_pattern.append(t == "sliding_attention") + self.gguf_writer.add_sliding_window_pattern(sliding_window_pattern) + + def _stack_matrices(self, matrices: list[Tensor]) -> Tensor | None: + has_all = all(m.numel() > 0 for m in matrices) + if not has_all: + return None + else: + return torch.stack(matrices, dim=0) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.endswith("_scale"): + name = name + ".weight" + + # TODO: implement self.prediction_coefs.weight.clamp_(...) + + if "language_model." not in name: + return [] # skip non-language model tensors + + if "altup_unembed_projections" in name: + data_torch = data_torch.to(device="cpu") + if ".0." in name: + self._altup_unembd[0] = data_torch + elif ".1." in name: + self._altup_unembd[1] = data_torch + elif ".2." in name: + self._altup_unembd[2] = data_torch + else: + raise ValueError(f"Unknown name: {name}") + out = self._stack_matrices(self._altup_unembd) + if out is not None: + return [(self.map_tensor_name("model.altup_unembed_projections.weight"), out)] + else: + return [] + + if "altup_projections" in name: + data_torch = data_torch.to(device="cpu") + if ".0." in name: + self._altup_proj[0] = data_torch + elif ".1." in name: + self._altup_proj[1] = data_torch + elif ".2." in name: + self._altup_proj[2] = data_torch + else: + raise ValueError(f"Unknown name: {name}") + out = self._stack_matrices(self._altup_proj) + if out is not None: + return [(self.map_tensor_name("model.altup_projections.weight"), out)] + else: + return [] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("Starcoder2ForCausalLM") +class StarCoder2Model(TextModel): + model_arch = gguf.MODEL_ARCH.STARCODER2 + + +@ModelBase.register("Rwkv6ForCausalLM") +class Rwkv6Model(TextModel): + model_arch = gguf.MODEL_ARCH.RWKV6 + + def set_vocab(self): + self._set_vocab_rwkv_world() + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + head_size = self.hparams["head_size"] + hidden_size = self.hparams["hidden_size"] + layer_norm_eps = self.hparams["layer_norm_epsilon"] + rescale_every_n_layers = self.hparams["rescale_every"] + intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else int((hidden_size * 3.5) // 32 * 32) + time_mix_extra_dim = 64 if hidden_size == 4096 else 32 + time_decay_extra_dim = 128 if hidden_size == 4096 else 64 + + # RWKV isn't context limited + self.gguf_writer.add_context_length(1048576) + self.gguf_writer.add_embedding_length(hidden_size) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_layer_norm_eps(layer_norm_eps) + self.gguf_writer.add_rescale_every_n_layers(rescale_every_n_layers) + self.gguf_writer.add_wkv_head_size(head_size) + self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim) + self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim) + self.gguf_writer.add_feed_forward_length(intermediate_size) + self.gguf_writer.add_file_type(self.ftype) + + # required by llama.cpp, unused + self.gguf_writer.add_head_count(0) + + lerp_weights: dict[int, dict[str, Tensor]] = {} + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + new_name = self.map_tensor_name(name) + + if not (new_name.endswith(".weight") or new_name.endswith(".bias")): + new_name += ".weight" + + if new_name.endswith("time_mix_w1.weight") or new_name.endswith("time_mix_decay_w1.weight") or new_name.endswith("time_mix_decay_w2.weight"): + data_torch = data_torch.transpose(0, 1) + + if new_name.endswith("time_mix_w2.weight"): + data_torch = data_torch.permute(0, 2, 1) + + if new_name.endswith("time_mix_decay.weight") or "lerp" in new_name: + data_torch = data_torch.squeeze() + + try: + rescale_every_n_layers = self.hparams["rescale_every"] + if rescale_every_n_layers > 0: + if new_name.endswith("time_mix_output.weight") or new_name.endswith("channel_mix_value.weight"): + data_torch = data_torch.div_(2 ** int(bid // rescale_every_n_layers)) + except KeyError: + pass + + # concat time_mix_lerp weights to reduce some cpu overhead + # also reduces the number of tensors in the model + if bid is not None and "time_mix_lerp" in new_name and "time_mix_lerp_x" not in new_name: + try: + self.lerp_weights[bid][new_name] = data_torch + except KeyError: + self.lerp_weights[bid] = {new_name: data_torch} + if all(f"blk.{bid}.time_mix_lerp_{i}.weight" in self.lerp_weights[bid].keys() for i in ["w", "k", "v", "r", "g"]): + new_name = f"blk.{bid}.time_mix_lerp_fused.weight" + data = torch.stack([self.lerp_weights[bid][f"blk.{bid}.time_mix_lerp_{i}.weight"].unsqueeze(0) for i in ["w", "k", "v", "r", "g"]], dim=0).unsqueeze(1) + yield (new_name, data) + return + + yield (new_name, data_torch) + + +@ModelBase.register("RWKV6Qwen2ForCausalLM") +class RWKV6Qwen2Model(Rwkv6Model): + model_arch = gguf.MODEL_ARCH.RWKV6QWEN2 + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + num_attention_heads = self.hparams["num_attention_heads"] + num_key_value_heads = self.hparams["num_key_value_heads"] + hidden_size = self.hparams["hidden_size"] + head_size = hidden_size // num_attention_heads + rms_norm_eps = self.hparams["rms_norm_eps"] + intermediate_size = self.hparams["intermediate_size"] + time_mix_extra_dim = self.hparams.get("lora_rank_tokenshift", 64 if hidden_size >= 4096 else 32) + time_decay_extra_dim = self.hparams.get("lora_rank_decay", 128 if hidden_size >= 4096 else 64) + + # RWKV isn't context limited + self.gguf_writer.add_context_length(1048576) + self.gguf_writer.add_embedding_length(hidden_size) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_wkv_head_size(head_size) + self.gguf_writer.add_time_mix_extra_dim(time_mix_extra_dim) + self.gguf_writer.add_time_decay_extra_dim(time_decay_extra_dim) + self.gguf_writer.add_feed_forward_length(intermediate_size) + self.gguf_writer.add_file_type(self.ftype) + + # special parameters for time_mixing in RWKV6QWEN2 + self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps) + self.gguf_writer.add_token_shift_count(1) + # RWKV6QWEN2 use grouped key/value like GQA + self.gguf_writer.add_head_count_kv(num_key_value_heads) + + # required by llama.cpp, unused + self.gguf_writer.add_head_count(0) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + for new_name, data in super().modify_tensors(data_torch, name, bid): + if "time_mix_w1" in new_name or "time_mix_w2" in new_name: + data = data.view(5, -1, data.shape[-1]) + # rwkv6qwen2 has a different order of rkvwg instead of the original wkvrg + # permute them here to avoid code changes + data = torch.stack([data[3], data[1], data[2], data[0], data[4]], dim=0).view(-1, data.shape[-1]) + if "w2" in new_name: + data = data.view(5, -1, data.shape[-1]) + yield (new_name, data) + continue + yield (new_name, data) + + +@ModelBase.register("Rwkv7ForCausalLM", "RWKV7ForCausalLM") +class Rwkv7Model(TextModel): + model_arch = gguf.MODEL_ARCH.RWKV7 + + def set_vocab(self): + self._set_vocab_rwkv_world() + + def calc_lora_rank(self, hidden_size, exponent, multiplier): + return max(1, round(hidden_size ** exponent * multiplier / 32)) * 32 + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + try: + head_size = self.hparams["head_size"] + layer_norm_eps = self.hparams["layer_norm_epsilon"] + except KeyError: + head_size = self.hparams["head_dim"] + layer_norm_eps = self.hparams["norm_eps"] + hidden_size = self.hparams["hidden_size"] + intermediate_size = self.hparams["intermediate_size"] if self.hparams["intermediate_size"] is not None else (hidden_size * 4) + + # ICLR: In-Context-Learning-Rate + try: + lora_rank_decay = self.hparams["lora_rank_decay"] if self.hparams["lora_rank_decay"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8) + lora_rank_iclr = self.hparams["lora_rank_iclr"] if self.hparams["lora_rank_iclr"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8) + lora_rank_value_residual_mix = self.hparams["lora_rank_value_residual_mix"] if self.hparams["lora_rank_value_residual_mix"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3) + lora_rank_gate = self.hparams["lora_rank_gate"] if self.hparams["lora_rank_gate"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6) + except KeyError: + lora_rank_decay = self.hparams["decay_low_rank_dim"] if self.hparams["decay_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8) + lora_rank_iclr = self.hparams["a_low_rank_dim"] if self.hparams["a_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.8) + lora_rank_value_residual_mix = self.hparams["v_low_rank_dim"] if self.hparams["v_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.5, 1.3) + lora_rank_gate = self.hparams["gate_low_rank_dim"] if self.hparams["gate_low_rank_dim"] is not None else self.calc_lora_rank(hidden_size, 0.8, 0.6) + + # RWKV isn't context limited + self.gguf_writer.add_context_length(1048576) + self.gguf_writer.add_embedding_length(hidden_size) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_layer_norm_eps(layer_norm_eps) + self.gguf_writer.add_wkv_head_size(head_size) + self.gguf_writer.add_decay_lora_rank(lora_rank_decay) + self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr) + self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix) + self.gguf_writer.add_gate_lora_rank(lora_rank_gate) + self.gguf_writer.add_feed_forward_length(intermediate_size) + self.gguf_writer.add_file_type(self.ftype) + + # required by llama.cpp, unused + self.gguf_writer.add_head_count(0) + + lerp_weights: dict[int, dict[str, Tensor]] = {} + lora_needs_transpose: bool = True + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # unify tensor names here to make life easier + name = name.replace("blocks", "layers").replace("ffn", "feed_forward") + name = name.replace("self_attn", "attention").replace("attn", "attention") + name = name.replace("time_mixer.", "") + # lora layer names in fla-hub's impl + if "_lora.lora" in name: + self.lora_needs_transpose = False + name = name.replace("_lora.lora.0.weight", "1.weight") + name = name.replace("_lora.lora.2.weight", "2.weight") + name = name.replace("_lora.lora.2.bias", "0.weight") + + name = name.replace("feed_forward_norm", "ln2") + name = name.replace("g_norm", "ln_x") + + if "attention.v" in name and "value" not in self.map_tensor_name(name) and bid == 0: + # some models have dummy v0/v1/v2 on first layer while others don't + # ignore them all since they are not used + return + + wkv_has_gate = self.hparams.get("wkv_has_gate", True) + lerp_list = ["r", "w", "k", "v", "a", "g"] if wkv_has_gate else ["r", "w", "k", "v", "a"] + + if bid is not None and "attention.x_" in name: + if "attention.x_x" in name: + # already concatenated + new_name = f"blk.{bid}.time_mix_lerp_fused.weight" + data = data_torch.reshape(len(lerp_list), 1, 1, -1) + yield (new_name, data) + else: + try: + self.lerp_weights[bid][name] = data_torch + except KeyError: + self.lerp_weights[bid] = {name: data_torch} + if all(f"model.layers.{bid}.attention.x_{i}" in self.lerp_weights[bid].keys() for i in lerp_list): + new_name = f"blk.{bid}.time_mix_lerp_fused.weight" + data = torch.stack([self.lerp_weights[bid][f"model.layers.{bid}.attention.x_{i}"] for i in lerp_list], dim=0) + yield (new_name, data) + return + else: + data_torch = data_torch.squeeze() + new_name = self.map_tensor_name(name) + + if not (new_name.endswith(".weight") or new_name.endswith(".bias")): + new_name += ".weight" + + if self.lora_needs_transpose and any( + new_name.endswith(t) for t in [ + "time_mix_w1.weight", "time_mix_w2.weight", + "time_mix_a1.weight", "time_mix_a2.weight", + "time_mix_v1.weight", "time_mix_v2.weight", + "time_mix_g1.weight", "time_mix_g2.weight", + ] + ): + data_torch = data_torch.transpose(0, 1) + + if 'r_k' in new_name: + data_torch = data_torch.flatten() + + if bid == 0 and "time_mix_a" in new_name: + # dummy v0/v1/v2 on first layer + # easist way to make llama happy + yield (new_name.replace("time_mix_a", "time_mix_v"), data_torch) + + yield (new_name, data_torch) + + +@ModelBase.register("RwkvHybridForCausalLM") +class ARwkv7Model(Rwkv7Model): + model_arch = gguf.MODEL_ARCH.ARWKV7 + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + block_count = self.hparams["num_hidden_layers"] + hidden_size = self.hparams["hidden_size"] + head_size = self.hparams["head_size"] + rms_norm_eps = self.hparams["rms_norm_eps"] + intermediate_size = self.hparams["intermediate_size"] + wkv_has_gate = self.hparams["wkv_has_gate"] + assert self.hparams["wkv_version"] == 7 + + # ICLR: In-Context-Learning-Rate + lora_rank_decay = 64 + lora_rank_iclr = 64 + lora_rank_value_residual_mix = 32 + lora_rank_gate = 128 if wkv_has_gate else 0 + + # RWKV isn't context limited + self.gguf_writer.add_context_length(1048576) + self.gguf_writer.add_embedding_length(hidden_size) + self.gguf_writer.add_block_count(block_count) + self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps) + self.gguf_writer.add_wkv_head_size(head_size) + self.gguf_writer.add_decay_lora_rank(lora_rank_decay) + self.gguf_writer.add_iclr_lora_rank(lora_rank_iclr) + self.gguf_writer.add_value_residual_mix_lora_rank(lora_rank_value_residual_mix) + self.gguf_writer.add_gate_lora_rank(lora_rank_gate) + self.gguf_writer.add_feed_forward_length(intermediate_size) + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_token_shift_count(1) + + # required by llama.cpp, unused + self.gguf_writer.add_head_count(0) + + +@ModelBase.register("MambaForCausalLM", "MambaLMHeadModel", "FalconMambaForCausalLM") +class MambaModel(TextModel): + model_arch = gguf.MODEL_ARCH.MAMBA + + def __init__(self, dir_model: Path, *args, **kwargs): + # Avoid using AutoConfig for hparams + hparams = kwargs.pop("hparams", None) + if hparams is None: + with open(dir_model / "config.json", "r", encoding="utf-8") as f: + hparams = json.load(f) + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + def set_vocab(self): + vocab_size = self.hparams["vocab_size"] + # Round vocab size to next multiple of 8 + pad_vocab = self.hparams.get("pad_vocab_size_multiple", 8) + # pad using ceiling division + # ref: https://stackoverflow.com/a/17511341/22827863 + vocab_size = -(vocab_size // -pad_vocab) * pad_vocab + self.hparams["vocab_size"] = vocab_size + + if (self.dir_model / "tokenizer.json").is_file(): + self._set_vocab_gpt2() + elif (self.dir_model / "tokenizer.model").is_file(): + self._set_vocab_sentencepiece() + else: + # Use the GPT-NeoX tokenizer when no tokenizer files are present + self._set_vocab_builtin("gpt-neox", vocab_size) + + def set_gguf_parameters(self): + d_model = self.find_hparam(["hidden_size", "d_model"]) + d_conv = self.find_hparam(["conv_kernel", "d_conv"], optional=True) or 4 + d_inner = self.find_hparam(["intermediate_size", "d_inner"], optional=True) or 2 * d_model + d_state = self.find_hparam(["state_size", "d_state"], optional=True) or 16 + # ceiling division + # ref: https://stackoverflow.com/a/17511341/22827863 + # ref: https://github.com/state-spaces/mamba/blob/ce59daea3a090d011d6476c6e5b97f6d58ddad8b/mamba_ssm/modules/mamba_simple.py#L58 + dt_rank = self.find_hparam(["time_step_rank", "dt_rank"], optional=True) or -(d_model // -16) + rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-5 + use_dt_b_c_norm = False + # For falconmamba we do apply RMS norm on B / DT and C layers + if self.find_hparam(["model_type"], optional=True) in ("falcon_mamba",): + use_dt_b_c_norm = True + # Fail early for models which don't have a block expansion factor of 2 + assert d_inner == 2 * d_model + + self.gguf_writer.add_context_length(2**20) # arbitrary value; for those who use the default + self.gguf_writer.add_embedding_length(d_model) + self.gguf_writer.add_feed_forward_length(0) # unused, but seemingly required when loading + self.gguf_writer.add_head_count(0) # unused, but seemingly required when loading + self.gguf_writer.add_block_count(self.block_count) + self.gguf_writer.add_ssm_conv_kernel(d_conv) + self.gguf_writer.add_ssm_inner_size(d_inner) + self.gguf_writer.add_ssm_state_size(d_state) + self.gguf_writer.add_ssm_time_step_rank(dt_rank) + self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps) + self.gguf_writer.add_ssm_dt_b_c_rms(use_dt_b_c_norm) # For classic Mamba we don't apply rms norm on B / DT layers + self.gguf_writer.add_file_type(self.ftype) + + _tok_embd = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + output_name = self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT) + tok_embd_name = self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD) + + new_name = self.map_tensor_name(name) + + if name.endswith(".A_log"): + logger.debug("A_log --> A ==> " + new_name) + data_torch = -torch.exp(data_torch) + + # [4 1 8192 1] -> [4 8192 1 1] + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_CONV1D, bid): + data_torch = data_torch.squeeze() + + # assuming token_embd.weight is seen before output.weight + if self._tok_embd is not None and new_name == output_name: + if torch.equal(self._tok_embd, data_torch): + logger.debug(f"{output_name} is equivalent to {tok_embd_name}, omitting") + return [] + elif new_name == tok_embd_name: + self._tok_embd = data_torch + + return [(new_name, data_torch)] + + +@ModelBase.register("Mamba2ForCausalLM") +class Mamba2Model(TextModel): + model_arch = gguf.MODEL_ARCH.MAMBA2 + + def __init__(self, dir_model: Path, *args, **kwargs): + # Avoid using AutoConfig for hparams + # It wrongly assumes all Mamba2 models are Mamba-Codestral-7B-v0.1 + hparams = kwargs.pop("hparams", None) + if hparams is None: + with open(dir_model / "config.json", "r", encoding="utf-8") as f: + hparams = json.load(f) + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + def set_vocab(self): + vocab_size = self.hparams["vocab_size"] + # Round vocab size to next multiple of 16 + pad_vocab = self.hparams.get("pad_vocab_size_multiple", 16) + # pad using ceiling division + # ref: https://stackoverflow.com/a/17511341/22827863 + vocab_size = -(vocab_size // -pad_vocab) * pad_vocab + self.hparams["vocab_size"] = vocab_size + + if (self.dir_model / "tokenizer.model").is_file(): + self._set_vocab_sentencepiece() + elif (self.dir_model / "tokenizer.model.v3").is_file(): + # mamba-codestral + raise NotImplementedError(f"Please rename {self.dir_model / 'tokenizer.model.v3'} to {self.dir_model / 'tokenizer.model'}") + elif (self.dir_model / "tokenizer.json").is_file(): + self._set_vocab_gpt2() + else: + # Use the GPT-NeoX tokenizer when no tokenizer files are present + self._set_vocab_builtin("gpt-neox", vocab_size) + + def set_gguf_parameters(self): + d_model = self.find_hparam(["hidden_size", "d_model", "dim"]) + d_conv = self.find_hparam(["conv_kernel", "d_conv"], optional=True) or 4 + d_inner = self.find_hparam(["intermediate_size", "d_inner"], optional=True) or 2 * d_model + d_state = self.find_hparam(["state_size", "d_state"], optional=True) or 128 + head_dim = self.find_hparam(["head_dim"], optional=True) or 64 + n_group = self.find_hparam(["n_groups"], optional=True) or 1 + + rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-5 + + # Fail early for models which don't have a block expansion factor of 2 + # TODO: does this really matter? + assert d_inner == 2 * d_model + assert d_inner % head_dim == 0 + + self.gguf_writer.add_context_length(2**20) # arbitrary value; for those who use the default + self.gguf_writer.add_embedding_length(d_model) + self.gguf_writer.add_feed_forward_length(0) # unused, but seemingly required when loading + self.gguf_writer.add_head_count(0) # unused, but seemingly required when loading + self.gguf_writer.add_block_count(self.block_count) + self.gguf_writer.add_ssm_conv_kernel(d_conv) + self.gguf_writer.add_ssm_inner_size(d_inner) + self.gguf_writer.add_ssm_state_size(d_state) + self.gguf_writer.add_ssm_time_step_rank(d_inner // head_dim) + self.gguf_writer.add_ssm_group_count(n_group) + self.gguf_writer.add_layer_norm_rms_eps(rms_norm_eps) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + + if name.startswith("model.backbone") or name.startswith("model.lm_head"): + # map Mamba-Codestral-7B-v0.1 tensor names to the names used by Mamba-2 + name = name.removeprefix("model.") + + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + new_name = self.map_tensor_name(name) + + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_CONV1D, bid): + data_torch = data_torch.squeeze() + elif any(self.match_model_tensor_name(new_name, t, bid, suffix="") for t in [ + gguf.MODEL_TENSOR.SSM_A, + gguf.MODEL_TENSOR.SSM_D, + ]): + # unsqueeze A to use similar shape semantics as Mamba-1 + # (D is also unsqueezed, but for more straightforward broadcast internally) + data_torch = data_torch.reshape((*data_torch.shape, 1)) + elif self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.SSM_NORM, bid): + d_model = self.find_hparam(["hidden_size", "d_model", "dim"]) + d_inner = self.find_hparam(["intermediate_size", "d_inner"], optional=True) or 2 * d_model + n_group = self.hparams.get("n_groups", 1) + data_torch = data_torch.reshape((n_group, d_inner // n_group)) + + if name.endswith(".A_log"): + logger.debug("A_log --> A ==> " + new_name) + data_torch = -torch.exp(data_torch) + + yield (new_name, data_torch) + + +@ModelBase.register("CohereForCausalLM") +class CommandR2Model(TextModel): + model_arch = gguf.MODEL_ARCH.COMMAND_R + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # max_position_embeddings = 8192 in config.json but model was actually + # trained on 128k context length + # aya-23 models don't have model_max_length specified + self.hparams["max_position_embeddings"] = self.find_hparam(["model_max_length", "max_position_embeddings"]) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_logit_scale(self.hparams["logit_scale"]) + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + + +@ModelBase.register("Cohere2ForCausalLM") +class Cohere2Model(TextModel): + model_arch = gguf.MODEL_ARCH.COHERE2 + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + self.gguf_writer.add_logit_scale(self.hparams["logit_scale"]) + self.gguf_writer.add_sliding_window(self.hparams["sliding_window"]) + self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) + + rotary_pct = self.hparams["rotary_pct"] + hidden_size = self.hparams["hidden_size"] + num_attention_heads = self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(int(rotary_pct * (hidden_size // num_attention_heads))) + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + + +@ModelBase.register("OlmoForCausalLM") +@ModelBase.register("OLMoForCausalLM") +class OlmoModel(TextModel): + model_arch = gguf.MODEL_ARCH.OLMO + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_layer_norm_eps(1e-5) + clip_qkv = self.hparams.get("clip_qkv") + if clip_qkv is not None: + self.gguf_writer.add_clamp_kqv(clip_qkv) + + # Same as super class, but permuting q_proj, k_proj + # Copied from: LlamaModel + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + + if name.endswith("q_proj.weight"): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith("k_proj.weight"): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("Olmo2ForCausalLM") +class Olmo2Model(TextModel): + model_arch = gguf.MODEL_ARCH.OLMO2 + + +@ModelBase.register("OlmoeForCausalLM") +class OlmoeModel(TextModel): + model_arch = gguf.MODEL_ARCH.OLMOE + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_layer_norm_rms_eps(1e-5) + if (n_experts := self.hparams.get("num_experts")) is not None: + self.gguf_writer.add_expert_count(n_experts) + + _experts: list[dict[str, Tensor]] | None = None + + # Copied from: Qwen2MoeModel + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # process the experts separately + if name.find("experts") != -1: + n_experts = self.hparams["num_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + # Copied from: Qwen2MoeModel + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("JinaBertModel", "JinaBertForMaskedLM") +class JinaBertV2Model(BertModel): + model_arch = gguf.MODEL_ARCH.JINA_BERT_V2 + + def set_vocab(self): + tokenizer_class = 'BertTokenizer' + with open(self.dir_model / "tokenizer_config.json", "r", encoding="utf-8") as f: + tokenizer_class = json.load(f)['tokenizer_class'] + + if tokenizer_class == 'BertTokenizer': + super().set_vocab() + elif tokenizer_class == 'RobertaTokenizer': + self._set_vocab_gpt2() + self.gguf_writer.add_token_type_count(2) + else: + raise NotImplementedError(f'Tokenizer {tokenizer_class} is not supported for JinaBertModel') + + +@ModelBase.register("OpenELMForCausalLM") +class OpenELMModel(TextModel): + model_arch = gguf.MODEL_ARCH.OPENELM + + @staticmethod + def _make_divisible(v: float | int, divisor: int) -> int: + # ref: https://huggingface.co/apple/OpenELM-270M-Instruct/blob/eb111ff2e6724348e5b905984063d4064d4bc579/configuration_openelm.py#L34-L38 + new_v = max(divisor, int(v + divisor / 2) // divisor * divisor) + # Make sure that round down does not go down by more than 10%. + if new_v < 0.9 * v: + new_v += divisor + return new_v + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + ffn_multipliers: list[float] = self.hparams["ffn_multipliers"] + ffn_dim_divisor: int = self.hparams["ffn_dim_divisor"] + self._n_embd: int = self.hparams["model_dim"] + self._num_kv_heads: list[int] = self.hparams["num_kv_heads"] + self._num_query_heads: list[int] = self.hparams["num_query_heads"] + self._ffn_dims: list[int] = [ + OpenELMModel._make_divisible(multiplier * self._n_embd, ffn_dim_divisor) + for multiplier in ffn_multipliers + ] + assert isinstance(self._num_kv_heads, list) and isinstance(self._num_kv_heads[0], int) + assert isinstance(self._num_query_heads, list) and isinstance(self._num_query_heads[0], int) + + # Uses the tokenizer from meta-llama/Llama-2-7b-hf + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_builtin("llama-spm", self.hparams["vocab_size"]) + + def set_gguf_parameters(self): + n_embd = self._n_embd + head_dim = self.hparams["head_dim"] + rot_pct = 1.0 + assert self.block_count == len(self._num_kv_heads) + assert self.block_count == len(self._num_query_heads) + assert self.block_count == len(self._ffn_dims) + + self.gguf_writer.add_block_count(self.block_count) + self.gguf_writer.add_context_length(self.hparams["max_context_length"]) + self.gguf_writer.add_embedding_length(n_embd) + self.gguf_writer.add_feed_forward_length(self._ffn_dims) + self.gguf_writer.add_head_count(self._num_query_heads) + self.gguf_writer.add_head_count_kv(self._num_kv_heads) + self.gguf_writer.add_rope_freq_base(self.hparams["rope_freq_constant"]) + # https://huggingface.co/apple/OpenELM-270M-Instruct/blob/c401df2/modeling_openelm.py#L30 + self.gguf_writer.add_layer_norm_rms_eps(1e-6) + self.gguf_writer.add_rope_dimension_count(int(rot_pct * head_dim)) + self.gguf_writer.add_key_length(head_dim) + self.gguf_writer.add_value_length(head_dim) + self.gguf_writer.add_file_type(self.ftype) + + def find_hparam(self, keys: Iterable[str], optional: bool = False) -> Any: + if "n_layers" in keys: + return self.hparams["num_transformer_layers"] + + return super().find_hparam(keys, optional) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + + # split ff + if bid is not None and name == f"transformer.layers.{bid}.ffn.proj_1.weight": + ff_dim = self._ffn_dims[bid] + yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE, bid), data_torch[:ff_dim]) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP, bid), data_torch[ff_dim:]) + return + + yield (self.map_tensor_name(name), data_torch) + + +@ModelBase.register("ArcticForCausalLM") +class ArcticModel(TextModel): + model_arch = gguf.MODEL_ARCH.ARCTIC + + def set_vocab(self): + # The reason for using a custom implementation here is that the + # snowflake-arctic-instruct model redefined tokens 31998 and 31999 from + # tokenizer.model and used them as BOS and EOS instead of adding new tokens. + from sentencepiece import SentencePieceProcessor + + tokenizer_path = self.dir_model / 'tokenizer.model' + + if not tokenizer_path.is_file(): + logger.error(f'Error: Missing {tokenizer_path}') + sys.exit(1) + + # Read the whole vocabulary from the tokenizer.model file + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + for token_id in range(tokenizer.vocab_size()): + + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + # Use the added_tokens_decoder field from tokeniser_config.json as the source + # of information about added/redefined tokens and modify them accordingly. + tokenizer_config_file = self.dir_model / 'tokenizer_config.json' + if tokenizer_config_file.is_file(): + with open(tokenizer_config_file, "r", encoding="utf-8") as f: + tokenizer_config_json = json.load(f) + + if "added_tokens_decoder" in tokenizer_config_json: + added_tokens_decoder = tokenizer_config_json["added_tokens_decoder"] + for token_id, token_json in added_tokens_decoder.items(): + token_id = int(token_id) + if token_id >= vocab_size: + logger.debug(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + + token_content = token_json["content"] + token_type = SentencePieceTokenTypes.USER_DEFINED + token_score = -10000.0 + + # Map unk_token to UNKNOWN, other special tokens to CONTROL + # Set the score to 0.0 as in the original tokenizer.model + if ("special" in token_json) and token_json["special"]: + if token_content == tokenizer_config_json["unk_token"]: + token_type = SentencePieceTokenTypes.UNKNOWN + else: + token_type = SentencePieceTokenTypes.CONTROL + token_score = 0.0 + + logger.info(f"Setting added token {token_id} to '{token_content}' (type: {token_type}, score: {token_score:.2f})") + tokens[token_id] = token_content.encode("utf-8") + toktypes[token_id] = token_type + scores[token_id] = token_score + + self.gguf_writer.add_tokenizer_model("llama") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_rope_dimension_count(hparams["hidden_size"] // hparams["num_attention_heads"]) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + + if name.endswith("q_proj.weight"): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith("k_proj.weight"): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + + # process the experts separately + if name.find("block_sparse_moe.experts") != -1: + n_experts = self.hparams["num_local_experts"] + + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for wid in ["w1", "w2", "w3"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.block_sparse_moe.experts.{xid}.{wid}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"layers.{bid}.feed_forward.experts.{wid}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("DeepseekForCausalLM") +class DeepseekModel(TextModel): + model_arch = gguf.MODEL_ARCH.DEEPSEEK + + def set_vocab(self): + try: + self._set_vocab_sentencepiece() + except FileNotFoundError: + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + if (rope_dim := hparams.get("head_dim")) is None: + rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"] + + self.gguf_writer.add_rope_dimension_count(rope_dim) + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_weights_scale(1.0) + self.gguf_writer.add_expert_count(hparams["n_routed_experts"]) + self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"]) + + _experts: list[dict[str, Tensor]] | None = None + + @staticmethod + def permute(weights: Tensor, n_head: int, n_head_kv: int | None): + if n_head_kv is not None and n_head != n_head_kv: + n_head = n_head_kv + return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = DeepseekModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = DeepseekModel.permute(data_torch, n_head, n_kv_head) + + # process the experts separately + if name.find("mlp.experts") != -1: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("DeepseekV2ForCausalLM") +@ModelBase.register("DeepseekV3ForCausalLM") +class DeepseekV2Model(TextModel): + model_arch = gguf.MODEL_ARCH.DEEPSEEK2 + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + + # note: deepseek2 using MLA converts into MQA (ie: GQA with 1 group) + self.hparams["num_key_value_heads"] = 1 + + super().set_gguf_parameters() + hparams = self.hparams + + self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + if "q_lora_rank" in hparams and hparams["q_lora_rank"] is not None: + self.gguf_writer.add_q_lora_rank(hparams["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) + + # note: deepseek2 using MLA converts into MQA with larger heads, then decompresses to MHA + self.gguf_writer.add_key_length(hparams["kv_lora_rank"] + hparams["qk_rope_head_dim"]) + self.gguf_writer.add_value_length(hparams["kv_lora_rank"]) + self.gguf_writer.add_key_length_mla(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]) + self.gguf_writer.add_value_length_mla(hparams["v_head_dim"]) + + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_count(hparams["n_routed_experts"]) + self.gguf_writer.add_expert_shared_count(hparams["n_shared_experts"]) + self.gguf_writer.add_expert_weights_scale(hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"]) + + if hparams["scoring_func"] == "sigmoid": + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + elif hparams["scoring_func"] == "softmax": + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX) + else: + raise ValueError(f"Unsupported scoring_func value: {hparams['scoring_func']}") + + self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) + + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + self.gguf_writer.add_rope_scaling_yarn_log_mul(0.1 * rope_scaling["mscale_all_dim"]) + + _experts: list[dict[str, Tensor]] | None = None + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # rename e_score_correction_bias tensors + if name.endswith("e_score_correction_bias"): + name = name.replace("e_score_correction_bias", "e_score_correction.bias") + + # skip Multi-Token Prediction (MTP) layers + block_count = self.hparams["num_hidden_layers"] + match = re.match(r"model.layers.(\d+)", name) + if match and int(match.group(1)) >= block_count: + return [] + + # process the experts separately + if name.find("mlp.experts") != -1: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + tensors: list[tuple[str, Tensor]] = [] + + # merge the experts into a single 3d tensor + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + return tensors + else: + return [] + + # note: MLA with the absorption optimization, needs these two split and k_b_proj transposed + if name.endswith("kv_b_proj.weight"): + name_kb = name.replace("kv_b_proj", "k_b_proj") + name_vb = name.replace("kv_b_proj", "v_b_proj") + + n_head_kv = self.hparams["num_key_value_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope_head_dim) + + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + k_b = k_b.transpose(1, 2) + + return [ + (self.map_tensor_name(name_kb), k_b), + (self.map_tensor_name(name_vb), v_b) + ] + + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("Dots1ForCausalLM") +class Dots1Model(Qwen2MoeModel): + model_arch = gguf.MODEL_ARCH.DOTS1 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.hparams["num_experts"] = self.hparams["n_routed_experts"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"]) + self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"]) + + if self.hparams["scoring_func"] == "noaux_tc": + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + else: + raise ValueError(f"Unsupported scoring_func value: {self.hparams['scoring_func']}") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): + if name.endswith("e_score_correction_bias"): + name = name.replace("e_score_correction_bias", "e_score_correction.bias") + if "shared_experts" in name: + return [(self.map_tensor_name(name), data_torch)] + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("PLMForCausalLM") +class PLMModel(TextModel): + model_arch = gguf.MODEL_ARCH.PLM + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_kv_lora_rank(hparams["kv_lora_rank"]) + self.gguf_writer.add_key_length(hparams["qk_nope_head_dim"] + hparams["qk_rope_head_dim"]) + self.gguf_writer.add_value_length(hparams["v_head_dim"]) + self.gguf_writer.add_rope_dimension_count(hparams["qk_rope_head_dim"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + return [(self.map_tensor_name(name), data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + +@ModelBase.register("T5WithLMHeadModel") +@ModelBase.register("T5ForConditionalGeneration") +@ModelBase.register("MT5ForConditionalGeneration") +@ModelBase.register("UMT5ForConditionalGeneration") +class T5Model(TextModel): + model_arch = gguf.MODEL_ARCH.T5 + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.shared_token_embeddings_found = False + + def set_vocab(self): + # to avoid TypeError: Descriptors cannot be created directly + # exception when importing sentencepiece_model_pb2 + os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" + from sentencepiece import SentencePieceProcessor + from sentencepiece import sentencepiece_model_pb2 as model + + tokenizer_path = self.dir_model / 'tokenizer.model' + + # many older models use spiece.model tokenizer model filename + if not tokenizer_path.is_file(): + tokenizer_path = self.dir_model / 'spiece.model' + + if not tokenizer_path.is_file(): + raise FileNotFoundError(f"File not found: {tokenizer_path}") + + sentencepiece_model = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] + sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read()) + + # some models like Pile-T5 family use BPE tokenizer instead of Unigram + if sentencepiece_model.trainer_spec.model_type == 2: # BPE + # assure the tokenizer model file name is correct + assert tokenizer_path.name == 'tokenizer.model' + return self._set_vocab_sentencepiece() + else: + assert sentencepiece_model.trainer_spec.model_type == 1 # UNIGRAM + + add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix + remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces + precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + for token_id in range(tokenizer.vocab_size()): + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + added_tokens_file = self.dir_model / 'added_tokens.json' + if added_tokens_file.is_file(): + with open(added_tokens_file, "r", encoding="utf-8") as f: + added_tokens_json = json.load(f) + for key in added_tokens_json: + token_id = added_tokens_json[key] + if token_id >= vocab_size: + logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + + tokens[token_id] = key.encode("utf-8") + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + + if vocab_size > len(tokens): + pad_count = vocab_size - len(tokens) + logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]") + for i in range(1, pad_count + 1): + tokens.append(bytes(f"[PAD{i}]", encoding="utf-8")) + scores.append(-1000.0) + toktypes.append(SentencePieceTokenTypes.UNUSED) + + self.gguf_writer.add_tokenizer_model("t5") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(add_prefix) + self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces) + if precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None: + logger.warning("Couldn't find context length in config.json, assuming default value of 512") + n_ctx = 512 + self.gguf_writer.add_context_length(n_ctx) + self.gguf_writer.add_embedding_length(self.hparams["d_model"]) + self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"]) + self.gguf_writer.add_block_count(self.hparams["num_layers"]) + self.gguf_writer.add_head_count(self.hparams["num_heads"]) + self.gguf_writer.add_key_length(self.hparams["d_kv"]) + self.gguf_writer.add_value_length(self.hparams["d_kv"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_decoder_start_token_id(self.hparams["decoder_start_token_id"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight", + # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored + # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder + # and decoder and ignore the remaining ones. + if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]: + if not self.shared_token_embeddings_found: + name = "shared.weight" + self.shared_token_embeddings_found = True + else: + logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.") + return [] + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("T5EncoderModel") +class T5EncoderModel(TextModel): + model_arch = gguf.MODEL_ARCH.T5ENCODER + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.shared_token_embeddings_found = False + + def set_vocab(self): + # to avoid TypeError: Descriptors cannot be created directly + # exception when importing sentencepiece_model_pb2 + os.environ["PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION"] = "python" + from sentencepiece import SentencePieceProcessor + from sentencepiece import sentencepiece_model_pb2 as model + + tokenizer_path = self.dir_model / 'tokenizer.model' + + # many older models use spiece.model tokenizer model filename + if not tokenizer_path.is_file(): + tokenizer_path = self.dir_model / 'spiece.model' + + if not tokenizer_path.is_file(): + raise FileNotFoundError(f"File not found: {tokenizer_path}") + + sentencepiece_model = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] + sentencepiece_model.ParseFromString(open(tokenizer_path, "rb").read()) + + # some models like Pile-T5 family use BPE tokenizer instead of Unigram + if sentencepiece_model.trainer_spec.model_type == 2: # BPE + # assure the tokenizer model file name is correct + assert tokenizer_path.name == 'tokenizer.model' + return self._set_vocab_sentencepiece() + else: + assert sentencepiece_model.trainer_spec.model_type == 1 # UNIGRAM + + add_prefix = sentencepiece_model.normalizer_spec.add_dummy_prefix + remove_whitespaces = sentencepiece_model.normalizer_spec.remove_extra_whitespaces + precompiled_charsmap = sentencepiece_model.normalizer_spec.precompiled_charsmap + + tokenizer = SentencePieceProcessor() + tokenizer.LoadFromFile(str(tokenizer_path)) + + vocab_size = self.hparams.get('vocab_size', tokenizer.vocab_size()) + + tokens: list[bytes] = [f"[PAD{i}]".encode("utf-8") for i in range(vocab_size)] + scores: list[float] = [-10000.0] * vocab_size + toktypes: list[int] = [SentencePieceTokenTypes.UNUSED] * vocab_size + + for token_id in range(tokenizer.vocab_size()): + piece = tokenizer.IdToPiece(token_id) + text = piece.encode("utf-8") + score = tokenizer.GetScore(token_id) + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.IsUnknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.IsControl(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.IsUnused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.IsByte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens[token_id] = text + scores[token_id] = score + toktypes[token_id] = toktype + + added_tokens_file = self.dir_model / 'added_tokens.json' + if added_tokens_file.is_file(): + with open(added_tokens_file, "r", encoding="utf-8") as f: + added_tokens_json = json.load(f) + for key in added_tokens_json: + token_id = added_tokens_json[key] + if token_id >= vocab_size: + logger.warning(f'ignore token {token_id}: id is out of range, max={vocab_size - 1}') + continue + + tokens[token_id] = key.encode("utf-8") + scores[token_id] = -1000.0 + toktypes[token_id] = SentencePieceTokenTypes.USER_DEFINED + + if vocab_size > len(tokens): + pad_count = vocab_size - len(tokens) + logger.debug(f"Padding vocab with {pad_count} token(s) - [PAD1] through [PAD{pad_count}]") + for i in range(1, pad_count + 1): + tokens.append(bytes(f"[PAD{i}]", encoding="utf-8")) + scores.append(-1000.0) + toktypes.append(SentencePieceTokenTypes.UNUSED) + + self.gguf_writer.add_tokenizer_model("t5") + self.gguf_writer.add_tokenizer_pre("default") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + self.gguf_writer.add_add_space_prefix(add_prefix) + self.gguf_writer.add_remove_extra_whitespaces(remove_whitespaces) + if precompiled_charsmap: + self.gguf_writer.add_precompiled_charsmap(precompiled_charsmap) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + if (n_ctx := self.find_hparam(["n_positions"], optional=True)) is None: + logger.warning("Couldn't find context length in config.json, assuming default value of 512") + n_ctx = 512 + self.gguf_writer.add_context_length(n_ctx) + self.gguf_writer.add_embedding_length(self.hparams["d_model"]) + self.gguf_writer.add_feed_forward_length(self.hparams["d_ff"]) + self.gguf_writer.add_block_count(self.hparams["num_layers"]) + self.gguf_writer.add_head_count(self.hparams["num_heads"]) + self.gguf_writer.add_key_length(self.hparams["d_kv"]) + self.gguf_writer.add_value_length(self.hparams["d_kv"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_relative_attn_buckets_count(self.hparams["relative_attention_num_buckets"]) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + # T5 based models contain shared token embeddings tensors saved randomly as either "encoder.embed_tokens.weight", + # "decoder.embed_tokens.weight" or "shared.weight" tensor. In some models there are even multiple of them stored + # in the safetensors files. We use the first tensor from these three as the token embeddings for both encoder + # and decoder and ignore the remaining ones. + if name in ["decoder.embed_tokens.weight", "encoder.embed_tokens.weight", "shared.weight"]: + if not self.shared_token_embeddings_found: + name = "shared.weight" + self.shared_token_embeddings_found = True + else: + logger.debug(f"Skipping shared tensor {name!r} in safetensors so that convert can end normally.") + return [] + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("JAISLMHeadModel") +class JaisModel(TextModel): + model_arch = gguf.MODEL_ARCH.JAIS + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # SwigLU activation + assert self.hparams["activation_function"] == "swiglu" + # ALiBi position embedding + assert self.hparams["position_embedding_type"] == "alibi" + + # Embeddings scale + self.embeddings_scale = 1.0 + if 'mup_embeddings_scale' in self.hparams: + self.embeddings_scale = self.hparams['mup_embeddings_scale'] + elif 'embeddings_scale' in self.hparams: + self.embeddings_scale = self.hparams['embeddings_scale'] + else: + assert False + + self.width_scale = 1.0 + if 'mup_output_alpha' in self.hparams: + assert 'mup_width_scale' in self.hparams + self.width_scale = self.hparams['mup_output_alpha'] * self.hparams['mup_width_scale'] + elif 'width_scale' in self.hparams: + self.width_scale = self.hparams['width_scale'] + else: + assert False + + self.max_alibi_bias = 8.0 + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + self.gguf_writer.add_block_count(self.hparams["n_layer"]) + self.gguf_writer.add_context_length(self.hparams["n_positions"]) + self.gguf_writer.add_embedding_length(self.hparams["n_embd"]) + self.gguf_writer.add_feed_forward_length(self.hparams["n_inner"]) + self.gguf_writer.add_head_count(self.hparams["n_head"]) + self.gguf_writer.add_layer_norm_eps(self.hparams["layer_norm_epsilon"]) + self.gguf_writer.add_file_type(self.ftype) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + tensors: list[tuple[str, Tensor]] = [] + + # we don't need these + if name.endswith((".attn.bias")): + return tensors + + if name.endswith(("relative_pe.slopes")): + # Calculate max ALiBi bias (this is the inverse of the ALiBi calculation) + # Some other models has max_alibi_bias spelled out explicitly in the hyperparams, + # but Jais's PyTorch model simply precalculates the slope values and places them + # in relative_pes.slopes + n_head_closest_log2 = 2 ** math.floor(math.log2(self.hparams["n_head"])) + first_val = float(data_torch[0].item()) + self.max_alibi_bias = -round(math.log2(first_val) * n_head_closest_log2) + + return tensors + + if name.endswith((".c_attn.weight", ".c_proj.weight", ".c_fc.weight", ".c_fc2.weight")): + data_torch = data_torch.transpose(1, 0) + + new_name = self.map_tensor_name(name) + + if new_name == self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD): + tensors.append((new_name, data_torch * self.embeddings_scale)) + elif new_name == self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT): + tensors.append((new_name, data_torch * self.width_scale)) + else: + tensors.append((new_name, data_torch)) + + return tensors + + def prepare_tensors(self): + super().prepare_tensors() + self.gguf_writer.add_max_alibi_bias(self.max_alibi_bias) + + +@ModelBase.register("Glm4ForCausalLM") +class Glm4Model(TextModel): + model_arch = gguf.MODEL_ARCH.GLM4 + + def set_vocab(self): + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"]) + special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"]) + special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"]) + special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"]) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + rope_dim = self.hparams["head_dim"] + self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.hparams.get("partial_rotary_factor", 0.5))) + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + + +@ModelBase.register("GlmForCausalLM", "ChatGLMModel", "ChatGLMForConditionalGeneration") +class ChatGLMModel(TextModel): + model_arch = gguf.MODEL_ARCH.CHATGLM + + def set_vocab_chatglm3(self): + dir_model = self.dir_model + hparams = self.hparams + tokens: list[bytes] = [] + toktypes: list[int] = [] + scores: list[float] = [] + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True) + vocab_size = hparams.get("padded_vocab_size", len(tokenizer.get_vocab())) + assert max(tokenizer.get_vocab().values()) < vocab_size + role_special_tokens = ["<|system|>", "<|user|>", "<|assistant|>", "<|observation|>"] + special_tokens = ["[MASK]", "[gMASK]", "[sMASK]", "sop", "eop"] + role_special_tokens + for token_id in range(vocab_size): + piece = tokenizer._convert_id_to_token(token_id) + if token_id == 0: + piece = "" + elif token_id == 1: + piece = "" + elif token_id == 2: + piece = "" + + text = piece.encode("utf-8") + score = 0.0 + # Referencing the tokenizer Python implementation(https://huggingface.co/THUDM/chatglm3-6b/blob/main/tokenization_chatglm.py), + # it is only valid if it is less than tokenizer.tokenizer.sp_model.vocab_size() + if len(piece) != 0 and token_id < tokenizer.tokenizer.sp_model.vocab_size(): + score = tokenizer.tokenizer.sp_model.get_score(token_id) + + if token_id >= tokenizer.tokenizer.sp_model.vocab_size(): + if piece in special_tokens: + toktype = SentencePieceTokenTypes.CONTROL + elif len(piece) == 0: + text = f"[PAD{token_id}]".encode("utf-8") + toktype = SentencePieceTokenTypes.UNUSED + else: + toktype = SentencePieceTokenTypes.USER_DEFINED + tokens.append(text) + scores.append(score) + toktypes.append(toktype) + continue + + toktype = SentencePieceTokenTypes.NORMAL + if tokenizer.tokenizer.sp_model.is_unknown(token_id): + toktype = SentencePieceTokenTypes.UNKNOWN + elif tokenizer.tokenizer.sp_model.is_control(token_id): + toktype = SentencePieceTokenTypes.CONTROL + elif tokenizer.tokenizer.sp_model.is_unused(token_id): + toktype = SentencePieceTokenTypes.UNUSED + elif tokenizer.tokenizer.sp_model.is_byte(token_id): + toktype = SentencePieceTokenTypes.BYTE + + tokens.append(text) + scores.append(score) + toktypes.append(toktype) + + self.gguf_writer.add_tokenizer_model("llama") + # glm3 needs prefix and suffix formatted as: + # prompt = "[gMASK]sop<|user|>\n" + prompt + "<|assistant|>" + self.gguf_writer.add_tokenizer_pre("chatglm-spm") + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_scores(scores) + self.gguf_writer.add_token_types(toktypes) + + special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens)) + special_vocab.add_to_gguf(self.gguf_writer) + + @staticmethod + def token_bytes_to_string(b): + from transformers.models.gpt2.tokenization_gpt2 import bytes_to_unicode + byte_encoder = bytes_to_unicode() + return ''.join([byte_encoder[ord(char)] for char in b.decode('latin-1')]) + + @staticmethod + def bpe(mergeable_ranks: dict[bytes, int], token: bytes, max_rank: int | None = None) -> list[bytes]: + parts = [bytes([b]) for b in token] + while True: + min_idx = None + min_rank = None + for i, pair in enumerate(zip(parts[:-1], parts[1:])): + rank = mergeable_ranks.get(pair[0] + pair[1]) + if rank is not None and (min_rank is None or rank < min_rank): + min_idx = i + min_rank = rank + if min_rank is None or (max_rank is not None and min_rank >= max_rank): + break + assert min_idx is not None + parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx + 1]] + parts[min_idx + 2:] + return parts + + def set_vocab(self): + if "THUDM/chatglm3-6b" in self.hparams.get("_name_or_path", ""): + self.set_vocab_chatglm3() + return + + dir_model = self.dir_model + hparams = self.hparams + tokens: list[str] = [] + toktypes: list[int] = [] + + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(dir_model, trust_remote_code=True) + vocab_size = hparams.get("padded_vocab_size",hparams["vocab_size"]) + assert max(tokenizer.get_vocab().values()) < vocab_size + + tokens, toktypes, tokpre = self.get_vocab_base() + self.gguf_writer.add_tokenizer_model("gpt2") + self.gguf_writer.add_tokenizer_pre(tokpre) + self.gguf_writer.add_token_list(tokens) + self.gguf_writer.add_token_types(toktypes) + special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) + # only add special tokens when they were not already loaded from config.json + special_vocab._set_special_token("eos", tokenizer.get_added_vocab()["<|endoftext|>"]) + special_vocab._set_special_token("eot", tokenizer.get_added_vocab()["<|user|>"]) + # this one is usually not in config.json anyway + special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<|endoftext|>"]) + special_vocab.add_to_gguf(self.gguf_writer) + + def set_gguf_parameters(self): + n_embed = self.hparams.get("hidden_size", self.hparams.get("n_embed")) + n_head = self.hparams.get("n_head", self.hparams.get("num_attention_heads")) + n_head_kv = self.hparams.get("multi_query_group_num", self.hparams.get("num_key_value_heads", n_head)) + self.gguf_writer.add_context_length(self.hparams.get("seq_length", n_embed)) + self.gguf_writer.add_embedding_length(n_embed) + self.gguf_writer.add_feed_forward_length(self.hparams.get("ffn_hidden_size", self.hparams.get("intermediate_size", 4 * n_embed))) + self.gguf_writer.add_block_count(self.hparams.get("num_layers", self.hparams["num_hidden_layers"])) + self.gguf_writer.add_head_count(n_head) + self.gguf_writer.add_head_count_kv(n_head_kv) + self.gguf_writer.add_layer_norm_rms_eps(self.hparams.get("layernorm_epsilon",1e-5)) + self.gguf_writer.add_file_type(self.ftype) + if "attention_dim" in self.hparams: + rope_dim = self.hparams["attention_dim"] + else: + rope_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_count(int(rope_dim * self.hparams.get("partial_rotary_factor", 0.5))) + self.gguf_writer.add_add_bos_token(False) + rope_freq = 10000 + if "rope_ratio" in self.hparams: + rope_freq = rope_freq * self.hparams["rope_ratio"] + self.gguf_writer.add_rope_freq_base(rope_freq) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if name.endswith(".rotary_pos_emb.inv_freq") or name.startswith("model.vision."): + return [] + + name = name.removeprefix("transformer.") + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("NemotronForCausalLM") +class NemotronModel(TextModel): + model_arch = gguf.MODEL_ARCH.NEMOTRON + + def set_vocab(self): + self._set_vocab_sentencepiece() + self.gguf_writer.add_pad_token_id(0) + self.gguf_writer.add_unk_token_id(1) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + + f_norm_eps = self.find_hparam(["layer_norm_eps", "layer_norm_epsilon", "norm_epsilon", "norm_eps"]) + self.gguf_writer.add_layer_norm_eps(f_norm_eps) + + # * Partial RoPE + rot_pct = self.find_hparam(["partial_rotary_factor", "rope_pct", "rope_percent"]) + n_embd = self.find_hparam(["hidden_size", "n_embd"]) + n_head = self.find_hparam(["num_attention_heads", "n_head"]) + self.gguf_writer.add_rope_dimension_count(int(rot_pct * n_embd) // n_head) + + # * RopeScaling for Nemotron + if "rope_scaling" not in self.hparams or self.hparams["rope_scaling"] is None: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + else: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(self.hparams["factor"]) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # * Adding +1 to LayerNorm's weights here to implement layernorm1p w/o changing anything on the GGML engine side + # model.layers.{l}.input_layernorm.weight + # model.layers.{l}.post_attention_layernorm.weight + # model.norm.weight + if name.endswith("norm.weight"): + data_torch = data_torch + 1 + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("ExaoneForCausalLM") +class ExaoneModel(TextModel): + model_arch = gguf.MODEL_ARCH.EXAONE + + def set_gguf_parameters(self): + hparams = self.hparams + + assert (hparams["activation_function"] == "silu") + + max_position_embeddings = hparams["max_position_embeddings"] + embed_dim = hparams["hidden_size"] + num_heads = hparams["num_attention_heads"] + num_kv_heads = hparams.get("num_key_value_heads", num_heads) + layer_norm_eps = hparams["layer_norm_epsilon"] + intermediate_size = hparams["intermediate_size"] if "intermediate_size" in hparams else 4 * embed_dim + num_layers = hparams["num_layers"] + # ignore for now as EXAONE-3.0-7.8B-Instruct attentino_dropout is 0.0 + # attention_dropout_rate = hparams["attention_dropout"] + # ignore for now as EXAONE-3.0-7.8B-Instruct embed_dropout is 0.0 + # embed_dropout_rate = hparams["embed_dropout"] + self.gguf_writer.add_embedding_length(embed_dim) + self.gguf_writer.add_head_count(num_heads) + self.gguf_writer.add_head_count_kv(num_kv_heads) + self.gguf_writer.add_context_length(max_position_embeddings) + self.gguf_writer.add_layer_norm_rms_eps(layer_norm_eps) + self.gguf_writer.add_feed_forward_length(intermediate_size) + self.gguf_writer.add_block_count(num_layers) + self.gguf_writer.add_file_type(self.ftype) + + if (rope_theta := self.hparams.get("rope_theta")) is not None: + self.gguf_writer.add_rope_freq_base(rope_theta) + rotary_factor = self.find_hparam(["partial_rotary_factor", "rope_pct"], optional=True) + rotary_factor = rotary_factor if rotary_factor is not None else 1.0 + self.gguf_writer.add_rope_dimension_count(int(rotary_factor * (hparams["hidden_size"] // hparams["num_attention_heads"]))) + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "linear" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.LINEAR) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + if rope_scaling := self.find_hparam(["rope_scaling"], optional=True): + if rope_scaling.get("rope_type", '').lower() == "llama3": + base = self.hparams.get("rope_theta", 10000.0) + if (dim := self.hparams.get("head_dim")) is None: + dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + + factor = rope_scaling.get("factor", 8.0) + low_freq_factor = rope_scaling.get("low_freq_factor", 1.0) + high_freq_factor = rope_scaling.get("high_freq_factor", 4.0) + old_context_len = self.hparams.get("original_max_position_embeddings", 8192) + + low_freq_wavelen = old_context_len / low_freq_factor + high_freq_wavelen = old_context_len / high_freq_factor + assert low_freq_wavelen != high_freq_wavelen + + rope_factors = [] + for freq in freqs: + wavelen = 2 * math.pi / freq + if wavelen < high_freq_wavelen: + rope_factors.append(1) + elif wavelen > low_freq_wavelen: + rope_factors.append(factor) + else: + smooth = (old_context_len / wavelen - low_freq_factor) / (high_freq_factor - low_freq_factor) + rope_factors.append(1 / ((1 - smooth) / factor + smooth)) + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.tensor(rope_factors, dtype=torch.float32)) + + +@ModelBase.register("GraniteForCausalLM") +class GraniteModel(LlamaModel): + """Conversion for IBM's GraniteForCausalLM""" + model_arch = gguf.MODEL_ARCH.GRANITE + + def set_gguf_parameters(self): + """Granite uses standard llama parameters with the following differences: + + - No head_dim support + - New multiplier params: + - attention_scale + - embedding_scale + - residual_scale + - logits_scaling + """ + if head_dim := self.hparams.pop("head_dim", None): + logger.warning("Ignoring head_dim (%s) from config for Granite", head_dim) + super().set_gguf_parameters() + # NOTE: Convert _multiplier params to _scale params for naming + # consistency + if attention_scale := self.hparams.get("attention_multiplier"): + self.gguf_writer.add_attention_scale(attention_scale) + logger.info("gguf: (granite) attention_scale = %s", attention_scale) + if embedding_scale := self.hparams.get("embedding_multiplier"): + self.gguf_writer.add_embedding_scale(embedding_scale) + logger.info("gguf: (granite) embedding_scale = %s", embedding_scale) + if residual_scale := self.hparams.get("residual_multiplier"): + self.gguf_writer.add_residual_scale(residual_scale) + logger.info("gguf: (granite) residual_scale = %s", residual_scale) + if logits_scale := self.hparams.get("logits_scaling"): + self.gguf_writer.add_logit_scale(logits_scale) + logger.info("gguf: (granite) logits_scale = %s", logits_scale) + + +@ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM") +class GraniteMoeModel(GraniteModel): + """Conversion for IBM's GraniteMoeForCausalLM""" + model_arch = gguf.MODEL_ARCH.GRANITE_MOE + + def set_gguf_parameters(self): + """GraniteMoeShared uses GraniteMoe parameters plus the following: + - shared_intermediate_size + """ + super().set_gguf_parameters() + if shared_feed_forward_length := self.hparams.get("shared_intermediate_size"): + self.gguf_writer.add_expert_shared_feed_forward_length(shared_feed_forward_length) + logger.info("gguf: (granitemoeshared) shared_feed_forward_length = %s", shared_feed_forward_length) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + """In modeling_granitemoe, the JetMoe implementation of parallel experts + is used. This essentially merges w1 and w3 into a single tensor with 2x + the hidden size that is then split during forward. To keep compatibility + with existing mixtral support, we pull them apart here. + """ + + if name.endswith("block_sparse_moe.input_linear.weight"): + ffn_dim = self.hparams["intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size" + gate, up = data_torch.split(ffn_dim, dim=-2) + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_EXP, bid), gate), + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_EXP, bid), up), + ] + + if name.endswith("shared_mlp.input_linear.weight"): + ffn_dim = self.hparams["shared_intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size" + gate, up = data_torch.split(ffn_dim, dim=-2) + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_SHEXP, bid), gate), + (self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), up), + ] + + return super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("BailingMoeForCausalLM") +class BailingMoeModel(TextModel): + model_arch = gguf.MODEL_ARCH.BAILINGMOE + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + if (rope_dim := hparams.get("head_dim")) is None: + rope_dim = hparams["hidden_size"] // hparams["num_attention_heads"] + + self.gguf_writer.add_rope_dimension_count(rope_dim) + rope_scaling = self.hparams.get("rope_scaling") or {} + if rope_scaling.get("rope_type", rope_scaling.get("type")) == "yarn" and "factor" in rope_scaling: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.YARN) + self.gguf_writer.add_rope_scaling_factor(rope_scaling["factor"]) + self.gguf_writer.add_rope_scaling_orig_ctx_len(rope_scaling["original_max_position_embeddings"]) + else: + self.gguf_writer.add_rope_scaling_type(gguf.RopeScalingType.NONE) + self.gguf_writer.add_leading_dense_block_count(hparams["first_k_dense_replace"]) + self.gguf_writer.add_vocab_size(hparams["vocab_size"]) + self.gguf_writer.add_expert_feed_forward_length(hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_weights_scale(1.0) + self.gguf_writer.add_expert_count(hparams["num_experts"]) + self.gguf_writer.add_expert_shared_count(hparams["num_shared_experts"]) + self.gguf_writer.add_expert_weights_norm(hparams["norm_topk_prob"]) + + _experts: list[dict[str, Tensor]] | None = None + + @staticmethod + def permute(weights: Tensor, n_head: int, n_head_kv: int | None): + if n_head_kv is not None and n_head != n_head_kv: + n_head = n_head_kv + return (weights.reshape(n_head, 2, weights.shape[0] // n_head // 2, *weights.shape[1:]) + .swapaxes(1, 2) + .reshape(weights.shape)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + n_embd = self.hparams["hidden_size"] + if (head_dim := self.hparams.get("head_dim")) is None: + head_dim = n_embd // n_head + + output_name = self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT) + + if name.endswith("attention.dense.weight"): + return [(self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_OUT, bid), data_torch)] + elif name.endswith("query_key_value.weight"): + q, k, v = data_torch.split([n_head * head_dim, n_kv_head * head_dim, n_kv_head * head_dim], dim=-2) + + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), BailingMoeModel.permute(q, n_head, n_head)), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), BailingMoeModel.permute(k, n_head, n_kv_head)), + (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v) + ] + elif name.find("mlp.experts") != -1: + n_experts = self.hparams["num_experts"] + assert bid is not None + + tensors: list[tuple[str, Tensor]] = [] + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) >= n_experts * 3: + # merge the experts into a single 3d tensor + for w_name in ["down_proj", "gate_proj", "up_proj"]: + datas: list[Tensor] = [] + + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + + data_torch = torch.stack(datas, dim=0) + + merged_name = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + + new_name = self.map_tensor_name(merged_name) + + tensors.append((new_name, data_torch)) + + return tensors + + new_name = self.map_tensor_name(name) + + if new_name == output_name and self.hparams.get("norm_head"): + data_torch = data_torch.float() + data_torch /= torch.norm(data_torch, p=2, dim=0, keepdim=True) + 1e-7 + + return [(new_name, data_torch)] + + def prepare_tensors(self): + super().prepare_tensors() + + if self._experts is not None: + # flatten `list[dict[str, Tensor]]` into `list[str]` + experts = [k for d in self._experts for k in d.keys()] + if len(experts) > 0: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("ChameleonForConditionalGeneration") +@ModelBase.register("ChameleonForCausalLM") # obsolete +class ChameleonModel(TextModel): + model_arch = gguf.MODEL_ARCH.CHAMELEON + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_swin_norm(self.hparams.get("swin_norm", False)) + + def set_vocab(self): + self._set_vocab_gpt2() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # ignore image tokenizer for now + # TODO: remove this once image support is implemented for Chameleon + if name.startswith("model.vqmodel"): + return [] + + n_head = self.hparams["num_attention_heads"] + n_kv_head = self.hparams.get("num_key_value_heads") + hidden_dim = self.hparams.get("hidden_size") + + if name.endswith(("q_proj.weight", "q_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + if name.endswith(("k_proj.weight", "k_proj.bias")): + data_torch = LlamaModel.permute(data_torch, n_head, n_kv_head) + if name.endswith(("q_norm.weight", "q_norm.bias")): + data_torch = ChameleonModel._reverse_hf_permute(data_torch, n_head, hidden_dim) + if name.endswith(("k_norm.weight", "k_norm.bias")): + data_torch = ChameleonModel._reverse_hf_permute(data_torch, n_kv_head, hidden_dim) + + return [(self.map_tensor_name(name), data_torch)] + + # see: https://github.com/huggingface/transformers/blob/72fb02c47dbbe1999ae105319f24631cad6e2e00/src/transformers/models/chameleon/convert_chameleon_weights_to_hf.py#L176-L203 + @staticmethod + def _reverse_hf_permute(data_torch, n_heads, hidden_dim): + head_dim = hidden_dim // n_heads + data_torch = data_torch[0].view(2, head_dim // 2).t().reshape(1, -1) + data_torch = data_torch.repeat_interleave(n_heads, 0) + return data_torch + + +@ModelBase.register("UltravoxModel") +class UltravoxModel(TextModel): + model_arch = gguf.MODEL_ARCH.LLAMA # dummy + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + raise NotImplementedError("Ultravox does not have text decoder. Instead, it uses Llama or other models for text. If you want to get the audio encoder, please use --mmproj argument") + + +@ModelBase.register("Qwen2AudioForConditionalGeneration") +class WhisperEncoderModel(MmprojModel): + has_vision_encoder = False # no vision encoder + has_audio_encoder = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.hparams["hidden_size"] = self.hparams["d_model"] + self.hparams["intermediate_size"] = self.hparams["encoder_ffn_dim"] + self.hparams["num_attention_heads"] = self.hparams["encoder_attention_heads"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.QWEN2A) + self.gguf_writer.add_audio_num_mel_bins(self.hparams["num_mel_bins"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams.get("layer_norm_eps", 1e-5)) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + del bid, new_name, n_dims # unused + if ".conv" in name and ".weight" in name: + return gguf.GGMLQuantizationType.F16 + return False + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + del bid # unused + + if name.startswith("language_model."): + # skip language model tensors + return [] + + # prevent clash naming with vision tensors + if name.startswith("multi_modal_projector"): + name = "audio." + name + + if "conv1.bias" in name or "conv2.bias" in name: + # transpose conv1 and conv2 bias + data_torch = data_torch.unsqueeze(-1) + + return [(self.map_tensor_name(name), data_torch)] + + +@ModelBase.register("UltravoxModel") +class UltravoxWhisperEncoderModel(WhisperEncoderModel): + has_vision_encoder = False # no vision encoder + has_audio_encoder = True + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_audio_stack_factor(self.global_config["stack_factor"]) + +###### CONVERSION LOGIC ###### + + +# tree of lazy tensors +class LazyTorchTensor(gguf.LazyBase): + _tensor_type = torch.Tensor + # to keep the type-checker happy + dtype: torch.dtype + shape: torch.Size + + # only used when converting a torch.Tensor to a np.ndarray + _dtype_map: dict[torch.dtype, type] = { + torch.float16: np.float16, + torch.float32: np.float32, + } + + # used for safetensors slices + # ref: https://github.com/huggingface/safetensors/blob/079781fd0dc455ba0fe851e2b4507c33d0c0d407/bindings/python/src/lib.rs#L1046 + # TODO: uncomment U64, U32, and U16, ref: https://github.com/pytorch/pytorch/issues/58734 + _dtype_str_map: dict[str, torch.dtype] = { + "F64": torch.float64, + "F32": torch.float32, + "BF16": torch.bfloat16, + "F16": torch.float16, + # "U64": torch.uint64, + "I64": torch.int64, + # "U32": torch.uint32, + "I32": torch.int32, + # "U16": torch.uint16, + "I16": torch.int16, + "U8": torch.uint8, + "I8": torch.int8, + "BOOL": torch.bool, + "F8_E4M3": torch.float8_e4m3fn, + "F8_E5M2": torch.float8_e5m2, + } + + def numpy(self) -> gguf.LazyNumpyTensor: + dtype = self._dtype_map[self.dtype] + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(dtype, self.shape), + args=(self,), + func=(lambda s: s.numpy()) + ) + + @classmethod + def meta_with_dtype_and_shape(cls, dtype: torch.dtype, shape: tuple[int, ...]) -> Tensor: + return torch.empty(size=shape, dtype=dtype, device="meta") + + @classmethod + def from_safetensors_slice(cls, st_slice: Any) -> Tensor: + dtype = cls._dtype_str_map[st_slice.get_dtype()] + shape: tuple[int, ...] = tuple(st_slice.get_shape()) + lazy = cls(meta=cls.meta_with_dtype_and_shape(dtype, shape), args=(st_slice,), func=lambda s: s[:]) + return cast(torch.Tensor, lazy) + + @classmethod + def from_remote_tensor(cls, remote_tensor: gguf.utility.RemoteTensor): + dtype = cls._dtype_str_map[remote_tensor.dtype] + shape = remote_tensor.shape + meta = cls.meta_with_dtype_and_shape(dtype, shape) + lazy = cls(meta=meta, args=(remote_tensor,), func=lambda r: torch.frombuffer(r.data(), dtype=dtype).reshape(shape)) + return cast(torch.Tensor, lazy) + + @classmethod + def __torch_function__(cls, func, types, args=(), kwargs=None): + del types # unused + + if kwargs is None: + kwargs = {} + + if func is torch.Tensor.numpy: + return args[0].numpy() + + return cls._wrap_fn(func)(*args, **kwargs) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Convert a huggingface model to a GGML compatible file") + parser.add_argument( + "--vocab-only", action="store_true", + help="extract only the vocab", + ) + parser.add_argument( + "--outfile", type=Path, + help="path to write to; default: based on input. {ftype} will be replaced by the outtype.", + ) + parser.add_argument( + "--outtype", type=str, choices=["f32", "f16", "bf16", "q8_0", "tq1_0", "tq2_0", "auto"], default="f16", + help="output format - use f32 for float32, f16 for float16, bf16 for bfloat16, q8_0 for Q8_0, tq1_0 or tq2_0 for ternary, and auto for the highest-fidelity 16-bit float type depending on the first loaded tensor type", + ) + parser.add_argument( + "--bigendian", action="store_true", + help="model is executed on big endian machine", + ) + parser.add_argument( + "model", type=str, + help="directory containing model file or huggingface repository ID (if --remote)", + nargs="?", + ) + parser.add_argument( + "--use-temp-file", action="store_true", + help="use the tempfile library while processing (helpful when running out of memory, process killed)", + ) + parser.add_argument( + "--no-lazy", action="store_true", + help="use more RAM by computing all outputs before writing (use in case lazy evaluation is broken)", + ) + parser.add_argument( + "--model-name", type=str, default=None, + help="name of the model", + ) + parser.add_argument( + "--verbose", action="store_true", + help="increase output verbosity", + ) + parser.add_argument( + "--split-max-tensors", type=int, default=0, + help="max tensors in each split", + ) + parser.add_argument( + "--split-max-size", type=str, default="0", + help="max size per split N(M|G)", + ) + parser.add_argument( + "--dry-run", action="store_true", + help="only print out a split plan and exit, without writing any new files", + ) + parser.add_argument( + "--no-tensor-first-split", action="store_true", + help="do not add tensors to the first split (disabled by default)" + ) + parser.add_argument( + "--metadata", type=Path, + help="Specify the path for an authorship metadata override file" + ) + parser.add_argument( + "--print-supported-models", action="store_true", + help="Print the supported models" + ) + parser.add_argument( + "--remote", action="store_true", + help="(Experimental) Read safetensors file remotely without downloading to disk. Config and tokenizer files will still be downloaded. To use this feature, you need to specify Hugging Face model repo name instead of a local directory. For example: 'HuggingFaceTB/SmolLM2-1.7B-Instruct'. Note: To access gated repo, set HF_TOKEN environment variable to your Hugging Face token.", + ) + parser.add_argument( + "--mmproj", action="store_true", + help="(Experimental) Export multimodal projector (mmproj) for vision models. This will only work on some vision models. A prefix 'mmproj-' will be added to the output file name.", + ) + + args = parser.parse_args() + if not args.print_supported_models and args.model is None: + parser.error("the following arguments are required: model") + return args + + +def split_str_to_n_bytes(split_str: str) -> int: + if split_str.endswith("K"): + n = int(split_str[:-1]) * 1000 + elif split_str.endswith("M"): + n = int(split_str[:-1]) * 1000 * 1000 + elif split_str.endswith("G"): + n = int(split_str[:-1]) * 1000 * 1000 * 1000 + elif split_str.isnumeric(): + n = int(split_str) + else: + raise ValueError(f"Invalid split size: {split_str}, must be a number, optionally followed by K, M, or G") + + if n < 0: + raise ValueError(f"Invalid split size: {split_str}, must be positive") + + return n + + +def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> str: + # TODO @ngxson : this won't work correctly if the model has both audio & vision encoders + # maybe we should fallback to text model's arch in that case, since not many models have both + text_config = hparams.get("text_config", {}) + vision_config = hparams.get("vision_config", {}) + arch = None + if (arches := hparams.get("architectures")) is not None and len(arches) > 0: + arch = arches[0] + elif "ssm_cfg" in hparams: + # For non-hf Mamba and Mamba2 models + arch = hparams["ssm_cfg"].get("layer", "Mamba") + "ForCausalLM" + + # if "architectures" is found in the sub-config, use that instead + if model_type == ModelType.TEXT and text_config.get("architectures") is not None: + arch = text_config["architectures"][0] + elif model_type == ModelType.MMPROJ and vision_config.get("architectures") is not None: + arch = vision_config["architectures"][0] + if arch is None: + raise ValueError("Failed to detect model architecture") + return arch + + +def main() -> None: + args = parse_args() + + if args.print_supported_models: + logger.error("Supported models:") + ModelBase.print_registered_models() + sys.exit(0) + + if args.verbose: + logging.basicConfig(level=logging.DEBUG) + else: + logging.basicConfig(level=logging.INFO) + + if args.remote: + hf_repo_id = args.model + from huggingface_hub import snapshot_download + local_dir = snapshot_download( + repo_id=hf_repo_id, + allow_patterns=["LICENSE", "*.json", "*.md", "*.txt", "tokenizer.model"]) + dir_model = Path(local_dir) + logger.info(f"Downloaded config and tokenizer to {local_dir}") + else: + hf_repo_id = None + dir_model = Path(args.model) + + if not dir_model.is_dir(): + logger.error(f'Error: {dir_model} is not a directory') + sys.exit(1) + + ftype_map: dict[str, gguf.LlamaFileType] = { + "f32": gguf.LlamaFileType.ALL_F32, + "f16": gguf.LlamaFileType.MOSTLY_F16, + "bf16": gguf.LlamaFileType.MOSTLY_BF16, + "q8_0": gguf.LlamaFileType.MOSTLY_Q8_0, + "tq1_0": gguf.LlamaFileType.MOSTLY_TQ1_0, + "tq2_0": gguf.LlamaFileType.MOSTLY_TQ2_0, + "auto": gguf.LlamaFileType.GUESSED, + } + + is_split = args.split_max_tensors > 0 or args.split_max_size != "0" + if args.use_temp_file and is_split: + logger.error("Error: Cannot use temp file when splitting") + sys.exit(1) + + if args.outfile is not None: + fname_out = args.outfile + elif hf_repo_id: + # if remote, use the model ID as the output file name + fname_out = Path("./" + hf_repo_id.replace("/", "-") + "-{ftype}.gguf") + else: + fname_out = dir_model + + logger.info(f"Loading model: {dir_model.name}") + + if args.mmproj: + if "mmproj" not in fname_out.name: + fname_out = ModelBase.add_prefix_to_filename(fname_out, "mmproj-") + + with torch.inference_mode(): + output_type = ftype_map[args.outtype] + model_type = ModelType.MMPROJ if args.mmproj else ModelType.TEXT + hparams = ModelBase.load_hparams(dir_model) + model_architecture = get_model_architecture(hparams, model_type) + logger.info(f"Model architecture: {model_architecture}") + try: + model_class = ModelBase.from_model_architecture(model_architecture, model_type=model_type) + except NotImplementedError: + logger.error(f"Model {model_architecture} is not supported") + sys.exit(1) + + model_instance = model_class(dir_model, output_type, fname_out, + is_big_endian=args.bigendian, use_temp_file=args.use_temp_file, + eager=args.no_lazy, + metadata_override=args.metadata, model_name=args.model_name, + split_max_tensors=args.split_max_tensors, + split_max_size=split_str_to_n_bytes(args.split_max_size), dry_run=args.dry_run, + small_first_shard=args.no_tensor_first_split, + remote_hf_model_id=hf_repo_id) + + if args.vocab_only: + logger.info("Exporting model vocab...") + model_instance.write_vocab() + logger.info(f"Model vocab successfully exported to {model_instance.fname_out}") + else: + logger.info("Exporting model...") + model_instance.write() + out_path = f"{model_instance.fname_out.parent}{os.sep}" if is_split else model_instance.fname_out + logger.info(f"Model successfully exported to {out_path}") + + +if __name__ == '__main__': + main() diff --git a/src/xpu/flamegraph/qwen3.cu/extract_v_m.py b/src/xpu/flamegraph/qwen3.cu/extract_v_m.py new file mode 100644 index 0000000..c7c05e3 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/extract_v_m.py @@ -0,0 +1,143 @@ +#modified from a utility in llama.cpp +#!/usr/bin/env python3 +import string +import logging +import sys +from pathlib import Path + +logger = logging.getLogger("reader") + +# Necessary to load the local gguf package +# sys.path.insert(0, str(Path(__file__).parent.parent)) + +from gguf.gguf_reader import GGUFReader + +# 된다 +def extract_merges_to_txt(reader, output_file="merges.txt"): + parts = reader.fields["tokenizer.ggml.merges"].parts + + # Skip initial header/metadata parts + start_idx = 6 + + # Crop to full merge pairs only + if (len(parts) - start_idx) % 2 != 0: + print(f"Merges field has odd number of parts after header. Truncating last.") + parts = parts[:len(parts) - 1] + + with open(output_file, "w", encoding="utf-8") as f: + for i in range(start_idx, len(parts), 2): + merge_bytes = parts[i] + try: + merge_str = bytes(merge_bytes).decode("utf-8") + except Exception: + merge_str = bytes(merge_bytes).decode("utf-8", errors="replace") + f.write(merge_str + "\n") + + print(f"Extracted {((len(parts) - start_idx) //2)} merges to {output_file}") + + +def extract_vocab_to_txt(reader, output_file="vocab.txt"): + tokens = reader.fields["tokenizer.ggml.tokens"].parts + with open(output_file, "w", encoding="utf-8") as f: + # Start at 6 (where real tokens start) + for i in range(6, len(tokens), 2): + token_bytes = tokens[i] + # Only process tokens that are arrays of uint8 + if getattr(token_bytes, 'dtype', None) == 'uint8': + b = bytes(token_bytes) + b = b.rstrip(b'\x00') + if b: # skip empty + try: + token_str = b.decode("utf-8") + except Exception: + token_str = b.decode("utf-8", errors="replace") + f.write(token_str + "\n") + print(f"Extraction complete ({(len(tokens) -6) //2} tokens written).") + + +def read_gguf_file(gguf_file_path): + """ + Reads and prints key-value pairs and tensor information from a GGUF file in an improved format. + + Parameters: + - gguf_file_path: Path to the GGUF file. + """ + + reader = GGUFReader(gguf_file_path) + + extract_merges_to_txt(reader) + extract_vocab_to_txt(reader) + + # List all key-value pairs in a columnized format + print("Key-Value Pairs:") # noqa: NP100 + max_key_length = max(len(key) for key in reader.fields.keys()) + + for key, field in reader.fields.items(): + value = field.parts[field.data[0]] + print(f"{key:{max_key_length}} : {value}") + + try: + value1 = ''.join(chr(x) for x in value) # Convert [103, 112, 116, 50] to "gpt2" + print(f"{key:{max_key_length}} : {value1}") # Print key and value + except: + pass + #elif isinstance(value, bytes): + #value2 = value.tobytes().decode('utf-8') # If value is bytes, decode to string + #print(f"{key:{max_key_length}} : {value2}") # Print key and value + + + for key, field in reader.fields.items(): + value = field.parts[field.data[0]] + + # Try to convert to string if it looks like string data + if isinstance(value, list) and all(isinstance(x, int) for x in value): + # Try UTF-8 first, fallback to ASCII, else show the list + try: + value_str = bytes(value).decode('utf-8') + except (UnicodeDecodeError, ValueError, TypeError): + try: + if all(32 <= x <= 126 for x in value): # printable ASCII + value_str = ''.join(chr(x) for x in value) + else: + value_str = str(value) + except Exception: + value_str = str(value) + value = value_str + + elif isinstance(value, bytes): + try: + value = value.decode('utf-8') + except UnicodeDecodeError: + value = str(value) + + elif hasattr(value, 'tobytes'): # numpy ndarray/memmap/etc + try: + value = value.tobytes().decode('utf-8') + except UnicodeDecodeError: + value = repr(value) + # OR, for arrays: np.array2string(value) for small arrays + except Exception: + value = repr(value) + else: + value = str(value) + + print(f"{key:{max_key_length}} : {value}") + + # List all tensors + print("Tensors:") # noqa: NP100 + tensor_info_format = "{:<30} | Shape: {:<15} | Size: {:<12} | Quantization: {}" + print(tensor_info_format.format("Tensor Name", "Shape", "Size", "Quantization")) # noqa: NP100 + print("-" * 80) # noqa: NP100 + for tensor in reader.tensors: + shape_str = "x".join(map(str, tensor.shape)) + size_str = str(tensor.n_elements) + quantization_str = tensor.tensor_type.name + print(tensor_info_format.format(tensor.name, shape_str, size_str, quantization_str)) # noqa: NP100 + + +if __name__ == '__main__': + if len(sys.argv) < 2: + logger.info("Usage: reader.py ") + sys.exit(1) + gguf_file_path = sys.argv[1] + read_gguf_file(gguf_file_path) diff --git a/src/xpu/flamegraph/qwen3.cu/header.py b/src/xpu/flamegraph/qwen3.cu/header.py new file mode 100644 index 0000000..637be16 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/header.py @@ -0,0 +1,133 @@ +import struct +import sys +import logging + +logger = logging.getLogger(__name__) + +def read_string(f): + """Read a string from the file""" + length = struct.unpack(' 50 else ''}>" + +def read_gguf_file(file_path): + """Read GGUF file and extract header information""" + with open(file_path, 'rb') as f: + # Read magic number + magic = f.read(4) + if magic != b'GGUF': + raise ValueError("Not a valid GGUF file") + + # Read version + version = struct.unpack('") + sys.exit(1) + + gguf_file_path = sys.argv[1] + read_gguf_file(gguf_file_path) diff --git a/src/xpu/flamegraph/qwen3.cu/runcu.cu b/src/xpu/flamegraph/qwen3.cu/runcu.cu new file mode 100644 index 0000000..e69afe7 --- /dev/null +++ b/src/xpu/flamegraph/qwen3.cu/runcu.cu @@ -0,0 +1,1459 @@ +/* Inference for GGUF Qwen-3 models in pure C */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + + +#ifdef USE_CUBLAS +#include + +// cublas handle +cublasHandle_t g_cublas_handle = nullptr; + +void create_cublas_handle() +{ + cublasStatus_t stat = cublasCreate(&g_cublas_handle); // FIXME cublasDestroy + if (stat != CUBLAS_STATUS_SUCCESS) + { + printf("CUBLAS initialization failed\n"); + exit(EXIT_FAILURE); + } +} +void destroy_cublas_handle() +{ + cublasStatus_t stat = cublasDestroy(g_cublas_handle); + if (stat != CUBLAS_STATUS_SUCCESS) + { + printf("CUBLAS initialization failed\n"); + exit(EXIT_FAILURE); + } +} +#endif + +// ---------------------------------------------------------------------------- +// Transformer model +typedef struct { + int dim; // transformer dimension + int hidden_dim; // for ffn layers + int n_layers; // number of layers + int n_heads; // number of query heads + int n_kv_heads; // number of key/value heads (can be < query heads because of multiquery) + int vocab_size; // vocabulary size + int seq_len; // max sequence length + int head_dim; // attention dimension +} Config; + +typedef struct { + // token embedding table + float* token_embedding_table; // (vocab_size, dim) + // weights for rmsnorms in each layer + float* rms_att_weight; // (layer, dim) + float* rms_ffn_weight; // (layer, dim) + // weights for matmuls + float* wq; // (layer, dim, n_heads * head_dim) + float* wk; // (layer, dim, n_kv_heads * head_dim) + float* wv; // (layer, dim, n_kv_heads * head_dim) + float* wo; // (layer, n_heads * head_dim, dim) + float* wq_norm; // (layer, head_dim) + float* wk_norm; // (layer, head_dim) + // weights for ffn. w1 = up, w3 = gate, w2 = down + float* w1; // (layer, dim, hidden_dim) + float* w2; // (layer, hidden_dim, dim) + float* w3; // (layer, dim, hidden_dim) + // final rmsnorm + float* rms_final_weight; // (dim,) + // Same as token_embedding_table. GGUF has the final layer anyway + float* wcls; +} TransformerWeights; + +typedef struct { + // current wave of activations + float* x; // activation at current time stamp (dim,) + float* xb; // buffer (dim,) + float* xb2; // an additional buffer just for convenience (dim,) + float* xb3; // an additional buffer just for convenience (att_head_dim,) + float* hb; // buffer for hidden dimension in the ffn (hidden_dim,) + float* hb2; // buffer for hidden dimension in the ffn (hidden_dim,) + float* q; // query (att_head_dim,) + float* k; // key (dim,) + float* v; // value (dim,) + float* att; // buffer for scores/attention values (n_heads, seq_len) + float* logits; // output logits + + float *d_logits; // CUDA logits + + // kv cache + float* key_cache; // (layer, seq_len, dim) + float* value_cache; // (layer, seq_len, dim) +} RunState; + +typedef struct { + Config config; // the hyperparameters of the architecture (the blueprint) + TransformerWeights weights; // the weights of the model + RunState state; // buffers for the "wave" of activations in the forward pass + int fd; // file descriptor for memory mapping + float* data; // memory mapped data pointer + ssize_t file_size; // size of the checkpoint file in bytes +} Transformer; + + +// Macro for checking CUDA errors +#define CUDA_CHECK(call) do { \ + cudaError_t err = call; \ + if (err != cudaSuccess) { \ + fprintf(stderr, "CUDA error at %s:%d - %s\n", __FILE__, __LINE__, cudaGetErrorString(err)); \ + exit(EXIT_FAILURE); \ + } \ +} while(0) + +void malloc_run_state(RunState* s, Config *p) { + // we calloc instead of malloc to keep valgrind happy + int att_head_dim = p->n_heads * p->head_dim; + int kv_dim = p->n_kv_heads * p->head_dim; // 1024 + + CUDA_CHECK(cudaMalloc(&s->x, p->dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->xb, p->dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->xb2, p->dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->xb3, att_head_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->hb, p->hidden_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->hb2, p->hidden_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->q, att_head_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->k, kv_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->v, kv_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->att, p->n_heads * p->seq_len * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->d_logits, p->vocab_size * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->key_cache, p->n_layers * p->seq_len * kv_dim * sizeof(float))); + CUDA_CHECK(cudaMalloc(&s->value_cache, p->n_layers * p->seq_len * kv_dim * sizeof(float))); + + s->logits = (float *)calloc(p->vocab_size, sizeof(float)); + + // ensure all mallocs went fine + if (!s->x || !s->xb || !s->xb2 || !s->xb3 || !s->hb || !s->hb2 || !s->q || !s->k || !s->v || !s->att || !s->logits || !s->d_logits || !s->key_cache || !s->value_cache) { + fprintf(stderr, "malloc failed!\n"); + exit(EXIT_FAILURE); + } +} + +void free_run_state(RunState* s) { + CUDA_CHECK(cudaFree(s->x)); + CUDA_CHECK(cudaFree(s->xb)); + CUDA_CHECK(cudaFree(s->xb2)); + CUDA_CHECK(cudaFree(s->xb3)); + CUDA_CHECK(cudaFree(s->hb)); + CUDA_CHECK(cudaFree(s->hb2)); + CUDA_CHECK(cudaFree(s->q)); + CUDA_CHECK(cudaFree(s->k)); + CUDA_CHECK(cudaFree(s->v)); + CUDA_CHECK(cudaFree(s->att)); + CUDA_CHECK(cudaFree(s->d_logits)); + CUDA_CHECK(cudaFree(s->key_cache)); + CUDA_CHECK(cudaFree(s->value_cache)); + free(s->logits); +} + +// Map GGUF layers to transformer weights +void memory_map_weights(TransformerWeights* w, Config* p, void* pt) { + //unsigned long long n_layers = p->n_layers; + float *ptr = (float*) pt; + + w->wcls = ptr; // last layer in TR + ptr += p->vocab_size * p->dim; + w->rms_final_weight = ptr; // right before the last + ptr += p->dim; + w->token_embedding_table = ptr; // first layer + ptr += p->vocab_size * p->dim; + w->wk = ptr; + ptr += p->dim * (p->n_kv_heads * p->head_dim); // 1024 x 1024 = dim (1024) x num_kv_heads (8) x p->head_dim (128) + w->wk_norm = ptr; + ptr += p->head_dim; //head_dim (128) + w->rms_att_weight = ptr; + ptr += p->dim; //dimension (1024) + w->wo = ptr; + ptr += (p->n_heads * p->head_dim) * p->dim; // attention heads (16) x head dim (128) * dim + w->wq = ptr; + ptr += p->dim * (p->n_heads * p->head_dim); + w->wq_norm = ptr; + ptr += p->head_dim; //head_dim (128) + w->wv = ptr; + ptr += p->dim * (p->n_kv_heads * p->head_dim); // equal to wk + w->w2 = ptr; + ptr += p->hidden_dim * p->dim; //ffn.down 3072 *1024 + w->w3 = ptr; + ptr += p->dim * p->hidden_dim; // ffn.gate + w->rms_ffn_weight = ptr; + ptr += p->dim; // ffn.norm + w->w1 = ptr; + ptr += p->dim * p->hidden_dim; //ffn.up +} + +// -------------------------------------- +// read GGUF +void read_checkpoint(char *checkpoint, Config *config, TransformerWeights* weights, int* fd, float** data, ssize_t* file_size) { + FILE *file = fopen(checkpoint, "rb"); + if (!file) { fprintf(stderr, "Couldn't open file %s\n", checkpoint); exit(EXIT_FAILURE); } + fseek(file, 0, SEEK_END); // move file pointer to end of file + *file_size = (ssize_t)ftell(file); // get the file size, in bytes + fclose(file); + + //printf("file size is %zd", *file_size); + + // memory map the Transformer weights into the data pointer + *fd = open(checkpoint, O_RDONLY); // open in read only mode + if (*fd == -1) { fprintf(stderr, "open failed!\n"); exit(EXIT_FAILURE); } + + *data = (float*) mmap(NULL, *file_size, PROT_READ, MAP_PRIVATE, *fd, 0); + if (*data == MAP_FAILED) { fprintf(stderr, "mmap failed!\n"); exit(EXIT_FAILURE); } + + // CUDA + float* d_weights_ptr; + // gguf total header = file size - (last tensor size + last offset) + size_t weights_size = *file_size - 5951648; // skip header bytes. header_size = 5951648 TODO + CUDA_CHECK(cudaMalloc((void**)&d_weights_ptr, weights_size)); + CUDA_CHECK(cudaMemcpy(d_weights_ptr, (*data) + 5951648/4, weights_size, cudaMemcpyHostToDevice)); + + memory_map_weights(weights, config, d_weights_ptr); +} + +void build_transformer(Transformer *t, char *checkpoint_path) { + // read in the Weights from the GGUF + read_checkpoint(checkpoint_path, &t->config, &t->weights, &t->fd, &t->data, &t->file_size); + // allocate the RunState buffers + malloc_run_state(&t->state, &t->config); +} + +void free_transformer(Transformer *t) { + if (t->data && t->data != MAP_FAILED) { + munmap(t->data, t->file_size); + t->data = NULL; + } + // CUDA_CHECK(cudaFree(t->weights.token_embedding_table)); +} + +// load the GGUF config file +void load_config(Transformer *t) { + FILE *f = fopen("header.txt", "r"); + if (!f) {perror("Failed to open header.txt"); exit(1);} + + char line[512]; + int line_num = 0; + + while (fgets(line, sizeof(line), f)) { + line_num++; + + char *cfg = strtok(line, " "); + while (cfg) { + char *eq = strchr(cfg, '='); + if (!eq) { + //fprintf(stderr, "Warning: malformed cfg on line %d: %s\n", line_num, cfg); + cfg = strtok(NULL, " "); + continue; + } + + *eq = '\0'; + char *key = cfg; + char *val = eq + 1; + + // Match keys + if (strcmp(key, "QWEN3_EMBEDDING_LENGTH") == 0) { + t->config.dim = atoi(val); + } else if (strcmp(key, "QWEN3_FEED_FORWARD_LENGTH") == 0) { + t->config.hidden_dim = atoi(val); + } else if (strcmp(key, "QWEN3_BLOCK_COUNT") == 0) { + t->config.n_layers = atoi(val); + } else if (strcmp(key, "QWEN3_ATTENTION_HEAD_COUNT") == 0) { + t->config.n_heads = atoi(val); + } else if (strcmp(key, "QWEN3_ATTENTION_HEAD_COUNT_KV") == 0) { + t->config.n_kv_heads = atoi(val); + } else if (strcmp(key, "QWEN3_CONTEXT_LENGTH") == 0) { + t->config.seq_len = atoi(val); + t->config.seq_len = 4096; //OVERWRITE TEMP + } else if (strcmp(key, "QWEN3_ATTENTION_KEY_LENGTH") == 0) { + t->config.head_dim = atoi(val); + }else if (strncmp(key, "TOKENIZER_GGML_TOKENS", 22) == 0) { + char *len_ptr = strstr(val, "ARRAY_LENGTH="); //handling nested config + if (len_ptr) { + t->config.vocab_size = atoi(len_ptr + strlen("ARRAY_LENGTH=")); + } + } + cfg = strtok(NULL, " "); + } + } + fclose(f); +} + +// ---------------------------------------------------------------------------- +// neural net blocks; the dynamics of the Transformer +// CUDA version + +// Utility routine to divide a into ceiling of b parts +int divUp(int a, int b) +{ + return (a - 1) / b + 1; +} + +const int num_threads_lrg = 1024; +const int num_threads_med = 256; + +//========== RMS ===== +__global__ void rmsnorm_kernel(float *o, float *x, float *weight, int size, int elementsPerThread) +{ + __shared__ float sdata[num_threads_lrg]; + + // compute partial sum of squares + float ss = 0.0f; + for (int i = 0; i < elementsPerThread; i++) + { + int j = threadIdx.x + i * num_threads_lrg; + if (j < size) + ss += x[j] * x[j]; + } + + sdata[threadIdx.x] = ss; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) + sdata[threadIdx.x] += sdata[threadIdx.x + stride]; + __syncthreads(); + } + + if (threadIdx.x == 0) { + ss = sdata[0] / size + 1e-6f; + sdata[0] = 1.0f / sqrtf(ss); + } + __syncthreads(); + ss = sdata[0]; + + for (int i = 0; i < elementsPerThread; i++) + { + int j = threadIdx.x + i * num_threads_lrg; + if (j < size) + o[j] = weight[j] * (ss * x[j]); + } +} + +void rmsnorm(float *o, float *x, float *weight, int size) +{ + int elementsPerThread = divUp(size, num_threads_lrg); + rmsnorm_kernel<<<1, num_threads_lrg>>>(o, x, weight, size, elementsPerThread); +} + + +__global__ void rmsnorm_kernel_multihead(float *o, float *x, float *weight, + int head_dim, int elementsPerThread, + int n_heads) +{ + __shared__ float sdata[num_threads_lrg]; + + // Get head index from block ID + int head = blockIdx.x; + if (head >= n_heads) return; + + // Calculate offsets for this head + float *head_input = x + head * head_dim; + float *head_output = o + head * head_dim; + // Note: weight is shared across all heads (same normalization weights) + + // compute partial sum of squares + float ss = 0.0f; + for (int i = 0; i < elementsPerThread; i++) + { + int j = threadIdx.x + i * num_threads_lrg; + if (j < head_dim) + ss += head_input[j] * head_input[j]; + } + + sdata[threadIdx.x] = ss; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (threadIdx.x < stride) + sdata[threadIdx.x] += sdata[threadIdx.x + stride]; + __syncthreads(); + } + + if (threadIdx.x == 0) { + ss = sdata[0] / head_dim + 1e-6f; + sdata[0] = 1.0f / sqrtf(ss); + } + __syncthreads(); + ss = sdata[0]; + + for (int i = 0; i < elementsPerThread; i++) + { + int j = threadIdx.x + i * num_threads_lrg; + if (j < head_dim) + head_output[j] = weight[j] * (ss * head_input[j]); + } +} + + +void rmsnorm_multihead(float *o, float *x, float *weight, int head_dim, int n_heads) +{ + int elementsPerThread = divUp(head_dim, num_threads_lrg); + + // Launch one block per head + rmsnorm_kernel_multihead<<>>( + o, x, weight, head_dim, elementsPerThread, n_heads + ); +} + + +//========== Softmax in GPU ===== +__device__ void softmax_gpu(float *__restrict__ x, int size) +{ + __shared__ float sdata[num_threads_lrg]; + int tid = threadIdx.x; + int step = blockDim.x; + + // find max value (for numerical stability) + float max_val = tid < size ? x[tid] : -INFINITY; + for (int i = tid + step; i < size; i += step) + { + if (x[i] > max_val) + max_val = x[i]; + } + + sdata[tid] = max_val; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) + sdata[tid] = fmaxf(sdata[tid], sdata[tid + stride]); + __syncthreads(); + } + max_val = sdata[0]; + + // exp and sum + float sum = 0.0f; + for (int i = tid; i < size; i += step) + { + x[i] = expf(x[i] - max_val); + sum += x[i]; + } + + sdata[tid] = sum; + __syncthreads(); + + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) + sdata[tid] += sdata[tid + stride]; + __syncthreads(); + } + sum = sdata[0]; + + // normalize + for (int i = tid; i < size; i += step) + x[i] /= sum; +} + +void softmax(float* x, int size) { + // find max value (for numerical stability) + float max_val = x[0]; + for (int i = 1; i < size; i++) { + if (x[i] > max_val) { + max_val = x[i]; + } + } + // exp and sum + float sum = 0.0f; + for (int i = 0; i < size; i++) { + x[i] = expf(x[i] - max_val); + sum += x[i]; + } + + // normalize + for (int i = 0; i < size; i++) { + x[i] /= sum; + } +} + + +// Use cuBLAS for matmul to leverage this included, high-performance library. +#ifdef USE_CUBLAS +void matmul(float *xout, float *x, float *w, int n, int d) +{ + // W (d,n) @ x (n,) -> xout (d,) + // W is stored in this order: (n=0,d=0), (n=1,d=0), (n=2,d=0), ... + // so W is n x d in cublas terms & we'll need to transpose. + // Sgemv does y = alpha * op(A) * x + beta * y (modifying y) + // where op can transpose the matrix A + // Translating to our local vars, that is + // xout = 1.0*op(w)*x + 0.0*xout + float alpha = 1.0f; + float beta = 0.0f; // when this is 0, xout will not be used for input + cublasSgemv(g_cublas_handle, CUBLAS_OP_T, n, d, &alpha, w, n, x, 1, &beta, xout, 1); +} +#else + +__global__ void matmul_kernel(float *xout, float *x, float *w, int n, int d) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + int tid = threadIdx.x; + + extern __shared__ float shared_x[]; + + // Load x into shared memory in chunks + for (int offset = 0; offset < n; offset += blockDim.x) { + if (offset + tid < n) { + shared_x[tid] = x[offset + tid]; + } + __syncthreads(); + + if (i < d) { + float sum = 0.0f; + int chunk_size = min(blockDim.x, n - offset); + + // Vectorized loads and computation + float4 *w_vec = (float4*)(w + i * n + offset); + float4 *x_vec = (float4*)shared_x; + + int vec_ops = chunk_size / 4; + for (int v = 0; v < vec_ops; v++) { + float4 w4 = w_vec[v]; + float4 x4 = x_vec[v]; + sum += w4.x * x4.x + w4.y * x4.y + w4.z * x4.z + w4.w * x4.w; + } + + // Handle remaining elements + for (int j = vec_ops * 4; j < chunk_size; j++) { + sum += w[i * n + offset + j] * shared_x[j]; + } + + if (offset == 0) xout[i] = sum; + else xout[i] += sum; + } + __syncthreads(); + } +} + +void matmul(float *xout, float *x, float *w, int n, int d) { + int block_size = 256; + int grid_size = (d + block_size - 1) / block_size; + int shared_mem = block_size * sizeof(float); + matmul_kernel<<>>(xout, x, w, n, d); +} +#endif + +// multihead attention +__global__ void multi_head_attention_kernel(int pos, int seq_len, float *sq, float *satt, float *sxb, float *key_cache, float *value_cache, int kv_dim, int kv_mul, int head_size, int loff) +{ + int h = blockIdx.x; + // get the query vector for this head + float *q = sq + h * head_size; + // attention scores for this head + float *att = satt + h * seq_len; + // iterate over all timesteps, including the current one + // In CUDA, each thread does a small portion of the calc + for (int t = threadIdx.x; t <= pos; t += blockDim.x) + { + // get the key vector for this head and at this timestep + float *k = key_cache + loff + t * kv_dim + (h / kv_mul) * head_size; + // calculate the attention score as the dot product of q and k + float score = 0.0f; + for (int i = 0; i < head_size; i++) + { + score += q[i] * k[i]; + } + score /= sqrtf(head_size); + // save the score to the attention buffer + att[t] = score; + } + // above was this threads portion of the iteration. wait for all threads to finish + __syncthreads(); + + // softmax the scores to get attention weights, from 0..pos inclusively + softmax_gpu(att, pos + 1); + __syncthreads(); + + // weighted sum of the values, store back into xb + // NOTE: by swapping the order of the for loops (vs. C) a simpler + // version of the code accomplishes the same task and fits more + // naturally with the CUDA way of subdividing the problem. + float *xb = sxb + h * head_size; + for (int i = threadIdx.x; i < head_size; i += blockDim.x) + { + float val = 0.0f; + for (int t = 0; t <= pos; t++) + { + // get the value vector for this head and at this timestep + float *v = value_cache + loff + t * kv_dim + (h / kv_mul) * head_size; + // get the attention weight for this timestep + float a = att[t]; + val += a * v[i]; + } + xb[i] = val; + } +} +void multi_head_attention(int pos, Config *p, RunState *s, int kv_dim, int kv_mul, int head_size, int loff) +{ + multi_head_attention_kernel<<n_heads, num_threads_lrg>>>(pos, p->seq_len, s->q, s->att, s->xb3, s->key_cache, s->value_cache, kv_dim, kv_mul, head_size, loff); +} // Fixed xb to xb3 + +//-------Activation----------- +__global__ void f_silu_elementwise_mul_w3_kernel(float *shb, float *shb2, int hidden_dim) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < hidden_dim) + { + float val = shb2[i]; + // silu(x)=x*σ(x), where σ(x) is the logistic sigmoid + val *= (1.0f / (1.0f + expf(-val))); + // elementwise multiply with w3(x) + val *= shb[i]; + shb2[i] = val; + } +} +void f_silu_elementwise_mul_w3(RunState *s, int hidden_dim) +{ + f_silu_elementwise_mul_w3_kernel<<>>(s->hb, s->hb2, hidden_dim); +} + +//------Residual Connection---------- +__global__ void accum_kernel(float *a, float *b, int size) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size) + { + a[i] += b[i]; + } +} + +void accum(float *a, float *b, int size) +{ + accum_kernel<<>>(a, b, size); +} + + +// Unified kernel for both Q and K RoPE rotation +__global__ void RoPe_rotation_kernel_multihead(int pos, float *tensor, int head_dim, + int kv_dim, int att_head_dim, int n_heads) +{ + // Get head and element indices + int head = blockIdx.x; + int i = threadIdx.x; + + if (head >= n_heads || i >= head_dim / 2) return; + + // Calculate offset for this head + float *head_tensor = tensor + head * head_dim; + + float freq = 1.0f / powf(1000000.0f, (float)i / (head_dim / 2)); + float fcr = cosf(pos * freq); + float fci = sinf(pos * freq); + + // Rotate tensor head (works for both Q and K) + float x = head_tensor[i]; + float y = head_tensor[i + head_dim / 2]; + head_tensor[i] = x * fcr - y * fci; + head_tensor[i + head_dim / 2] = x * fci + y * fcr; +} + +// Unified host function for both Q and K RoPE rotation +void RoPe_rotation_multihead(int pos, float *tensor, int head_dim, int kv_dim, + int att_head_dim, int n_heads) +{ + // Launch one block per head, threads per element + RoPe_rotation_kernel_multihead<<>>( + pos, tensor, head_dim, kv_dim, att_head_dim, n_heads + ); +} + + +//====================================== +// FORWARD PASS +float *forward(Transformer *transformer, int token, int pos) +{ + Config* p = &transformer->config; + TransformerWeights* w = &transformer->weights; + RunState* s = &transformer->state; + + int kv_dim = p->n_kv_heads * p->head_dim; + int kv_mul = p->n_heads / p->n_kv_heads; + int att_head_dim = p->n_heads * p->head_dim; + + int layer_offset = 62923776/4; // offset to the GGUF next layer for the same tensor type TODO + + // copy the token embedding into s->x in GPU, STARTING POINT - x is passing through. + float *content_row = w->token_embedding_table + token * p->dim; + CUDA_CHECK(cudaMemcpy(s->x, content_row, p->dim * sizeof(*s->x), cudaMemcpyHostToDevice)); + + // forward all the layers + for (int l = 0; l < p->n_layers; l++) { + // kv cache + int loff = l * p->seq_len * kv_dim; + s->k = s->key_cache + loff + pos * kv_dim; + s->v = s->value_cache + loff + pos * kv_dim; + + // attention rmsnorm + rmsnorm(s->xb, s->x, w->rms_att_weight + l * layer_offset, p->dim); + + // Query, key, value computation + matmul(s->q, s->xb, w->wq + l *layer_offset, p->dim, att_head_dim); + matmul(s->k, s->xb, w->wk + l *layer_offset, p->dim, kv_dim); + matmul(s->v, s->xb, w->wv + l *layer_offset, p->dim, kv_dim); + + // Apply RMSNorm to ALL query heads + rmsnorm_multihead(s->q, s->q, w->wq_norm + l * layer_offset, p->head_dim, p->n_heads); + // Apply RoPE to Q + RoPe_rotation_multihead(pos, s->q, p->head_dim, kv_dim, att_head_dim, p->n_heads); + + // Apply RMSNorm to ALL key heads + rmsnorm_multihead(s->k, s->k, w->wk_norm + l * layer_offset, p->head_dim, p->n_kv_heads); + // Apply RoPE to K + RoPe_rotation_multihead(pos, s->k, p->head_dim, kv_dim, att_head_dim, p->n_kv_heads); + + multi_head_attention(pos, p, s, kv_dim, kv_mul, p->head_dim, loff); + + // Output projection + matmul(s->xb2, s->xb3, w->wo + l *layer_offset, att_head_dim, p->dim); + + // Grouped-query attetion ends. Residual connection + accum(s->x, s->xb2, p->dim); + + + // FFN + rmsnorm(s->xb, s->x, w->rms_ffn_weight + l *layer_offset /* * p->dim*/ , p->dim); + matmul(s->hb, s->xb, w->w1 + l *layer_offset, p->dim, p->hidden_dim); // Up + matmul(s->hb2, s->xb, w->w3 + l *layer_offset, p->dim, p->hidden_dim); // Gate + f_silu_elementwise_mul_w3(s, p->hidden_dim); // Activation + matmul(s->xb, s->hb2, w->w2 + l *layer_offset, p->hidden_dim, p->dim); // Down + + // residual connection + accum(s->x, s->xb, p->dim); + } + + // Final norm after layers and right before logiting + rmsnorm(s->x, s->x, w->rms_final_weight, p->dim); + + // Calc logits in GPU and move it to CPU + matmul(s->d_logits, s->x, w->wcls, p->dim, p->vocab_size); + CUDA_CHECK(cudaMemcpy(s->logits, s->d_logits, p->vocab_size * sizeof(float), cudaMemcpyDeviceToHost)); + + return s->logits; +} + +// ---------------------------------------------------------------------------- +// TOKENIZER +// The Byte Pair Encoding (BPE) Tokenizer that translates strings <-> tokens + +typedef struct { + char **vocab; + float *merge_scores; + int vocab_size; + unsigned int max_token_length; + unsigned int bos_token_id; + unsigned int eos_token_id; + char prompt_template[1024]; + char system_prompt_template[1024]; +} Tokenizer; + +// ========== GLOBALS ========== + +// TODO +#define MAX_VOCAB 151936 +#define MAX_MERGES 151386 +#define MAX_TOKENS 1024 +#define MAX_TOKEN_LEN 32 + +char *vocab[MAX_VOCAB]; +int num_vocab = 0; + +typedef struct { + char left[MAX_TOKEN_LEN]; + char right[MAX_TOKEN_LEN]; + char merged[MAX_TOKEN_LEN * 2]; + int rank; +} MergeRule; + +MergeRule merges[MAX_MERGES]; +int merge_count = 0; + +unsigned int byte_to_unicode[256]; +char unicode_bytes[256][5]; + +//special token +const char *special_tokens[] = { + "<|im_start|>", + "<|im_end|>", + "", + "" +}; +int num_special_tokens = sizeof(special_tokens) / sizeof(special_tokens[0]); + +typedef struct { + int *data; // token ID array + size_t size; // number of tokens stored + size_t capacity; // total allocated capacity +} TokenBuffer; + +// ========== LOAD VOCAB ========== +void load_vocab(const char *path) { + FILE *f = fopen(path, "r"); + if (!f) { + perror("fopen"); + exit(1); + } + + char line[512]; // large enough to hold most tokens safely + int line_num = 0; + + while (fgets(line, sizeof(line), f)) { + line_num++; + + // Remove trailing newline / carriage return + line[strcspn(line, "\r\n")] = '\0'; + + // Skip empty lines + if (line[0] == '\0') continue; + + if (num_vocab >= MAX_VOCAB) { + fprintf(stderr, "Error: vocab exceeds MAX_VOCAB (%d)\n", MAX_VOCAB); + break; + } + + vocab[num_vocab] = strdup(line); + if (!vocab[num_vocab]) { + fprintf(stderr, "Error: strdup failed on line %d\n", line_num); + exit(1); + } + num_vocab++; + } + fclose(f); +} + +void free_vocab() { + for (int i = 0; i < num_vocab; i++) { + free(vocab[i]); // Free each string + vocab[i] = NULL; + } + num_vocab = 0; +} + +// ========== LOAD MERGES ========== +void load_merges(const char *path) { + FILE *f = fopen(path, "r"); + if (!f) { perror("merges.txt"); exit(1); } + char line[256]; + while (fgets(line, sizeof(line), f)) { + if (line[0] == '#' || isspace(line[0])) continue; + char *a = strtok(line, " \t\r\n"); + char *b = strtok(NULL, " \t\r\n"); + if (!a || !b) continue; + strncpy(merges[merge_count].left, a, MAX_TOKEN_LEN); + strncpy(merges[merge_count].right, b, MAX_TOKEN_LEN); + snprintf(merges[merge_count].merged, sizeof(merges[merge_count].merged), "%s%s", a, b); + merges[merge_count].rank = merge_count; + merge_count++; + } + fclose(f); +} + +bool pair_equals(const char *a, const char *b, const char *left, const char *right) { + return strcmp(a, left) == 0 && strcmp(b, right) == 0; +} + +int get_merge_rank(const char *left, const char *right) { + for (int i = 0; i < merge_count; i++) { + if (pair_equals(left, right, merges[i].left, merges[i].right)) { + return merges[i].rank; + } + } + return INT_MAX; +} + +// ========== BYTE-TO-UNICODE ========== +void init_byte_unicode_map() { + int n = 0; + for (int b = 0; b < 256; b++) { + // party for gpt-2 based tokenizer + if ((b >= 33 && b <= 126) || // ASCII + (b >= 161 && b <= 172) || // Latin-1 + (b >= 174 && b <= 255)) { // Latin-1 except 173 + byte_to_unicode[b] = b; + } else { + byte_to_unicode[b] = 256 + n++; // Map out of Unicode range + } + + // Direct UTF-8 encoding + uint32_t cp = byte_to_unicode[b]; + if (cp < 128) { // 1-byte UTF-8 + unicode_bytes[b][0] = cp; + unicode_bytes[b][1] = '\0'; + } + else if (cp < 2048) { // 2-byte UTF-8 + unicode_bytes[b][0] = 0xC0 | (cp >> 6); + unicode_bytes[b][1] = 0x80 | (cp & 0x3F); + unicode_bytes[b][2] = '\0'; + } + else { // 3-byte UTF-8 + unicode_bytes[b][0] = 0xE0 | (cp >> 12); + unicode_bytes[b][1] = 0x80 | ((cp >> 6) & 0x3F); + unicode_bytes[b][2] = 0x80 | (cp & 0x3F); + unicode_bytes[b][3] = '\0'; + } + } +} + +// TODO. Returns -1 if not a special token, else index into vocab[] +int match_special_token(const char *str, int *match_len) { + for (int i = 0; i < num_special_tokens; i++) { + const char *tok = special_tokens[i]; + int len = strlen(tok); + if (strncmp(str, tok, len) == 0) { + // Now check if this token exists in vocab + for (int j = 0; j < 151936; j++) { + if (strcmp(vocab[j], tok) == 0) { + *match_len = len; + return j; // return token ID + } + } + } + } + return -1; +} + +void build_tokenizer(Tokenizer* t) { + load_vocab("vocab.txt"); + load_merges("merges.txt"); + init_byte_unicode_map(); +} + +// build buffer +void build_token_buffer(TokenBuffer *tb, size_t initial_capacity) { + tb->data = (int *) malloc(initial_capacity * sizeof(int)); + if (!tb->data) { perror("malloc"); exit(1); } + tb->size = 0; + tb->capacity = initial_capacity; +} + +void free_token_buffer(TokenBuffer *tb) { + free(tb->data); + tb->data = NULL; + tb->size = 0; + tb->capacity = 0; +} +void append_tokens(TokenBuffer *tb, const int *tokens, size_t n) { + if (tb->size + n > tb->capacity) { + while (tb->size + n > tb->capacity) tb->capacity *= 2; + tb->data = (int *) realloc(tb->data, tb->capacity * sizeof(int)); + if (!tb->data) { perror("realloc"); exit(1); } + } + memcpy(&tb->data[tb->size], tokens, n * sizeof(int)); + tb->size += n; +} + +void encode(Tokenizer* t, char* rendered_prompt, int* prompt_tokens, int* num_prompt_tokens, int multi_turn) { + if (rendered_prompt == NULL) { fprintf(stderr, "cannot encode NULL text\n"); exit(EXIT_FAILURE); } + char *tokens[MAX_TOKENS]; + int count = 0; + + const char *p = rendered_prompt; + + while (*p) { + int match_len = 0; + int special_id = match_special_token(p, &match_len); + if (special_id >= 0) { + // Special token found → store its string + tokens[count++] = strdup(vocab[special_id]); // Store as string, not ID + p += match_len; + continue; + } + + // Not a special token → convert byte to unicode + unsigned char b = *p++; + tokens[count++] = strdup(unicode_bytes[b]); + } + + // === BPE merge === + bool changed = true; + while (changed) { + int best_rank = INT_MAX; + int best_pos = -1; + + for (int i = 0; i < count - 1; i++) { + int rank = get_merge_rank(tokens[i], tokens[i + 1]); + if (rank < best_rank) { + best_rank = rank; + best_pos = i; + } + } + + if (best_pos == -1) break; + + // if either token is a special token, skip merge + if (tokens[best_pos][0] == '<' && strchr(tokens[best_pos], '|') && + tokens[best_pos + 1][0] == '<' && strchr(tokens[best_pos + 1], '|')) { + break; // don't merge special tokens + } + + // Merge + char *merged = (char *) malloc(MAX_TOKEN_LEN * 2); + snprintf(merged, MAX_TOKEN_LEN * 2, "%s%s", tokens[best_pos], tokens[best_pos + 1]); + free(tokens[best_pos]); + free(tokens[best_pos + 1]); + tokens[best_pos] = merged; + for (int i = best_pos + 1; i < count - 1; i++) { + tokens[i] = tokens[i + 1]; + } + count--; + changed = true; + } + + // === Map tokens to token IDs === + int token_ids[MAX_TOKENS]; + int token_id_count = 0; + + for (int i = 0; i < count; i++) { + int id = -1; + for (int j = 0; j < 151936; j++) { + if (strcmp(tokens[i], vocab[j]) == 0) { + id = j; + break; + } + } + + if (id == -1) { + fprintf(stderr, "Token not found in vocab: [%s]\n", tokens[i]); + } else { + token_ids[token_id_count++] = id; + //printf("[%s] → id = %d\n", tokens[i], id); + } + // TODO + for (int i = 0; i < token_id_count; i++) { + prompt_tokens[i] = token_ids[i]; + } + *num_prompt_tokens = token_id_count; + } +} + +// ========== Decoding logic========== +// Inverse byte-to-unicode: utf8 string → byte value (0–255) +int unicode_to_byte(const char *utf8) { + for (int b = 0; b < 256; b++) { + if (strcmp(utf8, unicode_bytes[b]) == 0) + return b; + } + return -1; // not found +} + +char *decode_token_id(int token_id) { + const char *encoded = vocab[token_id]; + char *out = (char *) malloc(1024); + int rlen = 0; + for (int i = 0; encoded[i]; ) { + int matched = 0; + for (int len = 1; len <= 3 && encoded[i + len - 1]; len++) { + char utf8[5] = {0}; + strncpy(utf8, &encoded[i], len); + int b = unicode_to_byte(utf8); + if (b >= 0) { + out[rlen++] = (char)b; + i += len; + matched = 1; + break; + } + } + if (!matched) break; + } + out[rlen] = '\0'; + return out; +} + +// ---------------------------------------------------------------------------- +// The Sampler, which takes logits and returns a sampled token +// sampling can be done in a few ways: greedy argmax, sampling, top-p sampling + +typedef struct { + float prob; + int index; +} ProbIndex; // struct used when sorting probabilities during top-p sampling + +typedef struct { + int vocab_size; + ProbIndex *probindex; // buffer used in top-p sampling + float temperature; + float topp; + unsigned long long rng_state; +} Sampler; + +int sample_argmax(float *probabilities, int n) { + // return the index that has the highest probability + int max_i = 0; + float max_p = probabilities[0]; + for (int i = 1; i < n; i++) { + if (probabilities[i] > max_p) { + max_i = i; + max_p = probabilities[i]; + } + } + return max_i; +} + +int sample_mult(float *probabilities, int n, float coin) { + // sample index from probabilities (they must sum to 1!) + // coin is a random number in [0, 1), usually from random_f32() + float cdf = 0; + for (int i = 0; i < n; i++) { + cdf += probabilities[i]; + if (coin < cdf) + return i; + } + return n - 1; // in case of rounding errors +} + +int compare(const void *a, const void *b) { + ProbIndex *a_ = (ProbIndex *) a; + ProbIndex *b_ = (ProbIndex *) b; + if (a_->prob > b_->prob) return -1; + if (a_->prob < b_->prob) return 1; + return 0; +} + +int sample_topp(float *probabilities, int n, float topp, ProbIndex *probindex, float coin) { + // top-p sampling (or "nucleus sampling") samples from the smallest set of + // tokens that exceed probability topp. This way we never sample tokens that + // have very low probabilities and are less likely to go "off the rails". + // coin is a random number in [0, 1), usually from random_f32() + + int n0 = 0; + // quicksort indices in descending order of probabilities + // values smaller than (1 - topp) / (n - 1) cannot be part of the result + // so for efficiency we crop these out as candidates before sorting + const float cutoff = (1.0f - topp) / (n - 1); + for (int i = 0; i < n; i++) { + if (probabilities[i] >= cutoff) { + probindex[n0].index = i; + probindex[n0].prob = probabilities[i]; + n0++; + } + } + qsort(probindex, n0, sizeof(ProbIndex), compare); + + // truncate the list where cumulative probability exceeds topp + float cumulative_prob = 0; + int last_idx = n0 - 1; // in case of rounding errors consider all elements + for (int i = 0; i < n0; i++) { + cumulative_prob += probindex[i].prob; + if (cumulative_prob > topp) { + last_idx = i; + break; // we've exceeded topp by including last_idx + } + } + + // sample from the truncated list + float r = coin * cumulative_prob; + float cdf = 0; + for (int i = 0; i <= last_idx; i++) { + cdf += probindex[i].prob; + if (r < cdf) + return probindex[i].index; + } + return probindex[last_idx].index; // in case of rounding errors +} + +void build_sampler(Sampler* sampler, int vocab_size, float temperature, float topp, unsigned long long rng_seed) { + sampler->vocab_size = vocab_size; + sampler->temperature = temperature; + sampler->topp = topp; + sampler->rng_state = rng_seed; + // buffer only used with nucleus sampling; may not need but it's ~small + sampler->probindex = (ProbIndex *) malloc(sampler->vocab_size * sizeof(ProbIndex)); +} + +void free_sampler(Sampler* sampler) { + free(sampler->probindex); +} + +unsigned int random_u32(unsigned long long *state) { + // xorshift rng: https://en.wikipedia.org/wiki/Xorshift#xorshift.2A + *state ^= *state >> 12; + *state ^= *state << 25; + *state ^= *state >> 27; + return (*state * 0x2545F4914F6CDD1Dull) >> 32; +} +float random_f32(unsigned long long *state) { // random float32 in [0,1) + return (random_u32(state) >> 8) / 16777216.0f; +} + +int sample(Sampler* sampler, float* logits) { + // sample the token given the logits and some hyperparameters + int next; + if (sampler->temperature == 0.0f) { + // greedy argmax sampling: take the token with the highest probability + next = sample_argmax(logits, sampler->vocab_size); + } else { + // apply the temperature to the logits + for (int q=0; qvocab_size; q++) { logits[q] /= sampler->temperature; } + // apply softmax to the logits to get the probabilities for next token + softmax(logits, sampler->vocab_size); + // flip a (float) coin (this is our source of entropy for sampling) + float coin = random_f32(&sampler->rng_state); + // we sample from this distribution to get the next token + if (sampler->topp <= 0 || sampler->topp >= 1) { + // simply sample from the predicted probability distribution + next = sample_mult(logits, sampler->vocab_size, coin); + } else { + // top-p (nucleus) sampling, clamping the least likely tokens to zero + next = sample_topp(logits, sampler->vocab_size, sampler->topp, sampler->probindex, coin); + } + } + return next; +} + +// ---------------------------------------------------------------------------- +// utilities: time +long time_in_ms() { + // return time in milliseconds, for benchmarking the model speed + struct timespec time; + clock_gettime(CLOCK_REALTIME, &time); + return time.tv_sec * 1000 + time.tv_nsec / 1000000; +} + +// ------------------------------------------------------------------------ +// read input +void read_stdin(const char* guide, char* buffer, size_t bufsize) { + // read a line from stdin, up to but not including \n + printf("%s", guide); + if (fgets(buffer, bufsize, stdin) != NULL) { + size_t len = strlen(buffer); + if (len > 0 && buffer[len - 1] == '\n') { + buffer[len - 1] = '\0'; + } + } +} + +// ---------------------------------------------------------------------------- +// chat loop +void chat(Transformer* transformer, Tokenizer* tokenizer, Sampler* sampler, char* cli_user_prompt, char* cli_system_prompt, int think_on, int multi_turn, int tps, TokenBuffer* tb, int single_prompt) { + // buffers for reading the system prompt and user prompt from stdin + char system_prompt[512]; + char user_prompt[8192]; + char rendered_prompt[8192]; + int num_prompt_tokens = 0; + int* prompt_tokens = (int* )malloc(8192 * sizeof(int)); + //int user_idx; + + // start the main loop + int8_t user_turn = 1; // user starts + int next; // will store the next token in the sequence + int token; // stores the current token to feed into the transformer + //int prev_token; + int pos = 0; // position in the sequence + double timer = -1.0; // TPS timer start + int count = 0; // decoded token + + while (1) { + if (user_turn) { + if (pos == 0){ + if (!single_prompt) { + read_stdin("Enter system prompt (or Enter to skip): ", system_prompt, sizeof(system_prompt)); + } else { + system_prompt[0] = '\0'; + } + } + if (!single_prompt) { + read_stdin("Q: ", user_prompt, sizeof(user_prompt)); + } else { + if (cli_user_prompt) { + strncpy(user_prompt, cli_user_prompt, sizeof(user_prompt) - 1); + user_prompt[sizeof(user_prompt) - 1] = '\0'; + } else { + user_prompt[0] = '\0'; + } + } + // terminate if user enters a blank prompt + if (!user_prompt[0]) { + break; + } + + // render user/system prompts for Qwen3 + if (pos == 0 && system_prompt[0] != '\0') { + char system_template[] = "<|im_start|>system\n%s<|im_end|>\n<|im_start|>user\n%s<|im_end|>\n<|im_start|>assistant\n"; + sprintf(rendered_prompt, system_template, system_prompt, user_prompt); + } else { + char user_template[] = "<|im_start|>user\n%s<|im_end|>\n<|im_start|>assistant\n"; + sprintf(rendered_prompt, user_template, user_prompt); + } + + if (!think_on) { + strcat(rendered_prompt, "\n\n\n"); + } + + // encode the rendered prompt into tokens + encode(tokenizer, rendered_prompt, prompt_tokens, &num_prompt_tokens, multi_turn); + pos = 0; // reset the user index + user_turn = 0; + if (multi_turn) { + append_tokens(tb, prompt_tokens, num_prompt_tokens); + for (size_t i = 0; i < tb->size; i++) { + // printf("%d ", tb->data[i]); + } + printf("\n"); + } + printf("A: "); + } + + if (pos < (multi_turn ? tb->size : num_prompt_tokens)) { + token = (multi_turn) ? tb->data[pos] : prompt_tokens[pos]; + } else { + token = next; + } + //printf("right before foreward: %d\n", token); + // forward the transformer to get logits for the next token + float* logits = forward(transformer, token, pos); + next = sample(sampler, logits); + + pos++; + //printf("num_prompt_tokens: %d \n", num_prompt_tokens); + // decoding and printing + if (pos >= (multi_turn ? tb->size : num_prompt_tokens)) { + if (multi_turn) { + append_tokens(tb, &next, 1); + //printf("next token: %d\n", next); + } + + if (next == 151645) { // EOS token ID - TODO + printf("\n"); + user_turn = 1; + + // TPS + if (tps) { + fprintf(stderr, "tok/s: %f\n", count / (double)(time_in_ms() - timer) * 1000); + timer = -1; + count = 0; + } + + // Exit after single prompt + if (single_prompt) { + break; + } + } + else { + char *decoded = decode_token_id(next); + printf("%s", decoded); + fflush(stdout); + free(decoded); + + if (tps) { + count += 1; + // timer starts after the first token generation + if (timer == -1.0) {timer = time_in_ms();} + } + } + } + } +} + +// ---------------------------------------------------------------------------- +// CLI +void error_usage() { + fprintf(stderr, "Usage: run [options]\n"); + fprintf(stderr, "Example: ./run Qwen3-0.6B-FP32.gguf\n"); + fprintf(stderr, " ./run Qwen3-0.6B-FP32.gguf -q \"What is CUDA?\"\n"); + fprintf(stderr, "Options:\n"); + fprintf(stderr, " -t temperature in [0,inf], default 0.6\n"); + fprintf(stderr, " -p p value in top-p (nucleus) sampling in [0,1] default 0.95\n"); + fprintf(stderr, " -s random seed, default time(NULL)\n"); + fprintf(stderr, " -m multi-turn: 0 = off (defualt), 1 = on\n"); + fprintf(stderr, " -k reasoning: 0 = off (defualt), 1 = on\n"); + fprintf(stderr, " -r TPS: 0 = off (defualt), 1 = on\n"); + fprintf(stderr, " -q single prompt mode (run once and exit)\n"); + exit(EXIT_FAILURE); +} + + +int main(int argc, char *argv[]) { + // default parameters + char *checkpoint_path = NULL; // e.g. out/model.bin + float temperature = 0.6f; // 0.0 = greedy deterministic. 1.0 = original. don't set higher + float topp = 0.95f; // top-p in nucleus sampling. 1.0 = off. 0.9 works well, but slower + int steps = 256; // number of steps to run for + char *prompt = NULL; // prompt string + unsigned long long rng_seed = 0; // seed rng with time by default + char *system_prompt = NULL; // the (optional) system prompt to use in chat mode + int multi_turn = 0; // multi-turn conversation + int think_on = 0; // reasoning on + int tps = 0; // TPS + int single_prompt = 0; // single prompt mode + + if (argc >= 2) { checkpoint_path = argv[1]; } else { error_usage(); } + for (int i = 2; i < argc; i+=2) { + // do some basic validation + if (i + 1 >= argc) { error_usage(); } // must have arg after flag + if (argv[i][0] != '-') { error_usage(); } // must start with dash + if (strlen(argv[i]) != 2) { error_usage(); } // must be -x (one dash, one letter) + // read in the args + if (argv[i][1] == 't') { temperature = atof(argv[i + 1]); } + else if (argv[i][1] == 'p') { topp = atof(argv[i + 1]); } + else if (argv[i][1] == 's') { rng_seed = atoi(argv[i + 1]); } + else if (argv[i][1] == 'q') { prompt = argv[i + 1]; single_prompt = 1; } + else if (argv[i][1] == 'm') {if ((argv[i+1][0] == '0' || argv[i+1][0] == '1') && argv[i+1][1] == '\0') { + multi_turn = argv[i+1][0] - '0';} else { error_usage(); } } + else if (argv[i][1] == 'k') {if ((argv[i+1][0] == '0' || argv[i+1][0] == '1') && argv[i+1][1] == '\0') { + think_on = argv[i+1][0] - '0';} else { error_usage(); } } + else if (argv[i][1] == 'r') {if ((argv[i+1][0] == '0' || argv[i+1][0] == '1') && argv[i+1][1] == '\0') { + tps = argv[i+1][0] - '0';} else { error_usage(); } } + else { error_usage(); } + } + + // parameter validation/overrides + if (rng_seed <= 0) rng_seed = (unsigned int)time(NULL); + if (temperature < 0.0) temperature = 0.0; + if (topp < 0.0 || 1.0 < topp) topp = 0.9; + if (steps < 0) steps = 0; + + // read config + Transformer transformer; + load_config(&transformer); + + // build the Transformer via the GGUF file + build_transformer(&transformer, checkpoint_path); + + // build the Tokenizer + Tokenizer tokenizer; + build_tokenizer(&tokenizer); + + // multi-turn buffer + TokenBuffer tb; + build_token_buffer(&tb, 1024); + + Sampler sampler; + build_sampler(&sampler, transformer.config.vocab_size, temperature, topp, rng_seed); + + #ifdef USE_CUBLAS + // cuBlas handle + create_cublas_handle(); + #endif + + if (!single_prompt) { + printf("Multi-turn = %s, thinKing = %s, tps(R) = %s, Temperature = %.2f, top-P = %.2f\n", multi_turn ? "on" : "off", think_on ? "on" : "off", tps ? "on" : "off", temperature, topp); + printf("Press Enter to exit the chat\n"); + } + + // run! + chat(&transformer, &tokenizer, &sampler, prompt, system_prompt, think_on, multi_turn, tps, &tb, single_prompt); + + // memory and file handles cleanup + free_sampler(&sampler); + free_vocab(); + free_token_buffer(&tb); + free_transformer(&transformer); + + #ifdef USE_CUBLAS + destroy_cublas_handle(); + #endif + + return 0; +}