dylanyunlong/project_6
1
0
Fork 0
Code Issues Pull Requests Actions Projects Releases Wiki Activity

ref(upstream): FULL TREE — Deep-Spark xllm (1470) + ds_vllm csrc/models (703)

Browse Source 
Replaces cherry-picked upstream_ref with complete source trees.

xllm/ — Iluvatar official C++ inference engine (15MB, 1470 files)
  Complete: kernels → layers → models → runtime → scheduler → api
  Excluded: .git, binary images, third_party submodule checkouts

ds_vllm/ — Iluvatar official vllm fork (8MB, 703 files)
  Included: csrc/ (ALL CUDA kernels), fused_moe/, qwen3_5 model, _custom_ops
  Excluded: tests, benchmarks, docs, examples (not needed for reference)

Critical call chains now fully traceable:
  MoE: moe_topk_softmax_kernels.cuh → ixformer.h → fused_moe.cpp → layer
  GDN: qwen3_gated_delta_net_base.cpp → qwen3_5_gated_delta_net.cpp
  Attention: ixformer.h → xllm_paged_attention → attention.cpp
This commit is contained in:
EX Engine
2026-08-10 02:53:54 +00:00
parent 9e4fb3712f
commit 002f9879b2
2179 changed files with 494021 additions and 79 deletions
Download Patch File Download Diff File Expand all files Collapse all files

30
upstream_ref/xllm/tools/README.md Normal file
View File

@@ -0,0 +1,30 @@
# NPU Timeline Generation Guide
## Prerequisites
- Python environment
- Chrome browser (for visualization)
## Implementation Steps
### 1. Code Modification
#### Register the subscriber
Add the following at the beginning of your program:
```cpp
MsptiMetrics::register_subscriber();
```
#### Add tracing to ACLNN functions (work for msprof as well)
Insert the following macro in your ACLNN functions where you want to measure performance:
```cpp
LLM_MSTX_RANGE();
```
#### Release the subscriber
Add this at the end of your program:
```cpp
MsptiMetrics::release_subscriber();
```
### 2. Log Processing
After running your program, process the generated log file using the timeline script:
```bash
python npu_timeline.py -i custom_log.log -o custom_output.json
```
### 3. Visualization
Open Chrome browser
Navigate to: chrome://tracing
Load the generated JSON file: custom_output.json

57
upstream_ref/xllm/tools/compare_tensor.py Normal file
View File

@@ -0,0 +1,57 @@
import torch
import numpy as np
def compare_tensors(
a: torch.Tensor,
b: torch.Tensor,
tol: float = 1e-6,
verbose: bool = False
) -> int:
"""
Compare two PyTorch tensors and count the number of elements whose absolute difference
exceeds the given tolerance.
Args:
a (torch.Tensor): The first tensor to compare.
b (torch.Tensor): The second tensor to compare.
tol (float, optional): The absolute tolerance threshold. Defaults to 1e-6.
verbose (bool, optional): If True, print the indices and values of differing elements. Defaults to False.
Returns:
int: The number of elements where abs(a - b) > tol.
Raises:
ValueError: If the shapes of the input tensors do not match.
"""
# Check if tensor shapes are the same
if a.shape != b.shape:
raise ValueError(f"Shape mismatch: {a.shape} vs {b.shape}")
# Create a boolean mask where differences exceed the tolerance
diff_mask = (a - b).abs() > tol
# Count the number of differing elements
diff_count = int(diff_mask.sum().item())
# If verbose, print details of differing elements
if verbose and diff_count > 0:
indices = torch.nonzero(diff_mask, as_tuple=False)
for idx in indices:
i, j = idx[0].item(), idx[1].item()
print(
f"diff at {i},{j}: "
f"{a[i, j].item():.6f} - {b[i, j].item():.6f} = "
f"{(a[i, j] - b[i, j]).item():.6f}"
)
return diff_count
if __name__ == "__main__":
# example:
# a = torch.load("/path/to/a.pt")
# b = torch.load("/path/to/b.pt")
# diff_count = compare_tensors(a, b)
# print(f"diff count: {diff_count}")
pass

283
upstream_ref/xllm/tools/export_mtp.py Normal file
View File

@@ -0,0 +1,283 @@
"""
Export MTP layer for multiple model types (DeepSeek-V3, DeepSeek-V3.2, DeepSeek-R1, GLM4.5, GLM4.7 etc.).
The exported model can be used for speculative decoding.
Usage:
# DeepSeek V3
python3 export_mtp.py --input-dir /path/to/DeepSeek-V3 --output-dir /path/to/DeepSeek-V3-mtp
# DeepSeek V3.2
python3 export_mtp.py --input-dir /path/to/DeepSeek-V3.2 --output-dir /path/to/DeepSeek-V3.2-mtp
# DeepSeek R1
python3 export_mtp.py --input-dir /path/to/DeepSeek-R1 --output-dir /path/to/DeepSeek-R1-mtp
# GLM4 MoE
python3 export_mtp.py --input-dir /path/to/GLM-4.5-Air --output-dir /path/to/GLM-4.5-Air-mtp
"""
# adapted from https://github.com/sgl-project/sglang/blob/main/scripts/export_deepseek_nextn.py
import argparse
import json
import os
import shutil
import torch
from safetensors import safe_open
from safetensors.torch import save_file
from transformers import AutoConfig
def detect_model_type(config):
"""Detect model type from config."""
model_type = getattr(config, "model_type", "").lower()
architectures = getattr(config, "architectures", [])
# Check for DeepSeek models
# Note: DeepSeek V3, V3.2, and R1 may all have model_type="deepseek_v3" in config
# V3.2 can be distinguished by index_head_dim, index_n_heads, index_topk fields
if "deepseek" in model_type or any("deepseek" in arch.lower() for arch in architectures):
# Check for V3.2 specific fields (index_head_dim, index_n_heads, index_topk)
if hasattr(config, "index_head_dim") or hasattr(config, "index_n_heads") or hasattr(config, "index_topk"):
# V3.2 has these fields, use deepseek_v32 for MTP export
return "deepseek_v32"
else:
# V3 or R1 (both use deepseek_v3 for MTP export)
return "deepseek_v3"
# Check for GLM4
if "glm4" in model_type.lower() or any("glm4" in arch.lower() for arch in architectures):
# Check if it's MoE variant
if hasattr(config, "n_routed_experts") and getattr(config, "n_routed_experts", 0) > 0:
return "glm4_moe"
else:
return "glm4"
# Fallback: try to infer from model_type
if model_type:
return model_type
raise ValueError(f"Unable to detect model type from config. model_type={model_type}, architectures={architectures}")
def get_mtp_layer_id(config, model_type):
"""Get MTP layer ID based on model type."""
if not hasattr(config, "num_hidden_layers"):
raise ValueError("'num_hidden_layers' not found in model config.")
# For DeepSeek V3/V3.2/R1, GLM4 and GLM5, MTP layer is the last layer
if model_type in ["deepseek_v3", "deepseek_v32", "glm4_moe", "glm_moe_dsa"]:
return config.num_hidden_layers
raise ValueError(f"Unsupported model type for MTP export: {model_type}")
def get_mtp_model_type(model_type):
"""Get the MTP model type name for the output config."""
mapping = {
"deepseek_v3": "deepseek_v3_mtp", # Used for V3 and R1
"deepseek_v32": "deepseek_v32_mtp", # Used for V3.2
"glm4_moe": "glm4_moe_mtp",
"glm_moe_dsa": "glm_moe_dsa_mtp",
}
return mapping.get(model_type, f"{model_type}_mtp")
def get_mtp_architecture(model_type):
"""Get the architecture name for the output config."""
mapping = {
"deepseek_v3": "DeepseekMTPForCausalLM", # Used for V3 and R1
"deepseek_v32": "DeepseekV32MtpForCausalLM", # Used for V3.2
"glm4_moe": "Glm4MoeMtpForCausalLM",
"glm_moe_dsa": "GlmMoeDsaMtpForCausalLM",
}
return mapping.get(model_type, "MtpForCausalLM")
def update_and_save_config(config, output_dir, model_type):
"""Update and save config for MTP model."""
new_config = config.to_dict()
mtp_model_type = get_mtp_model_type(model_type)
mtp_architecture = get_mtp_architecture(model_type)
# Common updates for all models
updates = {
"num_hidden_layers": 1,
"architectures": [mtp_architecture],
"model_type": mtp_model_type,
"quantization_config": "",
}
# Keep consistent with MTP exported config requirements.
updates["first_k_dense_replace"] = 0
new_config.update(updates)
with open(os.path.join(output_dir, "config.json"), "w") as f:
json.dump(new_config, f, indent=2, ensure_ascii=False, sort_keys=True)
def copy_non_safetensors_files(input_dir, output_dir):
for filename in os.listdir(input_dir):
src_file_path = os.path.join(input_dir, filename)
if (
os.path.isfile(src_file_path)
and not filename.endswith(".safetensors")
and not filename.endswith(".safetensors.index.json")
):
dst_file_path = os.path.join(output_dir, filename)
shutil.copy2(src_file_path, dst_file_path)
print(f"All non-safetensors files have been copied to {output_dir}")
def block_dequant(
x_q_block: torch.Tensor,
x_s: torch.Tensor,
block_size: list[int],
) -> torch.Tensor:
"""This function conducts block-wise dequantization.
The inputs are block-wise quantization tensor `x_q_block`,
block-wise quantization scale and the block size.
The outputs are dequantized tensor.
"""
block_n, block_k = block_size[0], block_size[1]
n, k = x_q_block.shape
n_tiles = (n + block_n - 1) // block_n
k_tiles = (k + block_k - 1) // block_k
assert n_tiles == x_s.shape[0]
assert k_tiles == x_s.shape[1]
x_dq_block = x_q_block.to(torch.float32)
for i in range(k_tiles):
for j in range(n_tiles):
x_dq_block[
j * block_n:min((j + 1) * block_n, n),
i * block_k:min((i + 1) * block_k, k),
] *= x_s[j][i]
return x_dq_block.to(torch.bfloat16)
def export_mtp_layer_parameters(input_dir, output_dir, mtp_layer_id, model_type):
"""Export MTP layer parameters for the specified model type."""
prefix = f"model.layers.{mtp_layer_id}"
output_path = os.path.join(output_dir, "mtp_layer_parameters.safetensors")
params = {}
for filename in os.listdir(input_dir):
if not filename.endswith(".safetensors"):
continue
file_path = os.path.join(input_dir, filename)
print(f"Processing: {filename}")
try:
with safe_open(file_path, framework="pt") as f:
matching_keys = [k for k in f.keys() if (k.startswith(prefix) or k == "rot.weight")]
if not matching_keys:
print(f" No parameters starting with '{prefix}' found")
continue
for key in matching_keys:
# Handle special keys that should be at model level
if key == "rot.weight":
new_key = "model.rot.weight"
elif any(special in key for special in ["embed_tokens", "shared_head", "enorm", "hnorm", "eh_proj"]):
new_key = key.replace(prefix, "model")
else:
# Map to layer 0 for MTP model
new_key = key.replace(prefix, "model.layers.0")
params[new_key] = f.get_tensor(key)
except Exception as e:
print(f" Error processing {filename}: {str(e)}")
if params:
new_params = {}
for key, w_tensor in params.items():
# Handle block-wise quantization for DeepSeek models (V3, V3.2, R1)
if "weight_scale_inv" in key and model_type in ["deepseek_v3", "deepseek_v32"]:
weight_scale = w_tensor
weight_key = key.replace("weight_scale_inv", "weight")
if weight_key in params:
weight = params[weight_key]
weight = block_dequant(weight, weight_scale, [128, 128])
new_params[weight_key] = weight
elif key not in new_params:
new_params[key] = params[key]
params = new_params
print(f"Saving {len(params)} parameters to {output_path}")
save_file(params, output_path)
else:
print("No matching parameters found.")
raise ValueError(f"No MTP layer parameters found at layer {mtp_layer_id}")
# Update safetensors index
index_path = os.path.join(output_dir, "model.safetensors.index.json")
print(f"Updating safetensors index to {index_path}")
index_data = {"weight_map": {}}
for key in params:
index_data["weight_map"][key] = "mtp_layer_parameters.safetensors"
with open(index_path, "w") as f:
json.dump(index_data, f, indent=4)
print("All done.")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Export MTP layer parameters for multiple model types (DeepSeek-V3, DeepSeek-V3.2, DeepSeek-R1, GLM4, etc.)"
)
parser.add_argument(
"--input-dir",
type=str,
required=True,
help="Input HuggingFace model directory.",
)
parser.add_argument(
"--output-dir",
type=str,
required=True,
help="Output MTP model directory.",
)
parser.add_argument(
"--model-type",
type=str,
default=None,
help="Model type (deepseek_v3, deepseek_v32, glm4_moe, glm_moe_dsa). If not specified, will auto-detect. Note: DeepSeek V3 and R1 use 'deepseek_v3', V3.2 uses 'deepseek_v32'.",
)
args = parser.parse_args()
# Load config
config = AutoConfig.from_pretrained(args.input_dir, trust_remote_code=True)
# Detect or use specified model type
if args.model_type:
model_type = args.model_type.lower()
else:
model_type = detect_model_type(config)
print(f"Detected model type: {model_type}")
# Verify MTP support
if not hasattr(config, "num_nextn_predict_layers"):
raise ValueError("Model does not have 'num_nextn_predict_layers' attribute. This model may not support MTP.")
if config.num_nextn_predict_layers != 1:
raise ValueError(f"Only 1 MTP layer is supported, but found {config.num_nextn_predict_layers}.")
# Get MTP layer ID
mtp_layer_id = get_mtp_layer_id(config, model_type)
print(f"MTP layer ID: {mtp_layer_id}")
# Create output directory
os.makedirs(args.output_dir, exist_ok=True)
# Copy non-safetensors files
copy_non_safetensors_files(args.input_dir, args.output_dir)
# Update and save config
update_and_save_config(config, args.output_dir, model_type)
# Export MTP layer parameters
export_mtp_layer_parameters(args.input_dir, args.output_dir, mtp_layer_id, model_type)
print(f"\nMTP model exported successfully to: {args.output_dir}")

511
upstream_ref/xllm/tools/npu_timeline.py Normal file
View File

@@ -0,0 +1,511 @@
# Copyright 2016 The xLLM Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Timeline visualization for xLLM using Chrome Trace Format."""
import collections
import copy
import json
import re
import argparse
from typing import Any, Dict, List, Optional, Tuple, Union
class _ChromeTraceFormatter(object):
"""A helper class for generating traces in Chrome Trace Format."""
def __init__(self, show_memory: bool = False) -> None:
"""Constructs a new Chrome Trace formatter."""
self._show_memory = show_memory
self._events = []
self._metadata = []
def _create_event(
self,
ph: str,
category: str,
name: str,
pid: int,
tid: int,
timestamp: int,
) -> Dict[str, Union[str, int]]:
"""Creates a new Chrome Trace event.
For details of the file format, see:
https://github.com/catapult-project/catapult/blob/master/tracing/README.md
Args:
ph: The type of event - usually a single character.
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
Returns:
A JSON compatible event object.
"""
event = {}
event['ph'] = ph
event['cat'] = category
event['name'] = name
event['pid'] = pid
event['tid'] = tid
event['ts'] = timestamp
return event
def emit_pid(self, name: str, pid: int) -> None:
"""Adds a process metadata event to the trace.
Args:
name: The process name as a string.
pid: Identifier of the process as an integer.
"""
event = {}
event['name'] = 'process_name'
event['ph'] = 'M'
event['pid'] = pid
event['args'] = {'name': name}
self._metadata.append(event)
def emit_tid(self, name, pid, tid):
"""Adds a thread metadata event to the trace.
Args:
name: The thread name as a string.
pid: Identifier of the process as an integer.
tid: Identifier of the thread as an integer.
"""
event = {}
event['name'] = 'thread_name'
event['ph'] = 'M'
event['pid'] = pid
event['tid'] = tid
event['args'] = {'name': name}
self._metadata.append(event)
def emit_region(
self,
timestamp: int,
duration: int,
pid: int,
tid: int,
category: str,
name: str,
args: Dict[str, Any],
) -> None:
"""Adds a region event to the trace.
Args:
timestamp: The start timestamp of this region as a long integer.
duration: The duration of this region as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
category: The event category as a string.
name: The event name as a string.
args: A JSON-compatible dictionary of event arguments.
"""
event = self._create_event('X', category, name, pid, tid, timestamp)
event['dur'] = duration
event['args'] = args
self._events.append(event)
def emit_obj_create(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
) -> None:
"""Adds an object creation event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
"""
event = self._create_event('N', category, name, pid, tid, timestamp)
event['id'] = object_id
self._events.append(event)
def emit_obj_delete(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
) -> None:
"""Adds an object deletion event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
"""
event = self._create_event('D', category, name, pid, tid, timestamp)
event['id'] = object_id
self._events.append(event)
def emit_obj_snapshot(
self,
category: str,
name: str,
timestamp: int,
pid: int,
tid: int,
object_id: int,
snapshot: Dict[str, Any],
) -> None:
"""Adds an object snapshot event to the trace.
Args:
category: The event category as a string.
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
object_id: Identifier of the object as an integer.
snapshot: A JSON-compatible representation of the object.
"""
event = self._create_event('O', category, name, pid, tid, timestamp)
event['id'] = object_id
event['args'] = {'snapshot': snapshot}
self._events.append(event)
def emit_flow_start(
self, name: str, timestamp: int, pid: int, tid: int, flow_id: int
) -> None:
"""Adds a flow start event to the trace.
When matched with a flow end event (with the same 'flow_id') this will
cause the trace viewer to draw an arrow between the start and end events.
Args:
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
flow_id: Identifier of the flow as an integer.
"""
event = self._create_event('s', 'DataFlow', name, pid, tid, timestamp)
event['id'] = flow_id
self._events.append(event)
def emit_flow_end(
self, name: str, timestamp: int, pid: int, tid: int, flow_id: int
) -> None:
"""Adds a flow end event to the trace.
When matched with a flow start event (with the same 'flow_id') this will
cause the trace viewer to draw an arrow between the start and end events.
Args:
name: The event name as a string.
timestamp: The timestamp of this event as a long integer.
pid: Identifier of the process generating this event as an integer.
tid: Identifier of the thread generating this event as an integer.
flow_id: Identifier of the flow as an integer.
"""
event = self._create_event('t', 'DataFlow', name, pid, tid, timestamp)
event['id'] = flow_id
self._events.append(event)
def emit_counter(
self,
category: str,
name: str,
pid: int,
timestamp: int,
counter: str,
value: int,
) -> None:
"""Emits a record for a single counter.
Args:
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
counter: Name of the counter as a string.
value: Value of the counter as an integer.
"""
event = self._create_event('C', category, name, pid, 0, timestamp)
event['args'] = {counter: value}
self._events.append(event)
def emit_counters(self, category, name, pid, timestamp, counters):
"""Emits a counter record for the dictionary 'counters'.
Args:
category: The event category as a string.
name: The event name as a string.
pid: Identifier of the process generating this event as an integer.
timestamp: The timestamp of this event as a long integer.
counters: Dictionary of counter values.
"""
event = self._create_event('C', category, name, pid, 0, timestamp)
event['args'] = counters.copy()
self._events.append(event)
def format_to_string(self, pretty: bool = False) -> str:
"""Formats the chrome trace to a string.
Args:
pretty: (Optional.) If True, produce human-readable JSON output.
Returns:
A JSON-formatted string in Chrome Trace format.
"""
trace = {}
trace['traceEvents'] = self._metadata + self._events
if pretty:
return json.dumps(trace, indent=4, separators=(',', ': '))
else:
return json.dumps(trace, separators=(',', ':'))
class Timeline(object):
"""A class for visualizing execution timelines of xLLM steps."""
def __init__(self, log_file_path: str) -> None:
"""Constructs a new Timeline.
A 'Timeline' is used for visualizing the execution of a xLLM
computation. It shows the timings and concurrency of execution at
the granularity of xLLM Ops.
This class is not thread safe.
"""
self._step_stats = self.parse_log(log_file_path)
self._chrome_trace = _ChromeTraceFormatter()
self._next_pid = 0
self._marker_names = {} # id -> trace name for marker.
self._marker_end_ts = {} # id -> (deviceId, end timestamp) for marker.
self._device_pids = {} # device id -> trace pid for marker.
self._memory_pids = {} # device id -> trace pid for memory.
self._kernel_pids = {} # device id -> trace pid for kernel.
self._next_flow_id = 0
self._flow_starts = {} # tensor_name -> (timestamp, pid, tid)
def parse_log(self, log_file_path: str) -> List:
import json
step_stats = []
with open(log_file_path, 'r') as f:
lines = f.readlines()
for line in lines:
if "AscendKind" in line:
start_idx = line.find('{')
line = line[start_idx:].strip()
stats = json.loads(line)
if stats["AscendKind"] in ['MARKER', 'MEMORY', 'KERNEL']:
step_stats.append(stats)
assert len(step_stats) > 0, "step_stats is empty"
return step_stats
def _alloc_pid(self) -> int:
"""Allocate a process Id."""
pid = self._next_pid
self._next_pid += 1
return pid
def _alloc_flow_id(self) -> int:
"""Allocate a flow Id."""
flow_id = self._next_flow_id
self._next_flow_id += 1
return flow_id
def _emit_marker(
self, stats: Dict, pid: int
) -> None:
"""Generates a Chrome Trace event to show marker event.
Args:
stats: The log recording marker event.
pid: The pid assigned for the device where this marker stat ran.
"""
name = stats['name']
start = stats['timestamp'] / 1000 #microsecond
duration = stats['duration'] / 1000 #microsecond
tid = stats['streamId']
sourceKind = stats['sourceKind']
flag = stats['flag']
args = {'sourceKind': sourceKind, 'flag': flag}
self._chrome_trace.emit_region(start, duration, pid, tid, 'Marker', name, args)
def _emit_memory(
self, stats: Dict, pid: int
) -> None:
"""Generates a Chrome Trace event to show memory event.
Args:
stats: The log recording memory event.
pid: The pid assigned for the device where this memory stat ran.
"""
name = "memory_alloc" if 1 == stats['memoryKind'] else "memory_free"
start = stats['start'] / 1000 #microsecond
duration = stats['duration'] / 1000 #microsecond
tid = stats['streamId']
address = stats['address']
bytes_ = stats['bytes'] / 1024 / 1024
args = {'address': address, 'bytes': bytes_}
self._chrome_trace.emit_region(start, duration, pid, tid, 'Memory', name, args)
def _emit_kernel(
self, stats: Dict, pid: int
) -> None:
"""Generates a Chrome Trace event to show kernel event.
Args:
stats: The log recording kernel event.
pid: The pid assigned for the device where this kernel stat ran.
"""
name = stats['name'] if stats['name'] != "" else stats['type']
start = stats['start'] / 1000 #microsecond
duration = stats['duration'] / 1000 #microsecond
tid = stats['streamId']
type_ = stats['type']
args = {'type': type_}
self._chrome_trace.emit_region(start, duration, pid, tid, 'Kernel', name, args)
def _allocate_pids(self) -> None:
"""Allocate fake process ids for each device in the step_stats_pb2.StepStats."""
# Add processes in the Chrome trace to show compute and data activity.
for dev_stats in self._step_stats:
deviceId = dev_stats['deviceId']
if dev_stats['AscendKind'] == "MARKER":
if dev_stats['name'] != "":
self._marker_names[dev_stats['id']] = dev_stats['name']
else:
if dev_stats['flag'] == 32 or dev_stats['flag'] == 4: # mstxRangeEnd
if dev_stats['id'] not in self._marker_end_ts:
self._marker_end_ts[dev_stats['id']] = [(dev_stats['deviceId'], dev_stats['timestamp'])]
else:
self._marker_end_ts[dev_stats['id']].append((dev_stats['deviceId'], dev_stats['timestamp']))
if deviceId not in self._device_pids:
device_pid = self._alloc_pid()
self._device_pids[deviceId] = device_pid
if deviceId < 50:
self._chrome_trace.emit_pid('CPU Process ' + str(deviceId), device_pid)
else:
self._chrome_trace.emit_pid('NPU Device ' + str(deviceId), device_pid)
elif dev_stats['AscendKind'] == "MEMORY":
if deviceId not in self._memory_pids:
device_pid = self._alloc_pid()
self._memory_pids[deviceId] = device_pid
self._chrome_trace.emit_pid('Memory ' + str(deviceId), device_pid)
elif dev_stats['AscendKind'] == "KERNEL":
if deviceId not in self._kernel_pids:
device_pid = self._alloc_pid()
self._kernel_pids[deviceId] = device_pid
self._chrome_trace.emit_pid('Kernel ' + str(deviceId), device_pid)
else:
print("Unsupport AscendKind ", dev_stats['AscendKind'])
def _get_marker_end(self, deviceId:int, stat_id:int) -> Dict:
"""Get the end marker stats."""
for dev_stats in self._step_stats:
if 'MARKER' not in dev_stats['AscendKind']:
continue
if dev_stats['flag'] == 32 or dev_stats['flag'] == 4: # mstxRangeEnd
cur_deviceId = dev_stats['deviceId']
cur_id = dev_stats['id']
if cur_deviceId == deviceId and cur_id == stat_id:
return dev_stats
return None
def _show_marker(self, show_flow: bool = False) -> None:
"""Visualize the marker activity."""
for dev_stats in self._step_stats:
if 'MARKER' in dev_stats['AscendKind']:
deviceId = dev_stats['deviceId']
device_pid = self._device_pids[deviceId]
if dev_stats['flag'] == 32 or dev_stats['flag'] == 4: # mstxRangeEnd
continue
start_time = dev_stats['timestamp']
stats_id = dev_stats['id']
end_time = 0
# end_marker_stat = self._get_marker_end(deviceId, stats_id)
for end_ts in self._marker_end_ts[stats_id]:
cur_deviceId, cur_end_time = end_ts
if cur_deviceId == deviceId:
end_time = cur_end_time
if end_time == 0:
print(f"end marker not found: deviceId:{deviceId} id:{stats_id}")
continue
# end_time = end_marker_stat['timestamp']
dev_stats['duration'] = end_time - start_time
dev_stats['name'] = self._marker_names[stats_id]
self._emit_marker(dev_stats, device_pid)
def _show_memory(self) -> None:
"""Visualize the memory activity."""
for dev_stats in self._step_stats:
if 'MEMORY' in dev_stats['AscendKind']:
deviceId = dev_stats['deviceId']
device_pid = self._memory_pids[deviceId]
start_time = dev_stats['start']
end_time = dev_stats['end']
dev_stats['duration'] = end_time - start_time
self._emit_memory(dev_stats, device_pid)
def _show_kernel(self) -> None:
"""Visualize the kernel activity."""
for dev_stats in self._step_stats:
if 'KERNEL' in dev_stats['AscendKind']:
deviceId = dev_stats['deviceId']
device_pid = self._kernel_pids[deviceId]
start_time = dev_stats['start']
end_time = dev_stats['end']
dev_stats['duration'] = end_time - start_time
self._emit_kernel(dev_stats, device_pid)
def generate_chrome_trace_format(
self,
) -> str:
# pyformat: disable
"""Produces a trace in Chrome Trace Format.
Returns:
A JSON formatted string in Chrome Trace format.
"""
# pyformat: enable
self._allocate_pids()
self._show_marker()
self._show_memory()
self._show_kernel()
return self._chrome_trace.format_to_string(pretty=True)
def parse_args():
parser = argparse.ArgumentParser(description='Generate Chrome trace from log file')
parser.add_argument('--input', '-i', type=str, default='./node_0.log',
help='Path to input log file (default: ./log/node_0.log)')
parser.add_argument('--output', '-o', type=str, default='mspti_chrome_trace.json',
help='Path to output Chrome trace file (default: mspti_chrome_trace.json)')
return parser.parse_args()
# main
if __name__ == '__main__':
args = parse_args()
time_line = Timeline(args.input)
chrome_trace_str = time_line.generate_chrome_trace_format()
with open(args.output, 'w') as f:
f.write(chrome_trace_str)