minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity
Files
c32c4ef79cf5218c24f467aad7f3ff392a62bed7
sglang/.claude/skills/llm-torch-profiler-analysis/scripts/triage_kernel_helpers.py
T

2841 lines
92 KiB
Python
Raw Blame History

"""Internal kernel attribution helpers for triage-only torch-profiler analysis."""
from __future__ import annotations
import json
import re
from bisect import bisect_right
from collections import Counter, defaultdict
from dataclasses import dataclass, field
from functools import lru_cache
from pathlib import Path
from typing import DefaultDict, Dict, Iterable, List, Optional, Sequence, Tuple
from profile_common import (
coerce_optional_int,
contains_any_keyword,
extract_trace_events,
has_stream_marker,
is_annotation_event,
is_complete_duration_event,
is_non_kernel_trace_category,
is_trace_metadata_name,
looks_like_python_scope_name,
normalize_repo_relative_path,
normalize_text,
select_heaviest_pid,
)
CATEGORY_PATTERNS: List[Tuple[str, Tuple[str, ...]]] = [
(
"hybrid_linear",
(
"gdn",
"gated_delta",
"mamba",
"selective_scan",
"ssd",
"causal_conv",
"ssm",
),
),
(
"attention",
(
"flash_attn",
"flashattention",
"flash_attention",
"fmha",
"attention",
"mla",
"paged_attention",
"decode_attention",
),
),
(
"moe",
(
"fused_moe",
"grouped_mm",
"groupgemm",
"group_gemm",
"moe",
"expert",
"groupproblemshape",
),
),
(
"gemm",
(
"gemm",
"gemv",
"matmul",
"cublas",
"cutlass",
"wgmma",
"mma",
"bmm",
"nvjet",
),
),
(
"norm",
(
"rmsnorm",
"layernorm",
"_norm_",
" norm",
"normkernel",
),
),
("rope", ("rotary", "rope", "mrope")),
("softmax", ("softmax",)),
("activation", ("silu", "gelu", "relu", "act_and_mul", "sigmoid")),
("quantize", ("quant", "fp8", "mxfp", "nvfp4", "dequant", "cvt")),
(
"reduce_topk",
("topk", "reduce", "argmax", "argtopk", "sampling", "multinomial"),
),
(
"sampling_io",
(
"prepare_inputs",
"write_req_to",
"catarraybatched",
"prepare_next",
"copy_next",
),
),
(
"elementwise",
(
"elementwise",
"vectorized_elementwise_kernel",
"unrolled_elementwise_kernel",
"gpu_kernel_impl",
"binary_internal",
"unaryfunctor",
"add_kernel",
"sub_kernel",
"mul_kernel",
"div_",
"floor_kernel",
"log_kernel",
"neg_kernel",
),
),
]
COMMUNICATION_STRONG_KEYWORDS = (
"nccl",
"allreduce",
"all_reduce",
"reduce_scatter",
"allgather",
"all_gather",
"alltoall",
"all_to_all",
"cross_device_reduce",
"deepep",
"mooncake",
)
COMMUNICATION_WEAK_KEYWORDS = (
"broadcast",
"dispatch",
"combine",
)
MEMORY_STRONG_KEYWORDS = (
"memcpy",
"memset",
"dma",
"prefetch",
)
MEMORY_WEAK_KEYWORDS = (
"copy",
"fill",
)
COMPUTE_HINT_KEYWORDS = (
"gemm",
"gemv",
"matmul",
"cublas",
"cutlass",
"wgmma",
"mma",
"bmm",
"nvjet",
"fmha",
"attention",
"flash_attn",
"flashattention",
"flash_attention",
"grouped_mm",
"groupgemm",
"moe",
"expert",
)
NOISE_FRAME_PREFIXES = (
"threading.py(",
"multiprocessing/",
"contextlib.py(",
"torch/utils/_contextlib.py(",
"runpy.py(",
"asyncio/",
"selectors.py(",
"queue.py(",
"socket.py(",
"tqdm/_monitor.py(",
"<string>(",
"<built-in method ",
)
LOW_LEVEL_FRAME_PREFIXES = (
"triton/runtime/",
"triton/backends/",
"torch/_ops.py",
"torch/nn/modules/module.py",
)
LOW_SIGNAL_FUNCTION_TOKENS = (
"__torch_function__",
"__torch_dispatch__",
"__call__",
"_call_impl",
"_wrapped_call_impl",
)
LOW_SIGNAL_PATH_TOKENS = (
"model_executor/parameter.py:",
"model_executor/cuda_graph_runner.py:",
"compilation/cuda_graph.py:",
"pyexecutor/cuda_graph_runner.py:",
"pyexecutor/py_executor.py:",
"_torch/utils.py:",
"torch/fx/graph_module.py:",
)
@dataclass
class KernelEvent:
name: str
canonical_name: str
category: str
stage: str
pid: str
tid: str
ts: float
dur: float
external_id: Optional[int]
correlation: Optional[int] = None
@dataclass
class CpuOpEvent:
name: str
pid: str
tid: str
ts: float
dur: float
external_id: int
@dataclass
class LaunchEvent:
name: str
pid: str
tid: str
ts: float
dur: float
correlation: int
@dataclass
class PythonFrame:
name: str
normalized_name: str
pid: str
tid: str
ts: float
dur: float
python_id: Optional[int]
parent_id: Optional[int]
end_ts: float
priority: int
@dataclass
class TimedEventIndex:
events: List[object]
start_ts: List[float]
@dataclass
class FrameResolution:
location: str
stack: str
@dataclass(frozen=True)
class StageAnnotation:
stage: str
ts: float
end_ts: float
external_id: Optional[int]
is_gpu: bool
@dataclass(frozen=True)
class StageWindow:
stage: str
ts: float
end_ts: float
@dataclass
class Aggregate:
total_us: float = 0.0
count: int = 0
max_us: float = 0.0
@property
def avg_us(self) -> float:
return self.total_us / self.count if self.count else 0.0
@dataclass
class MappingSiteAggregate:
total_us: float = 0.0
count: int = 0
cpu_ops: Counter = field(default_factory=Counter)
stacks: Counter = field(default_factory=Counter)
@dataclass
class KernelRow:
name: str
category: str
aggregate: Aggregate
location: str
cpu_op: str
entry: Optional[dict]
@property
def total_us(self) -> float:
return self.aggregate.total_us
@dataclass
class FusionOpportunity:
pattern: str
status: str
confidence: str
related_us: float
evidence: str
current_locations: str
candidate_path: str
rationale: str
covered_row_keys: Tuple[Tuple[str, str, str], ...] = field(
default_factory=tuple, repr=False
)
pattern_span: int = field(default=1, repr=False)
has_active_match: bool = field(default=False, repr=False)
priority: int = field(default=0, repr=False)
subsumes: Tuple[str, ...] = field(default_factory=tuple, repr=False)
@dataclass(frozen=True)
class FusionPatternSpec:
pattern: str
candidate_path: str
active_keywords: Tuple[str, ...] = ()
split_groups: Tuple[Tuple[str, ...], ...] = ()
rationale_hint: str = ""
origin: str = "mainline"
model_include: Tuple[str, ...] = ()
model_exclude: Tuple[str, ...] = ()
min_tp_size: int = 1
require_tp: bool = False
min_share: float = 0.25
likely_share: float = 3.0
priority: int = 0
subsumes: Tuple[str, ...] = ()
FUSION_PATTERN_REGISTRY: Tuple[FusionPatternSpec, ...] = (
FusionPatternSpec(
pattern="Fused residual add + RMSNorm",
candidate_path=(
"python/sglang/srt/layers/layernorm.py"
"<br>python/sglang/srt/layers/quantization/modelslim/modelslim.py"
),
active_keywords=(
"fused_add_rmsnorm",
"gemma_fused_add_rmsnorm",
"npu_add_rms_norm",
"add_rmsnorm_bias",
),
rationale_hint=(
"Residual add plus RMSNorm already has fused implementations across"
" several backends."
),
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="FlashInfer unified allreduce_fusion",
candidate_path=(
"python/sglang/srt/layers/flashinfer_comm_fusion.py"
"<br>python/sglang/srt/layers/layernorm.py"
"<br>python/sglang/srt/layers/communicator.py"
),
active_keywords=(
"allreduce_fusion",
"fusedaddrmsnormkernel",
"flashinfer_comm_fusion.py",
),
split_groups=(
(
"cross_device_reduce",
"allreduce",
"all_reduce",
"custom_all_reduce_ops.py",
),
("rmsnorm", "layernorm", "fused_add_rmsnorm", "layernorm.py"),
),
rationale_hint=(
"FlashInfer has a TP all-reduce plus residual/RMSNorm fusion path."
),
require_tp=True,
min_tp_size=2,
min_share=0.5,
likely_share=4.0,
),
FusionPatternSpec(
pattern="AITER allreduce fusion",
candidate_path=(
"python/sglang/srt/distributed/communication_op.py"
"<br>python/sglang/srt/layers/communicator.py"
"<br>python/sglang/srt/layers/layernorm.py"
),
active_keywords=(
"tensor_model_parallel_fused_allreduce_rmsnorm",
"apply_aiter_all_reduce_fusion",
"custom_fused_ar_rms",
),
split_groups=(
("allreduce", "all_reduce", "cross_device_reduce"),
("rmsnorm", "layernorm"),
),
rationale_hint=(
"ROCm already has an AITER fused all-reduce plus RMSNorm family."
),
require_tp=True,
min_tp_size=2,
min_share=0.5,
likely_share=4.0,
),
FusionPatternSpec(
pattern="Fused activation-and-mul (SwiGLU / GeGLU)",
candidate_path="python/sglang/srt/layers/activation.py",
active_keywords=("silu_and_mul", "gelu_and_mul", "npu_swiglu"),
rationale_hint=(
"Packed MLP activation and multiply already has dedicated fused ops."
),
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="In-place QK RMSNorm",
candidate_path=(
"python/sglang/srt/models/utils.py"
"<br>python/sglang/kernels/ops/layernorm/norm.py"
),
active_keywords=("fused_inplace_qknorm", "minimaxm2rmsnormtp"),
split_groups=(("apply_qk_norm", "q_norm", "k_norm", "qknorm"),),
rationale_hint=(
"Q/K normalization already has in-place or model-specific fused"
" implementations."
),
min_share=0.3,
likely_share=2.0,
),
FusionPatternSpec(
pattern="Fused QK RMSNorm + RoPE",
candidate_path=(
"python/sglang/kernels/ops/attention/fused_qknorm_rope.py"
"<br>python/sglang/srt/models/qwen3_moe.py"
),
active_keywords=("fused_qknorm_rope", "fused_qk_norm_rope"),
split_groups=(
("apply_qk_norm", "q_norm", "k_norm", "qknorm"),
("apply_rope", "rotary", "rope", "mrope"),
),
rationale_hint=("SGLang has a fused QK-norm plus RoPE kernel family."),
min_share=0.3,
likely_share=2.0,
priority=30,
),
FusionPatternSpec(
pattern="Fused QK RoPE reshape + KV cache write",
candidate_path="python/sglang/srt/layers/attention/utils.py",
active_keywords=("fused_qk_rope_reshape_and_cache",),
split_groups=(
("rotary", "rope", "mrope"),
("reshape", "set_kv", "kv_cache", "cache write", "paged kv"),
),
rationale_hint=(
"Attention prep already has a fused RoPE plus reshape plus cache"
" write path."
),
min_share=0.4,
likely_share=2.0,
priority=40,
subsumes=("Fused RoPE + KV cache store",),
),
FusionPatternSpec(
pattern="Fused RoPE + KV cache store",
candidate_path=(
"python/sglang/kernels/ops/attention/rope.py"
"<br>python/sglang/srt/models/utils.py"
),
active_keywords=("fused_set_kv_buffer",),
split_groups=(
("rotary", "rope", "mrope"),
("set_kv_buffer", "kv cache write", "paged kv", "cache write"),
),
rationale_hint=(
"RoPE application and KV cache storage already have fused fast"
" paths in several models."
),
min_share=0.3,
likely_share=1.5,
priority=20,
),
FusionPatternSpec(
pattern="Fused decode metadata setup",
candidate_path=("python/sglang/srt/layers/attention/flashattention_backend.py"),
active_keywords=(
"normal_decode_set_metadata",
"cache_seqlens_int32",
"cu_seqlens_k",
"swa_page_table",
),
rationale_hint=(
"Decode metadata setup already has a fused Triton preparation path."
),
min_share=0.05,
likely_share=0.5,
),
FusionPatternSpec(
pattern="NSA fused metadata copy for graph replay",
candidate_path="python/sglang/kernels/ops/attention/fused_metadata_copy.py",
active_keywords=(
"fused_metadata_copy",
"fused_metadata_copy_multi",
"fused_nsa_cache_seqlens",
"fused_flashmla_metadata",
),
rationale_hint=(
"NSA replay metadata copies are already fused into one-kernel" " families."
),
min_share=0.02,
likely_share=0.2,
),
FusionPatternSpec(
pattern="DeepSeek MLA fused projection + norm + RoPE",
candidate_path=(
"python/sglang/srt/models/deepseek_common/attention_forward_methods/"
"forward_mla_fused_rope_cpu.py"
"<br>python/sglang/srt/models/deepseek_common/attention_forward_methods/"
"forward_mla_fused_rope_rocm.py"
),
active_keywords=(
"qkv_proj_with_rope_fused_weight",
"fused_qkv_a_proj_with_mqa",
"forward_absorb_fused_mla_rope",
),
split_groups=(
("mla", "qkv_a_proj", "q_a_proj"),
("qknorm", "rmsnorm", "apply_qk_norm"),
("rope", "rotary"),
),
rationale_hint=(
"DeepSeek MLA has backend-specific fused projection, norm, and"
" RoPE prep paths."
),
model_include=("deepseek", "glm"),
min_share=0.4,
likely_share=2.0,
priority=80,
subsumes=("Fused QK RMSNorm + RoPE",),
),
FusionPatternSpec(
pattern="Fused QK RoPE concat + MLA cache write",
candidate_path=(
"python/sglang/srt/layers/rocm_linear_utils.py"
"<br>python/sglang/srt/models/deepseek_common/attention_forward_methods/"
"forward_mla.py"
),
active_keywords=("fused_qk_rope_cat_and_cache_mla", "set_mla_kv_buffer"),
split_groups=(
("mla", "rope", "rotary"),
("cache", "kv_buffer", "concat"),
),
rationale_hint=(
"MLA RoPE packing and cache write already have fused backend paths."
),
model_include=("deepseek", "glm"),
min_share=0.3,
likely_share=1.5,
priority=85,
subsumes=("Fused RoPE + KV cache store",),
),
FusionPatternSpec(
pattern="Qwen3 decode fused QK norm + 3D mRoPE + KV cache write",
candidate_path="python/sglang/srt/models/qwen3.py",
active_keywords=("fused_qk_norm_mrope_3d_cache_pts_quant_shuffle",),
split_groups=(
("apply_qk_norm", "q_norm", "k_norm", "qknorm"),
("mrope", "3d rope", "rotary"),
("cache", "kv_buffer", "paged kv", "cache write"),
),
rationale_hint=(
"Qwen3-style decode already has a fused QK-norm plus 3D mRoPE plus"
" cache-write path."
),
model_include=("qwen3",),
model_exclude=("qwen3.5", "qwen3_5"),
min_share=0.4,
likely_share=2.0,
priority=90,
subsumes=(
"Fused QK RMSNorm + RoPE",
"Fused QK RoPE reshape + KV cache write",
"Fused RoPE + KV cache store",
),
),
FusionPatternSpec(
pattern="Fused MoE router / top-k / softcapping",
candidate_path="python/sglang/srt/layers/moe/router.py",
active_keywords=("fusedmoerouter", "fused_moe_router"),
split_groups=(
("router", "gate", "router logits"),
("topk", "softmax", "softcap", "tanh"),
),
rationale_hint=(
"MoE routing already has fused router, softcap, and top-k kernels."
),
min_share=0.3,
likely_share=1.5,
priority=30,
),
FusionPatternSpec(
pattern="Fused MoE grouped-topk / gate kernels",
candidate_path="python/sglang/srt/layers/moe/topk.py",
active_keywords=(
"fused_topk_deepseek",
"moe_fused_gate",
"aiter_fused_topk",
"kimi_k2_moe_fused_gate",
),
split_groups=(
("grouped_topk", "topk", "biased_grouped_topk"),
("gate", "router", "renorm", "routed scaling"),
),
rationale_hint=(
"Grouped-topk, bias handling, and routed scaling already have fused"
" gate kernels."
),
min_share=0.3,
likely_share=1.5,
priority=50,
subsumes=("Fused MoE router / top-k / softcapping",),
),
FusionPatternSpec(
pattern="Qwen-style shared-expert append into routed top-k output",
candidate_path=(
"python/sglang/srt/models/qwen2_moe.py"
"<br>python/sglang/srt/layers/moe/moe_runner/triton_utils/"
"fused_moe_triton_kernels.py"
),
active_keywords=(
"_append_shared_to_topk_output",
"fused_append_shared_experts_with_weights",
"_fused_append_shared_experts_with_weights_kernel",
),
split_groups=(
("_append_shared_to_topk_output", "topk", "grouped_topk"),
("shared_expert", "shared_expert_gate", "sigmoid"),
),
rationale_hint=(
"Qwen-style shared experts can already be appended into routed top-k"
" output in one Triton prep kernel before fused MoE execution."
),
min_share=0.05,
likely_share=0.5,
priority=55,
),
FusionPatternSpec(
pattern="Fused MoE sum + all-reduce",
candidate_path=("python/sglang/srt/layers/moe/fused_moe_triton/fused_moe.py"),
active_keywords=("fuse_sum_all_reduce", "enable_fused_moe_sum_all_reduce"),
split_groups=(
("fused_moe", "expert", "moe"),
("allreduce", "all_reduce", "cross_device_reduce"),
),
rationale_hint=(
"The second MoE GEMM already has a fused sum-plus-all-reduce path."
),
require_tp=True,
min_tp_size=2,
min_share=0.4,
likely_share=2.0,
),
FusionPatternSpec(
pattern="Fused MoE activation + quant / re-quant",
candidate_path=(
"python/sglang/srt/layers/moe/ep_moe/kernels.py"
"<br>python/sglang/kernels/ops/quantization/nvfp4_gemm_swiglu_nvfp4_quant.py"
"<br>python/sglang/srt/layers/moe/cutlass_w4a8_moe.py"
),
active_keywords=(
"silu_and_mul_scaled_fp4",
"npu_dequant_swiglu_quant",
"swiglu_quant",
),
split_groups=(
("silu", "gelu", "act_and_mul"),
("quant", "fp8", "mxfp", "nvfp4", "dequant"),
),
rationale_hint=(
"Quantized MoE backends already fuse activation with re-quantization."
),
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="DeepSeek comm-prep fused RMSNorm + quant / flatten-quant",
candidate_path=(
"python/sglang/srt/layers/communicator.py"
"<br>python/sglang/srt/models/deepseek_common/attention_forward_methods/"
"forward_mla.py"
"<br>python/sglang/srt/models/deepseek_common/attention_forward_methods/"
"forward_mha.py"
),
active_keywords=(
"fused_rms_fp8_group_quant",
"fused_rms_mxfp4_quant",
"fused_flatten_fp8_group_quant",
"fused_flatten_mxfp4_quant",
),
split_groups=(
("rmsnorm", "layernorm", "flatten"),
("fp8", "mxfp4", "quant"),
),
rationale_hint=(
"DeepSeek comm preparation already fuses norm or flatten work with"
" quantization."
),
model_include=("deepseek", "glm"),
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="NSA fused top-k transform / page-table build",
candidate_path="python/sglang/srt/layers/attention/nsa_backend.py",
active_keywords=(
"fast_topk_transform_fused",
"fast_topk_transform_ragged_fused",
),
rationale_hint=(
"NSA top-k metadata preparation already has fused transform kernels."
),
min_share=0.05,
likely_share=0.3,
),
FusionPatternSpec(
pattern="NSA fused quantize + indexed K-cache store",
candidate_path=(
"python/sglang/kernels/ops/attention/fused_store_index_cache.py"
"<br>python/sglang/srt/layers/attention/nsa/nsa_indexer.py"
),
active_keywords=("fused_store_index_k_cache",),
split_groups=(
("act_quant", "quant", "scale_buffer"),
("index_k", "cache", "store"),
),
rationale_hint=(
"NSA already has a fused quantize-and-indexed-store kernel family."
),
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="Fused sampling temperature + softmax",
candidate_path=(
"python/sglang/srt/layers/fused_sampling.py"
"<br>python/sglang/srt/layers/sampler.py"
),
active_keywords=("fused_temperature_softmax",),
split_groups=(
("temperature", "temp_scale"),
("softmax", "sampling"),
),
rationale_hint=(
"Decode-time sampling already has fused temperature and softmax" " kernels."
),
min_share=0.05,
likely_share=0.5,
),
FusionPatternSpec(
pattern="Fused logit softcap",
candidate_path=(
"python/sglang/srt/layers/elementwise.py"
"<br>python/sglang/srt/layers/logits_processor.py"
),
active_keywords=("fused_softcap", "final_logit_softcapping"),
rationale_hint=(
"Logit softcap math already has dedicated fused elementwise kernels."
),
min_share=0.02,
likely_share=0.2,
),
FusionPatternSpec(
pattern="PR #20667 Qwen3.5 fused QK norm + RoPE + KV cache write",
candidate_path=(
"PR #20667"
"<br>python/sglang/srt/models/qwen3_5.py"
"<br>python/sglang/srt/models/utils.py"
),
active_keywords=(
"fused_qk_norm_rope_cache_pts_quant_shuffle",
"fused_qk_norm_mrope_3d_cache_pts_quant_shuffle",
),
split_groups=(
("apply_qk_norm", "qknorm", "q_norm", "k_norm"),
("rotary", "rope", "mrope"),
("cache", "kv_buffer", "cache write"),
),
rationale_hint=(
"Open SGLang ROCm PR wires a fused QK-norm plus RoPE plus KV-cache"
" family for Qwen3.5."
),
origin="inflight",
model_include=("qwen3.5", "qwen3_5"),
min_share=0.4,
likely_share=2.0,
priority=100,
subsumes=(
"Fused QK RMSNorm + RoPE",
"Fused QK RoPE reshape + KV cache write",
"Fused RoPE + KV cache store",
),
),
FusionPatternSpec(
pattern="PR #22392 CUTLASS FP8 scaled MM replacing nvjet",
candidate_path=(
"PR #22392"
"<br>python/sglang/kernels/aot/python/sgl_kernel/gemm.py"
"<br>python/sglang/srt/layers/quantization/fp8_utils.py"
),
active_keywords=("cutlass_scaled_mm", "fp8_scaled_mm"),
split_groups=(
("nvjet", "_scaled_mm"),
("memset", "memcpy128"),
),
rationale_hint=(
"Open SGLang PR replaces nvjet FP8 GEMM with CUTLASS to remove"
" memset bubbles and extra copies."
),
origin="inflight",
min_share=0.2,
likely_share=1.0,
priority=90,
),
FusionPatternSpec(
pattern="SGLang LTX2 fused Ada values",
candidate_path=(
"PR #29390"
"<br>python/sglang/kernels/ops/diffusion/triton/ltx2_ada_values.py"
"<br>python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py"
),
active_keywords=(
"ltx2_ada_values9",
"ltx2_ada_values",
"LTX2TransformerBlock",
),
split_groups=(
("scale_shift_table", "timestep", "reshape"),
("get_ada_values", "ada", "adaln"),
("slice", "split", "unbind"),
),
rationale_hint=(
"SGLang mainline fuses LTX-2.3 Ada value materialization for"
" video/audio streams; split Ada table add/reshape/slice ladders"
" should be checked against this diffusion Triton kernel first."
),
origin="upstream",
model_include=("ltx", "ltx-2", "ltx2"),
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="SGLang LTX2 residual-gate add CUDA fast path",
candidate_path=(
"PR #29361"
"<br>python/sglang/kernels/ops/diffusion/residual_gate_add.py"
"<br>python/sglang/kernels/jit/csrc/diffusion/residual_gate_add.cuh"
"<br>python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py"
),
active_keywords=(
"diffusion_residual_gate_add",
"residual_gate_add",
"_ltx2_residual_gate_add",
),
split_groups=(
("add", "mul", "gate"),
("residual", "update", "gate"),
("hidden_states", "attn_hidden_states", "gate"),
),
rationale_hint=(
"SGLang mainline fuses LTX2 residual + update * gate sites into"
" a CUDA custom op; split add/mul gate ladders should be checked"
" against this path before proposing a new diffusion elementwise"
" fusion."
),
origin="upstream",
model_include=("ltx", "ltx-2", "ltx2"),
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TokenSpeed CuTe DSL MLA prefill / decode",
candidate_path=(
"python/tokenspeed/runtime/layers/attention/backends/tokenspeed_mla.py"
"<br>tokenspeed-mla/python/tokenspeed_mla/mla_decode.py"
"<br>tokenspeed-mla/python/tokenspeed_mla/mla_prefill.py"
"<br>tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/"
"tokenspeed_mla/__init__.py"
),
active_keywords=(
"tokenspeed_mla_decode",
"tokenspeed_mla_prefill",
"BlackwellMultiHeadLatentAttentionForward",
),
split_groups=(
("mla", "flashmla", "attention", "fmha"),
("prefill", "decode", "verify"),
("fp8", "kv_cache", "page_table"),
),
rationale_hint=(
"TokenSpeed ships Blackwell CuTe DSL MLA prefill/decode kernels;"
" split MLA support kernels should be checked against backend"
" selection before being called novel."
),
origin="upstream",
model_include=("deepseek", "kimi", "qwen3.5", "qwen3_5"),
min_share=0.4,
likely_share=2.0,
),
FusionPatternSpec(
pattern="TokenSpeed MLA KV pack + FP8 quantize",
candidate_path=(
"tokenspeed-mla/python/tokenspeed_mla/mla_kv_pack_quantize_fp8.py"
"<br>tokenspeed-kernel/python/tokenspeed_kernel/ops/attention/"
"tokenspeed_mla/__init__.py"
),
active_keywords=(
"_mla_kv_pack_quantize_fp8_kernel",
"mla_kv_pack_quantize_fp8",
),
split_groups=(
("k_nope", "k_pe", "cat", "concat", "pack"),
("quant", "fp8", "float8"),
("v", "kv", "cache"),
),
rationale_hint=(
"TokenSpeed fuses MLA K/V pack, concat, and FP8 quantization into"
" one Triton kernel for chunked prefill."
),
origin="upstream",
model_include=("deepseek", "kimi", "qwen3.5", "qwen3_5"),
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TokenSpeed fused top-k + top-p sampling",
candidate_path=(
"tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/" # codespell:ignore thirdparty
"fused_topk_topp.py"
"<br>tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/" # codespell:ignore thirdparty
"csrc/fused_topk_topp/fused_topk_topp.cu"
),
active_keywords=("fused_topk_topp", "fused_topk_topp_renorm"),
split_groups=(
("topk", "top_k"),
("topp", "top_p"),
("sampling", "renorm", "softmax"),
),
rationale_hint=(
"TokenSpeed has a fused top-k/top-p renormalization path for"
" decode sampling."
),
origin="upstream",
min_share=0.1,
likely_share=0.8,
),
FusionPatternSpec(
pattern="TokenSpeed persistent lm_head GEMM",
candidate_path=(
"tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/" # codespell:ignore thirdparty
"lm_head_gemm.py"
"<br>tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cuda/" # codespell:ignore thirdparty
"csrc/lm_head_gemm.cu"
),
active_keywords=("lm_head_gemm",),
split_groups=(
("lm_head", "logits", "vocab"),
("gemm", "matmul", "linear"),
),
rationale_hint=(
"TokenSpeed has a shape-gated persistent lm_head GEMM path; visible"
" lm_head matmul ladders should be compared against it."
),
origin="upstream",
model_include=("kimi", "qwen"),
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TokenSpeed NVFP4 GEMM + SwiGLU + quant",
candidate_path=(
"tokenspeed-kernel/python/tokenspeed_kernel/thirdparty/cute_dsl/" # codespell:ignore thirdparty
"nvfp4_gemm_swiglu_nvfp4_quant.py"
),
active_keywords=("nvfp4_gemm_swiglu_nvfp4_quant",),
split_groups=(
("gemm", "nvfp4", "fp4"),
("swiglu", "silu", "activation", "mul"),
("quant", "scale", "sfc"),
),
rationale_hint=(
"TokenSpeed's CuTe DSL kernel fuses NVFP4 GEMM, SwiGLU, and"
" optional output quantization in one expert-style path."
),
origin="upstream",
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin Attention + Quantization",
candidate_path=(
"vllm/compilation/passes/fusion/attn_quant_fusion.py"
"<br>vllm/v1/attention/ops/merge_attn_states.py"
"<br>vllm/csrc/attention/merge_attn_states.cu"
"<br>vllm/docs/design/fusions.md"
),
active_keywords=(
"merge_attn_states",
"attn_quant_fusion",
"output_scale",
"output_group_scale",
),
split_groups=(
("attention", "flash_attn", "flashattention", "mla"),
("quant", "fp8", "nvfp4", "group_scale"),
),
rationale_hint=(
"vLLM combines attention merge with attention-epilogue quantization."
),
origin="upstream",
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin DSV3.2 fused indexer projections",
candidate_path=(
"vllm/model_executor/models/deepseek_v2.py"
"<br>vllm/model_executor/models/deepseek_mtp.py"
),
active_keywords=("wk_weights_proj",),
split_groups=(
("wk_weights_proj", "wk", "weights_proj"),
("mergedcolumnparallellinear", "gemm", "matmul"),
),
rationale_hint=(
"vLLM already fuses the paired `wk` and `weights_proj` indexer"
" projections into one DSV3.2 linear family."
),
origin="upstream",
min_share=0.2,
likely_share=1.0,
),
FusionPatternSpec(
pattern="vLLM-origin RMSNorm + Quantization",
candidate_path=(
"vllm/compilation/passes/fusion/rms_quant_fusion.py"
"<br>vllm/docs/design/fusions.md"
),
active_keywords=(
"fused_add_rms_norm_static_fp8_quant",
"rms_quant_fusion",
"norm_quant",
),
split_groups=(
("rmsnorm", "layernorm", "fused_add_rms_norm"),
("quant", "fp8", "fp4", "per-group"),
),
rationale_hint=(
"vLLM already has a compile-time norm-plus-quant fusion family."
),
origin="upstream",
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin SiLU+Mul + Quantization",
candidate_path=(
"vllm/compilation/passes/fusion/act_quant_fusion.py"
"<br>vllm/docs/design/fusions.md"
),
active_keywords=(
"silu_mul_quant_fp4",
"fused_silu_mul_block_quant",
"act_quant_fusion",
),
split_groups=(
("silu", "gelu", "act_and_mul"),
("quant", "fp8", "fp4", "block_quant"),
),
rationale_hint=("vLLM has an activation-plus-quant fusion family."),
origin="upstream",
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin DSV3 router GEMM",
candidate_path=(
"vllm/model_executor/layers/fused_moe/router/gate_linear.py"
"<br>vllm/csrc/moe/dsv3_router_gemm_entry.cu"
),
active_keywords=("dsv3_router_gemm", "fp32_router_gemm"),
split_groups=(
("router", "gate", "router logits"),
("gemm", "matmul", "cublas", "cutlass"),
),
rationale_hint=(
"vLLM has a specialized DeepSeek router GEMM family for small"
" decode batches."
),
origin="upstream",
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin GPT-OSS router GEMM",
candidate_path=(
"vllm/_custom_ops.py"
"<br>vllm/model_executor/layers/fused_moe/router/gate_linear.py"
"<br>vllm/csrc/moe/gpt_oss_router_gemm.cu"
),
active_keywords=("gpt_oss_router_gemm",),
split_groups=(
("router", "gate", "router logits", "gpt_oss"),
("gemm", "matmul", "cublas", "cutlass"),
),
rationale_hint=("vLLM has a GPT-OSS-specific router GEMM path."),
origin="upstream",
model_include=("gpt-oss", "gpt_oss"),
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM-origin DeepSeek min-latency fused QKV-A projection",
candidate_path=(
"vllm/model_executor/models/deepseek_v2.py"
"<br>vllm/csrc/dsv3_fused_a_gemm.cu"
),
active_keywords=("dsv3_fused_a_gemm", "fused_qkv_a_proj"),
split_groups=(
("q_a_proj", "kv_a_proj", "weights_proj"),
("gemm", "matmul", "cutlass", "cublas"),
),
rationale_hint=(
"vLLM has a fused DeepSeek QKV-A projection family for decode"
" latency reduction."
),
origin="upstream",
model_include=("deepseek", "glm"),
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="PR #38621 fused QK norm + RoPE + cache + quant",
candidate_path=(
"PR #38621"
"<br>vllm/csrc/fused_qk_norm_rope_cache_quant.cu"
"<br>vllm/compilation/passes/fusion/qk_norm_rope_cache_quant_fusion.py"
),
active_keywords=("fused_qk_norm_rope_cache_quant",),
split_groups=(
("qknorm", "q_norm", "k_norm"),
("rope", "rotary", "mrope"),
("cache", "kv_buffer", "cache write"),
("quant", "fp8", "nvfp4"),
),
rationale_hint=(
"Open vLLM PR covers QK-norm plus RoPE plus cache plus quant as"
" one fusion family."
),
origin="inflight",
min_share=0.4,
likely_share=2.0,
priority=100,
subsumes=("vLLM-origin Attention + Quantization",),
),
FusionPatternSpec(
pattern="vLLM-origin MiniMax allreduce_rms kernels",
candidate_path="vllm/model_executor/models/minimax_m2.py",
active_keywords=("minimax_allreduce_rms", "minimax_allreduce_rmsnorm"),
split_groups=(
("q_norm", "k_norm", "rmsnorm", "minimax"),
("allreduce", "all_reduce", "cross_device_reduce"),
),
rationale_hint=(
"vLLM includes the TRTLLM-derived MiniMax allreduce-plus-RMSNorm"
" kernel family."
),
origin="upstream",
model_include=("minimax",),
min_share=0.3,
likely_share=1.5,
),
FusionPatternSpec(
pattern="vLLM fused residual add + RMSNorm",
candidate_path=(
"vllm/_custom_ops.py"
"<br>vllm/compilation/passes/fusion/rms_quant_fusion.py"
),
active_keywords=(
"fused_add_rms_norm",
"fused_add_rms_norm_static_fp8_quant",
),
rationale_hint=(
"vLLM exposes fused residual-add-plus-RMSNorm kernels and matching"
" compile-time hooks."
),
origin="upstream",
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="vLLM fused activation-and-mul",
candidate_path=(
"vllm/_custom_ops.py"
"<br>vllm/compilation/passes/fusion/act_quant_fusion.py"
),
active_keywords=(
"silu_and_mul",
"silu_and_mul_quant",
"silu_and_mul_per_block_quant",
"act_and_mul",
),
rationale_hint=(
"vLLM ships fused activation-and-multiply kernels plus quantized"
" variants for the MLP epilogue."
),
origin="upstream",
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TensorRT-LLM FlashInfer residual add + RMSNorm",
candidate_path=(
"tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py"
"<br>tensorrt_llm/_torch/modules/rms_norm.py"
"<br>tensorrt_llm/_torch/auto_deploy/transform/library/fused_add_rms_norm.py"
),
active_keywords=(
"flashinfer_fused_add_rmsnorm",
"flashinfer_gemma_fused_add_rmsnorm",
"flashinfer::norm::FusedAddRMSNormKernel",
"FusedAddRMSNormKernel",
"auto_deploy::flashinfer_fused_add_rms_norm_inplace",
),
rationale_hint=(
"TensorRT-LLM exposes a FlashInfer fused residual-add plus RMSNorm"
" family, including AutoDeploy rewrites."
),
origin="upstream",
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TensorRT-LLM Triton fused residual add + RMSNorm + FP8 quant",
candidate_path=(
"tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/"
"triton_fused_add_rms_norm_quant_fp8.py"
"<br>tensorrt_llm/_torch/auto_deploy/transform/library/"
"fuse_rmsnorm_quant_fp8.py"
),
active_keywords=(
"triton_fused_add_rms_norm_quant_fp8",
"fuse_rmsnorm_quant_fp8",
),
rationale_hint=(
"TensorRT-LLM mainline has a Triton residual-add plus RMSNorm plus"
" FP8-quant family in AutoDeploy."
),
origin="upstream",
min_share=0.2,
likely_share=1.0,
priority=20,
),
FusionPatternSpec(
pattern="TensorRT-LLM FlashInfer RMSNorm family",
candidate_path=(
"tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py"
"<br>tensorrt_llm/_torch/modules/rms_norm.py"
"<br>tensorrt_llm/_torch/auto_deploy/custom_ops/normalization/rms_norm.py"
),
active_keywords=(
"flashinfer_rmsnorm",
"flashinfer_gemma_rmsnorm",
"auto_deploy::flashinfer_rms_norm",
),
rationale_hint=(
"TensorRT-LLM lowers RMSNorm-style ladders to FlashInfer kernels"
" and AutoDeploy custom ops."
),
origin="upstream",
min_share=0.1,
likely_share=1.0,
),
FusionPatternSpec(
pattern="TensorRT-LLM FlashInfer activation / gate epilogues",
candidate_path=(
"tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py"
"<br>tensorrt_llm/_torch/auto_deploy/transform/library/fuse_silu_mul.py"
"<br>tensorrt_llm/_torch/models/modeling_gemma3.py"
),
active_keywords=(
"flashinfer_silu_and_mul",
"flashinfer_gelu_tanh_and_mul",
"auto_deploy::silu_and_mul",
),
rationale_hint=(
"TensorRT-LLM already rewrites gate activation plus multiply"
" ladders into FlashInfer epilogue kernels."
),
origin="upstream",
min_share=0.1,
likely_share=1.0,
),
)
def short_name(name: str, max_len: int = 96) -> str:
text = normalize_text(name)
if len(text) <= max_len:
return text
return text[: max_len - 3] + "..."
@lru_cache(maxsize=65536)
def canonicalize_name(name: str) -> str:
text = normalize_text(name)
text = re.sub(r"0x[0-9a-fA-F]+", "0xADDR", text)
if text.startswith("void ") and text.endswith(")"):
depth = 0
split_idx: Optional[int] = None
for idx in range(len(text) - 1, -1, -1):
char = text[idx]
if char == ")":
depth += 1
elif char == "(":
depth -= 1
if depth == 0:
split_idx = idx
break
if split_idx is not None:
text = text[:split_idx]
return text
@lru_cache(maxsize=65536)
def classify_kernel(name: str) -> str:
# Keep the matching order explicit: strong communication/memory signals win
# first, then we fall back to weaker category hints.
lowered = name.lower()
if contains_any_keyword(lowered, COMMUNICATION_STRONG_KEYWORDS):
return "communication"
if contains_any_keyword(lowered, MEMORY_STRONG_KEYWORDS):
return "memory"
looks_compute_like = contains_any_keyword(lowered, COMPUTE_HINT_KEYWORDS)
if contains_any_keyword(lowered, MEMORY_WEAK_KEYWORDS) and not looks_compute_like:
return "memory"
for category, keywords in CATEGORY_PATTERNS:
if contains_any_keyword(lowered, keywords):
return category
if (
contains_any_keyword(lowered, COMMUNICATION_WEAK_KEYWORDS)
and not looks_compute_like
):
return "communication"
return "other"
@lru_cache(maxsize=65536)
def normalize_source_location(name: str) -> str:
text = normalize_text(name)
match = re.match(r"(?P<path>.+?)\((?P<line>\d+)\): (?P<func>.+)$", text)
if not match:
return text
path = normalize_repo_relative_path(match.group("path"))
return f"{path}:{match.group('line')} {match.group('func')}"
def source_location_priority(location: str) -> int:
text = str(location).strip()
if not text or text == "unresolved":
return -100
penalty = 80 if is_low_signal_source_location(text) else 0
if text.startswith("python/sglang/"):
return 300 - penalty
if text.startswith("sglang/"):
return 290 - penalty
if text.startswith("vllm/"):
return 285 - penalty
if text.startswith("python/tokenspeed/") or text.startswith("tokenspeed/"):
return 283 - penalty
if text.startswith("tensorrt_llm/"):
return 280 - penalty
if text.startswith("sgl_kernel/"):
return 260 - penalty
if text.startswith("python/"):
return 180 - penalty
if text.startswith("torch/") or "/torch/" in text:
return 20
if ".py:" in text:
return 120 - penalty
return 0
def is_preferred_source_location(location: str) -> bool:
text = str(location).strip()
return (
text.startswith("python/sglang/")
or text.startswith("sglang/")
or text.startswith("vllm/")
or text.startswith("python/tokenspeed/")
or text.startswith("tokenspeed/")
or text.startswith("tensorrt_llm/")
or text.startswith("sgl_kernel/")
)
def extract_preferred_stack_location(stack: Optional[str]) -> Optional[str]:
if not stack:
return None
parts = [str(part).strip() for part in str(stack).split("->")]
ranked: List[Tuple[int, int, str]] = []
for index, part in enumerate(parts):
normalized = normalize_source_location(part)
priority = source_location_priority(normalized)
if priority <= 0:
continue
ranked.append((priority, index, normalized))
if not ranked:
return None
ranked.sort(key=lambda item: (item[0], item[1]), reverse=True)
return ranked[0][2]
def site_display_location(site: dict) -> str:
location = str(site.get("location") or "unresolved").strip()
if is_preferred_source_location(location) and not is_low_signal_source_location(
location
):
return location
stack_location = extract_preferred_stack_location(site.get("stack"))
if stack_location:
return stack_location
return location
def choose_best_location(locations: Dict[str, MappingSiteAggregate]) -> str:
if not locations:
return "unresolved"
ranked = sorted(
locations.items(),
key=lambda pair: (
source_location_priority(pair[0]),
pair[1].total_us,
pair[1].count,
),
reverse=True,
)
return ranked[0][0]
@lru_cache(maxsize=65536)
def frame_priority(frame_name: str) -> int:
raw_text = str(frame_name).strip()
normalized_text = normalize_source_location(raw_text)
penalty = 80 if is_low_signal_source_location(normalized_text) else 0
if raw_text.startswith(NOISE_FRAME_PREFIXES):
return -20
if normalized_text.startswith("python/sglang/"):
return 300 - penalty
if normalized_text.startswith("sglang/"):
return 290 - penalty
if normalized_text.startswith("vllm/"):
return 285 - penalty
if normalized_text.startswith("python/tokenspeed/") or normalized_text.startswith(
"tokenspeed/"
):
return 283 - penalty
if normalized_text.startswith("tensorrt_llm/"):
return 280 - penalty
if normalized_text.startswith("sgl_kernel/"):
return 260 - penalty
if normalized_text.startswith("triton_kernels/"):
return 220 - penalty
if normalized_text.startswith(LOW_LEVEL_FRAME_PREFIXES):
return 0
if raw_text.startswith("/data/") or raw_text.startswith("/Users/"):
if "/sglang/" in raw_text:
return 120
if "/vllm/" in raw_text:
return 118
if "/tokenspeed/" in raw_text or "/TokenSpeed/" in raw_text:
return 117
if "/TensorRT-LLM/" in raw_text or "/tensorrt_llm/" in raw_text:
return 116
return 100
if ".py(" in raw_text and "/sglang/" in raw_text:
return 110
if ".py(" in raw_text and "/vllm/" in raw_text:
return 108
if ".py(" in raw_text and (
"/tokenspeed/" in raw_text or "/TokenSpeed/" in raw_text
):
return 107
if ".py(" in raw_text and (
"/TensorRT-LLM/" in raw_text or "/tensorrt_llm/" in raw_text
):
return 106
if ".py:" in normalized_text and (
"site-packages" in raw_text or normalized_text.startswith("torch/")
):
return 45
if ".py:" in normalized_text:
return 35
if raw_text.startswith("<built-in method "):
return -10
return 0
@lru_cache(maxsize=65536)
def is_low_signal_source_location(location: str) -> bool:
lowered = str(location).strip().lower()
if not lowered:
return False
return any(token in lowered for token in LOW_SIGNAL_FUNCTION_TOKENS) or any(
token in lowered for token in LOW_SIGNAL_PATH_TOKENS
)
def stage_label(stage: str) -> str:
if stage == "extend":
return "extend/prefill"
return stage
def stage_aliases(stage: str) -> List[str]:
if stage == "extend":
return ["extend", "prefill", "all"]
if stage == "prefill":
return ["prefill", "extend", "all"]
if stage == "decode":
return ["decode", "all"]
return [stage, "all"]
def escape_md_cell(text: str) -> str:
return str(text).replace("|", "\\|").replace("\n", "<br>")
def pct(part: float, whole: float) -> float:
return 100.0 * part / whole if whole else 0.0
def format_ms(value_us: float) -> str:
return f"{value_us / 1000.0:.2f} ms"
@lru_cache(maxsize=16384)
def _is_cuda_launch_event_cached(name: str, cat: str) -> bool:
lowered_name = normalize_text(name).lower()
lowered_cat = normalize_text(cat).lower()
if lowered_cat not in {"cuda_runtime", "cuda_driver"}:
return False
return "launch" in lowered_name
def is_cuda_launch_event(name: str, cat: str) -> bool:
return _is_cuda_launch_event_cached(str(name), str(cat))
def is_gpu_kernel_event(event: dict) -> bool:
# Be conservative here: first drop trace metadata / Python scopes /
# annotations, then only accept entries with clear GPU-kernel markers.
if not is_complete_duration_event(event):
return False
name = normalize_text(event.get("name", ""))
if is_trace_metadata_name(name):
return False
cat = normalize_text(event.get("cat", "")).lower()
args = event.get("args") or {}
if is_non_kernel_trace_category(cat):
return False
if is_annotation_event(name, cat):
return False
if "kernel" in cat or cat.startswith("gpu_"):
return True
if looks_like_python_scope_name(name):
return False
return has_stream_marker(args)
def infer_stage_from_annotation_name(name: str) -> Optional[str]:
lowered = normalize_text(name).lower()
if not lowered:
return None
if "generation_1" in lowered or "decode" in lowered:
return "decode"
if "generation_0" in lowered or "prefill" in lowered:
return "extend"
return None
def build_stage_annotations(
raw_events: Sequence[dict],
) -> Tuple[
Dict[int, StageAnnotation],
List[StageWindow],
List[StageWindow],
]:
by_external_id: Dict[int, StageAnnotation] = {}
gpu_annotations: List[StageAnnotation] = []
cpu_annotations: List[StageAnnotation] = []
def should_replace(current: StageAnnotation, candidate: StageAnnotation) -> bool:
if candidate.is_gpu != current.is_gpu:
return candidate.is_gpu
return (candidate.end_ts - candidate.ts) > (current.end_ts - current.ts)
for event in raw_events:
if not is_complete_duration_event(event):
continue
category = normalize_text(event.get("cat", "")).lower()
if category not in {"user_annotation", "gpu_user_annotation"}:
continue
stage = infer_stage_from_annotation_name(str(event.get("name", "")))
if not stage:
continue
annotation = StageAnnotation(
stage=stage,
ts=float(event.get("ts", 0.0)),
end_ts=float(event.get("ts", 0.0)) + float(event.get("dur", 0.0)),
external_id=coerce_optional_int(
(event.get("args") or {}).get("External id")
),
is_gpu=(category == "gpu_user_annotation"),
)
if annotation.external_id is not None:
existing = by_external_id.get(annotation.external_id)
if existing is None or should_replace(existing, annotation):
by_external_id[annotation.external_id] = annotation
if annotation.is_gpu:
gpu_annotations.append(annotation)
else:
cpu_annotations.append(annotation)
gpu_annotations.sort(key=lambda item: (item.ts, item.end_ts))
cpu_annotations.sort(key=lambda item: (item.ts, item.end_ts))
return (
by_external_id,
merge_stage_windows(gpu_annotations),
merge_stage_windows(cpu_annotations),
)
def merge_stage_windows(annotations: Sequence[StageAnnotation]) -> List[StageWindow]:
merged: List[StageWindow] = []
for annotation in annotations:
if (
merged
and merged[-1].stage == annotation.stage
and annotation.ts <= merged[-1].end_ts + 1e-3
):
merged[-1] = StageWindow(
stage=merged[-1].stage,
ts=merged[-1].ts,
end_ts=max(merged[-1].end_ts, annotation.end_ts),
)
continue
merged.append(
StageWindow(
stage=annotation.stage,
ts=annotation.ts,
end_ts=annotation.end_ts,
)
)
return merged
def resolve_stage_from_windows(
probe_ts: float,
windows: Sequence[StageWindow],
) -> Tuple[Optional[str], Optional[float]]:
nearest_stage: Optional[str] = None
nearest_gap: Optional[float] = None
for window in windows:
if window.ts <= probe_ts <= window.end_ts + 1e-3:
return window.stage, 0.0
gap = min(abs(probe_ts - window.ts), abs(probe_ts - window.end_ts))
if nearest_gap is None or gap < nearest_gap:
nearest_gap = gap
nearest_stage = window.stage
return nearest_stage, nearest_gap
def resolve_kernel_stage(
*,
kernel_ts: float,
external_id: Optional[int],
annotations_by_external_id: Dict[int, StageAnnotation],
gpu_annotations: Sequence[StageWindow],
cpu_annotations: Sequence[StageWindow],
) -> str:
if external_id is not None:
annotation = annotations_by_external_id.get(external_id)
if annotation is not None:
return annotation.stage
probe_ts = kernel_ts + 1e-3
nearest_stage: Optional[str] = None
nearest_gap: Optional[float] = None
for windows in (gpu_annotations, cpu_annotations):
stage, gap = resolve_stage_from_windows(probe_ts, windows)
if gap == 0.0 and stage is not None:
return stage
if stage is not None and (
nearest_gap is None or (gap is not None and gap < nearest_gap)
):
nearest_stage = stage
nearest_gap = gap
if (
nearest_stage is not None
and nearest_gap is not None
and nearest_gap <= 20_000.0
):
return nearest_stage
return "all"
def extract_trace_data(
trace: dict,
) -> Tuple[
List[KernelEvent],
List[CpuOpEvent],
Dict[Tuple[str, str], List[PythonFrame]],
List[LaunchEvent],
Optional[str],
float,
]:
# Build the basic trace views in one pass so later stages can stay simple:
# GPU kernels for ranking, CPU ops for External-id mapping, Python frames for
# source attribution, and CUDA launch calls for correlation-based fallback.
raw_events = extract_trace_events(trace)
correlation_external = build_correlation_external_lookup(raw_events)
(
annotations_by_external_id,
gpu_stage_annotations,
cpu_stage_annotations,
) = build_stage_annotations(raw_events)
chosen_pid = select_heaviest_pid(
raw_events,
is_gpu_kernel_event,
preferred_substrings=("TP00", "TP-0"),
)
kernels: List[KernelEvent] = []
cpu_ops: List[CpuOpEvent] = []
launches: List[LaunchEvent] = []
python_frames: DefaultDict[Tuple[str, str], List[PythonFrame]] = defaultdict(list)
min_ts = None
max_end = None
for event in raw_events:
if event.get("ph") != "X":
continue
pid = str(event.get("pid"))
tid = str(event.get("tid"))
ts = float(event.get("ts", 0.0))
dur = float(event.get("dur", 0.0))
cat = str(event.get("cat", ""))
args = event.get("args") or {}
name = str(event.get("name", ""))
if cat == "python_function":
python_frames[(pid, tid)].append(
PythonFrame(
name=name,
normalized_name=normalize_source_location(name),
pid=pid,
tid=tid,
ts=ts,
dur=dur,
python_id=coerce_optional_int(args.get("Python id")),
parent_id=coerce_optional_int(args.get("Python parent id")),
end_ts=ts + dur,
priority=frame_priority(name),
)
)
correlation = coerce_optional_int(args.get("correlation"))
external_id = coerce_optional_int(args.get("External id"))
if external_id is None and correlation is not None:
external_id = correlation_external.get(correlation)
if cat == "cpu_op" and external_id is not None:
cpu_ops.append(
CpuOpEvent(
name=name,
pid=pid,
tid=tid,
ts=ts,
dur=dur,
external_id=external_id,
)
)
if is_cuda_launch_event(name, cat) and correlation is not None:
launches.append(
LaunchEvent(
name=name,
pid=pid,
tid=tid,
ts=ts,
dur=dur,
correlation=correlation,
)
)
if chosen_pid is None or not is_gpu_kernel_event(event) or pid != chosen_pid:
continue
min_ts = ts if min_ts is None else min(min_ts, ts)
max_end = ts + dur if max_end is None else max(max_end, ts + dur)
kernels.append(
KernelEvent(
name=name,
canonical_name=canonicalize_name(name),
category=classify_kernel(name),
stage=resolve_kernel_stage(
kernel_ts=ts,
external_id=external_id,
annotations_by_external_id=annotations_by_external_id,
gpu_annotations=gpu_stage_annotations,
cpu_annotations=cpu_stage_annotations,
),
pid=pid,
tid=tid,
ts=ts,
dur=dur,
external_id=external_id,
correlation=correlation,
)
)
for frames in python_frames.values():
frames.sort(key=lambda item: (item.ts, item.end_ts))
window_us = 0.0 if min_ts is None or max_end is None else max_end - min_ts
return kernels, cpu_ops, dict(python_frames), launches, chosen_pid, window_us
def build_correlation_external_lookup(raw_events: Sequence[dict]) -> Dict[int, int]:
lookup: Dict[int, int] = {}
for event in raw_events:
args = event.get("args", {}) or {}
correlation = coerce_optional_int(args.get("correlation"))
external_id = coerce_optional_int(args.get("External id"))
if correlation is not None and external_id is not None:
lookup[correlation] = external_id
return lookup
def build_timed_event_index(events: Sequence[object]) -> TimedEventIndex:
ordered = list(events)
ordered.sort(key=lambda item: item.ts)
return TimedEventIndex(
events=ordered,
start_ts=[float(item.ts) for item in ordered],
)
def build_cpu_op_index(cpu_ops: Sequence[CpuOpEvent]) -> Dict[int, TimedEventIndex]:
output: DefaultDict[int, List[CpuOpEvent]] = defaultdict(list)
for cpu_op in cpu_ops:
output[cpu_op.external_id].append(cpu_op)
return {
external_id: build_timed_event_index(items)
for external_id, items in output.items()
}
def match_cpu_op(
kernel: KernelEvent, cpu_ops_by_external_id: Dict[int, TimedEventIndex]
) -> Optional[CpuOpEvent]:
if kernel.external_id is None:
return None
return match_timed_event(
cpu_ops_by_external_id.get(kernel.external_id, []), kernel.ts
)
def build_launch_index(
launch_events: Sequence[LaunchEvent],
) -> Dict[int, TimedEventIndex]:
output: DefaultDict[int, List[LaunchEvent]] = defaultdict(list)
for launch in launch_events:
output[launch.correlation].append(launch)
return {
correlation: build_timed_event_index(items)
for correlation, items in output.items()
}
def match_launch_event(
kernel: KernelEvent, launches_by_correlation: Dict[int, TimedEventIndex]
) -> Optional[LaunchEvent]:
if kernel.correlation is None:
return None
return match_timed_event(
launches_by_correlation.get(kernel.correlation, []), kernel.ts
)
def match_timed_event(index: object, probe_ts: float):
if not index:
return None
if isinstance(index, TimedEventIndex):
events = index.events
if not events:
return None
right = bisect_right(index.start_ts, probe_ts + 1e-3)
candidates: List[object] = []
if right > 0:
candidates.extend(events[max(0, right - 4) : right])
if right < len(events):
candidates.extend(events[right : min(len(events), right + 2)])
if not candidates:
return None
earlier = [item for item in candidates if item.ts <= probe_ts + 1e-3]
if earlier:
return min(earlier, key=lambda item: abs((item.ts + item.dur) - probe_ts))
return min(candidates, key=lambda item: abs(item.ts - probe_ts))
events = list(index)
if not events:
return None
earlier = [item for item in events if item.ts <= probe_ts + 1e-3]
if earlier:
return min(earlier, key=lambda item: abs((item.ts + item.dur) - probe_ts))
return min(events, key=lambda item: abs(item.ts - probe_ts))
def resolve_active_frames_linear(
frames: Sequence[PythonFrame], probe_ts: float
) -> List[PythonFrame]:
active = [item for item in frames if item.ts <= probe_ts <= item.end_ts]
active.sort(key=lambda item: (item.ts, item.end_ts))
return active
def thread_has_crossing_frames(frames: Sequence[PythonFrame]) -> bool:
ordered_frames = sorted(frames, key=lambda item: (item.ts, -item.end_ts))
stack: List[PythonFrame] = []
for frame in ordered_frames:
while stack and stack[-1].end_ts < frame.ts:
stack.pop()
if stack and frame.end_ts > stack[-1].end_ts + 1e-3:
return True
stack.append(frame)
return False
def render_frame_resolution(
active_frames: Sequence[PythonFrame],
) -> Optional[FrameResolution]:
if not active_frames:
return None
chosen_frame = choose_mapping_frame(active_frames)
if chosen_frame is None:
return None
return FrameResolution(
location=chosen_frame.normalized_name,
stack=build_stack_display(active_frames),
)
def resolve_thread_query_times(
frames: Sequence[PythonFrame], query_times: Sequence[float]
) -> Dict[float, Optional[FrameResolution]]:
if not frames or not query_times:
return {}
ordered_frames = sorted(frames, key=lambda item: (item.ts, -item.end_ts))
ordered_queries = sorted(set(float(ts) for ts in query_times))
results: Dict[float, Optional[FrameResolution]] = {}
active_frames: List[PythonFrame] = []
frame_idx = 0
total_frames = len(ordered_frames)
for ts in ordered_queries:
while frame_idx < total_frames and ordered_frames[frame_idx].ts <= ts:
active_frames.append(ordered_frames[frame_idx])
frame_idx += 1
if active_frames:
active_frames = [
frame for frame in active_frames if frame.end_ts >= ts - 1e-3
]
results[ts] = render_frame_resolution(active_frames)
return results
def build_frame_resolution_index(
python_frames: Dict[Tuple[str, str], List[PythonFrame]],
query_times_by_thread: Dict[Tuple[str, str], Sequence[float]],
) -> Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]]:
output: Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]] = {}
for thread_key, query_times in query_times_by_thread.items():
frames = python_frames.get(thread_key, [])
output[thread_key] = resolve_thread_query_times(frames, query_times)
return output
def find_active_python_frames(
cpu_op: CpuOpEvent,
python_frames: Dict[Tuple[str, str], List[PythonFrame]],
) -> List[PythonFrame]:
frames = python_frames.get((cpu_op.pid, cpu_op.tid), [])
if not frames:
return []
probe_ts = cpu_op.ts + min(cpu_op.dur * 0.5, 1.0)
return resolve_active_frames_linear(frames, probe_ts)
def find_active_python_frames_at_ts(
*,
pid: str,
tid: str,
ts: float,
python_frames: Dict[Tuple[str, str], List[PythonFrame]],
) -> List[PythonFrame]:
frames = python_frames.get((pid, tid), [])
if not frames:
return []
return resolve_active_frames_linear(frames, ts)
def render_kernel_site(
active_frames: Sequence[PythonFrame], cpu_op_name: str
) -> Tuple[str, str, str]:
chosen_frame = choose_mapping_frame(active_frames)
if chosen_frame is None:
return "unresolved", "", cpu_op_name
return chosen_frame.normalized_name, build_stack_display(active_frames), cpu_op_name
def resolve_kernel_site_context(
kernel: KernelEvent,
cpu_ops_by_external_id: Dict[int, TimedEventIndex],
python_frames: Dict[Tuple[str, str], List[PythonFrame]],
launches_by_correlation: Dict[int, TimedEventIndex],
frame_resolution_index: Optional[
Dict[Tuple[str, str], Dict[float, Optional[FrameResolution]]]
] = None,
) -> Tuple[str, str, str]:
# Prefer the normal External-id path first. If the kernel dropped that link,
# fall back to the correlated CUDA launch and reuse the Python frames that
# were active when the launch happened.
cpu_op = match_cpu_op(kernel, cpu_ops_by_external_id)
if cpu_op is not None:
probe_ts = cpu_op.ts + min(cpu_op.dur * 0.5, 1.0)
if frame_resolution_index is not None:
resolved = frame_resolution_index.get((cpu_op.pid, cpu_op.tid), {}).get(
probe_ts
)
if resolved is not None:
return resolved.location, resolved.stack, cpu_op.name
active_frames = find_active_python_frames(cpu_op, python_frames)
if active_frames:
return render_kernel_site(active_frames, cpu_op.name)
launch_event = match_launch_event(kernel, launches_by_correlation)
if launch_event is not None:
if frame_resolution_index is not None:
resolved = frame_resolution_index.get(
(launch_event.pid, launch_event.tid), {}
).get(launch_event.ts)
if resolved is not None:
cpu_op_name = cpu_op.name if cpu_op is not None else launch_event.name
return resolved.location, resolved.stack, cpu_op_name
active_frames = find_active_python_frames_at_ts(
pid=launch_event.pid,
tid=launch_event.tid,
ts=launch_event.ts,
python_frames=python_frames,
)
if active_frames:
cpu_op_name = cpu_op.name if cpu_op is not None else launch_event.name
return render_kernel_site(active_frames, cpu_op_name)
return "unresolved", "", launch_event.name
cpu_op_name = cpu_op.name if cpu_op is not None else ""
return "unresolved", "", cpu_op_name
def choose_mapping_frame(active_frames: Sequence[PythonFrame]) -> Optional[PythonFrame]:
if not active_frames:
return None
best = active_frames[0]
best_key = (best.priority, best.ts, -best.dur)
for item in active_frames[1:]:
key = (item.priority, item.ts, -item.dur)
if key > best_key:
best = item
best_key = key
return best
def build_stack_display(active_frames: Sequence[PythonFrame]) -> str:
if not active_frames:
return ""
filtered = [item.normalized_name for item in active_frames if item.priority > 0]
if not filtered:
filtered = [active_frames[-1].normalized_name]
return " -> ".join(filtered[-4:])
def aggregate(events: Iterable[KernelEvent], key_fn) -> Dict[str, Aggregate]:
output: Dict[str, Aggregate] = defaultdict(Aggregate)
for event in events:
key = key_fn(event)
item = output[key]
item.total_us += event.dur
item.count += 1
item.max_us = max(item.max_us, event.dur)
return output
def group_kernels_by_stage(
kernels: Sequence[KernelEvent], default_stage: str
) -> Dict[str, List[KernelEvent]]:
grouped: DefaultDict[str, List[KernelEvent]] = defaultdict(list)
for kernel in kernels:
stage = default_stage if default_stage != "all" else (kernel.stage or "all")
grouped[stage].append(kernel)
return dict(grouped)
def aggregate_kernel_sites(
kernels: Sequence[KernelEvent],
cpu_ops_by_external_id: Dict[int, TimedEventIndex],
python_frames: Dict[Tuple[str, str], List[PythonFrame]],
launches_by_correlation: Optional[Dict[int, TimedEventIndex]] = None,
site_context_cache: Optional[
Dict[Tuple[str, str, float, Optional[int], Optional[int]], Tuple[str, str, str]]
] = None,
) -> Dict[str, Dict[str, MappingSiteAggregate]]:
# Each kernel is mapped independently so the fallback behavior stays easy to
# reason about and easy to regression-test.
output: DefaultDict[str, DefaultDict[str, MappingSiteAggregate]] = defaultdict(
lambda: defaultdict(MappingSiteAggregate)
)
launch_index = launches_by_correlation or {}
query_times_by_thread: DefaultDict[Tuple[str, str], List[float]] = defaultdict(list)
for kernel in kernels:
cpu_op = match_cpu_op(kernel, cpu_ops_by_external_id)
if cpu_op is not None:
query_times_by_thread[(cpu_op.pid, cpu_op.tid)].append(
cpu_op.ts + min(cpu_op.dur * 0.5, 1.0)
)
launch_event = match_launch_event(kernel, launch_index)
if launch_event is not None:
query_times_by_thread[(launch_event.pid, launch_event.tid)].append(
launch_event.ts
)
frame_resolution_index = build_frame_resolution_index(
python_frames, query_times_by_thread
)
resolved_cache = site_context_cache if site_context_cache is not None else {}
for kernel in kernels:
cache_key = (
kernel.pid,
kernel.tid,
kernel.ts,
kernel.external_id,
kernel.correlation,
)
cached = resolved_cache.get(cache_key)
if cached is None:
cached = resolve_kernel_site_context(
kernel,
cpu_ops_by_external_id,
python_frames,
launch_index,
frame_resolution_index=frame_resolution_index,
)
resolved_cache[cache_key] = cached
location, stack, cpu_op_name = cached
item = output[kernel.canonical_name][location]
item.total_us += kernel.dur
item.count += 1
if cpu_op_name:
item.cpu_ops[cpu_op_name] += 1
if stack:
item.stacks[stack] += 1
return {kernel_name: dict(locations) for kernel_name, locations in output.items()}
def merge_site_stats(
destination: DefaultDict[str, DefaultDict[str, MappingSiteAggregate]],
source: Dict[str, Dict[str, MappingSiteAggregate]],
) -> None:
for kernel_name, locations in source.items():
for location, aggregate_item in locations.items():
target = destination[kernel_name][location]
target.total_us += aggregate_item.total_us
target.count += aggregate_item.count
target.cpu_ops.update(aggregate_item.cpu_ops)
target.stacks.update(aggregate_item.stacks)
def build_stage_payload(
site_stats: Dict[str, Dict[str, MappingSiteAggregate]],
kernel_categories: Dict[str, str],
) -> Dict[str, dict]:
kernels_payload: Dict[str, dict] = {}
for kernel_name, locations in sorted(site_stats.items()):
total_us = sum(item.total_us for item in locations.values())
sites = []
for location, aggregate_item in sorted(
locations.items(),
key=lambda pair: pair[1].total_us,
reverse=True,
):
sites.append(
{
"location": location,
"display_location": extract_preferred_stack_location(
aggregate_item.stacks.most_common(1)[0][0]
if aggregate_item.stacks
else None
)
or location,
"launches": aggregate_item.count,
"total_us": round(aggregate_item.total_us, 3),
"share_pct_within_kernel": round(
pct(aggregate_item.total_us, total_us), 3
),
"top_cpu_op": (
aggregate_item.cpu_ops.most_common(1)[0][0]
if aggregate_item.cpu_ops
else None
),
"stack": (
aggregate_item.stacks.most_common(1)[0][0]
if aggregate_item.stacks
else None
),
}
)
sites.sort(
key=lambda site: (
source_location_priority(site_display_location(site)),
float(site.get("total_us", 0.0)),
int(site.get("launches", 0)),
),
reverse=True,
)
kernels_payload[kernel_name] = {
"category": kernel_categories.get(kernel_name, "other"),
"sites": sites,
"best_location": (
site_display_location(sites[0])
if sites
else choose_best_location(locations)
),
}
return {"kernels": kernels_payload}
def load_kernel_map(path: Path) -> dict:
with open(path, "r", encoding="utf-8") as handle:
return json.load(handle)
def relaxed_kernel_entry_lookup(
kernels: Dict[str, dict], kernel_name: str
) -> Optional[dict]:
if kernel_name in kernels:
return kernels[kernel_name]
lowered = kernel_name.lower()
best_key = None
best_score = -1
for candidate_key in kernels:
candidate_lowered = candidate_key.lower()
if candidate_lowered.startswith(lowered) or lowered.startswith(
candidate_lowered
):
score = min(len(candidate_lowered), len(lowered))
elif candidate_lowered in lowered or lowered in candidate_lowered:
score = min(len(candidate_lowered), len(lowered)) // 2
else:
continue
if score > best_score:
best_key = candidate_key
best_score = score
if best_key:
return kernels.get(best_key)
# Long auto-generated kernels such as CUTLASS / FlashAttention templates can
# differ in the middle of the symbol while still sharing the same high-level
# family. Fall back to a conservative common-prefix match so we can still
# recover the higher-level Python callsite from the mapping trace.
lowered_compact = normalize_match_text(kernel_name)
if len(lowered_compact) < 96:
return alias_kernel_entry_lookup(kernels, kernel_name)
def common_prefix_len(left: str, right: str) -> int:
count = 0
for left_ch, right_ch in zip(left, right):
if left_ch != right_ch:
break
count += 1
return count
best_key = None
best_score = -1
for candidate_key in kernels:
candidate_compact = normalize_match_text(candidate_key)
if len(candidate_compact) < 96:
continue
prefix_len = common_prefix_len(lowered_compact, candidate_compact)
shorter_len = min(len(lowered_compact), len(candidate_compact))
if prefix_len < 64 or prefix_len < int(shorter_len * 0.4):
continue
score = prefix_len
if lowered_compact.startswith(
"voidcutlassdevicekernelflash"
) and candidate_compact.startswith("voidcutlassdevicekernelflash"):
score += 32
if score > best_score:
best_key = candidate_key
best_score = score
if best_key:
return kernels.get(best_key)
return alias_kernel_entry_lookup(kernels, kernel_name)
def lookup_kernel_map_entry(
kernel_map: dict, stage: str, kernel_name: str
) -> Optional[dict]:
stage_map = kernel_map.get("stages", {})
for candidate_stage in stage_aliases(stage):
entry = relaxed_kernel_entry_lookup(
stage_map.get(candidate_stage, {}).get("kernels", {}),
kernel_name,
)
if entry:
return entry
return relaxed_kernel_entry_lookup(
kernel_map.get("global", {}).get("kernels", {}), kernel_name
)
def best_site_summary(kernel_entry: Optional[dict]) -> Tuple[str, str]:
if not kernel_entry:
return "unresolved", "-"
sites = kernel_entry.get("sites") or []
if not sites:
return kernel_entry.get("best_location", "unresolved"), "-"
preferred_sites = [
site
for site in sites
if is_preferred_source_location(site_display_location(site))
]
candidate_sites = preferred_sites or sites
rendered_locations = []
rendered_cpu_ops = []
for site in candidate_sites[:2]:
location = site_display_location(site)
share = site.get("share_pct_within_kernel")
if len(candidate_sites) > 1 and share is not None:
rendered_locations.append(f"{location} (site share {share:.0f}%)")
else:
rendered_locations.append(location)
cpu_op = site.get("top_cpu_op")
if cpu_op:
rendered_cpu_ops.append(cpu_op)
return "<br>".join(rendered_locations), (
"<br>".join(rendered_cpu_ops) if rendered_cpu_ops else "-"
)
def resolve_kernel_entry(
stage: str,
kernel_name: str,
local_stage_payload: dict,
external_kernel_map: Optional[dict],
) -> Optional[dict]:
if external_kernel_map:
kernel_entry = lookup_kernel_map_entry(external_kernel_map, stage, kernel_name)
if kernel_entry:
return kernel_entry
return relaxed_kernel_entry_lookup(
local_stage_payload.get("kernels", {}), kernel_name
)
def build_kernel_rows(
stage: str,
kernel_stats: Dict[str, Aggregate],
kernel_categories: Dict[str, str],
local_stage_payload: dict,
external_kernel_map: Optional[dict],
) -> List[KernelRow]:
rows: List[KernelRow] = []
for kernel_name, aggregate_item in sorted(
kernel_stats.items(),
key=lambda pair: pair[1].total_us,
reverse=True,
):
kernel_entry = resolve_kernel_entry(
stage, kernel_name, local_stage_payload, external_kernel_map
)
location, cpu_op = best_site_summary(kernel_entry)
rows.append(
KernelRow(
name=kernel_name,
category=kernel_categories.get(kernel_name, "other"),
aggregate=aggregate_item,
location=location,
cpu_op=cpu_op,
entry=kernel_entry,
)
)
return rows
def limit_kernel_rows(rows: Sequence[KernelRow], table_limit: int) -> List[KernelRow]:
if table_limit <= 0:
return list(rows)
return list(rows[:table_limit])
def entry_sites(kernel_entry: Optional[dict]) -> List[dict]:
if not kernel_entry:
return []
sites = kernel_entry.get("sites") or []
return [site for site in sites if site.get("location")]
def ordered_unique(values: Iterable[str], limit: int = 4) -> List[str]:
output: List[str] = []
seen = set()
for value in values:
item = str(value).strip()
if not item or item in seen:
continue
seen.add(item)
output.append(item)
if len(output) >= limit:
break
return output
def kernel_row_locations(row: KernelRow, limit: int = 4) -> List[str]:
values = [site_display_location(site) for site in entry_sites(row.entry)]
if not values and row.location and row.location != "unresolved":
values = [fragment.strip() for fragment in row.location.split("<br>")]
return ordered_unique(values, limit=limit)
def format_location_for_fusion_display(location: str) -> str:
text = normalize_text(location)
match = re.match(r"(?P<path>.+?):(?P<line>\d+)\s+(?P<func>.+)$", text)
if not match:
return text
return f"{match.group('func')} @ {match.group('path')}:{match.group('line')}"
def normalize_match_text(text: object) -> str:
return re.sub(r"[^0-9A-Za-z]+", "", normalize_text(text)).lower()
def kernel_entry_total_us(entry: Optional[dict]) -> float:
if not entry:
return 0.0
return sum(float(site.get("total_us", 0.0)) for site in entry.get("sites", []))
def kernel_entry_lookup_text(kernel_name: str, entry: Optional[dict]) -> str:
parts = [kernel_name]
if entry:
parts.append(str(entry.get("best_location") or ""))
for site in entry.get("sites", [])[:4]:
parts.append(str(site.get("location") or ""))
parts.append(str(site.get("display_location") or ""))
parts.append(str(site.get("top_cpu_op") or ""))
parts.append(str(site.get("stack") or ""))
return normalize_match_text(" ".join(parts))
def kernel_alias_token_groups(kernel_name: str) -> List[Tuple[str, ...]]:
lowered = normalize_match_text(kernel_name)
groups: List[Tuple[str, ...]] = []
if "flashattnfwdcombine" in lowered:
groups.append(
(
"flashattnfwdsm90",
"flashattnvarlenfunc",
"vllmflashattnflashattninterface",
"vllmfa3cfwd",
)
)
if "kernelmha" in lowered:
groups.append(
(
"maskedmultiheadattentionkernel",
"attentioninplace",
"attentionbackendtrtllm",
)
)
if "applybiasropeupdatekvcachev2" in lowered:
groups.append(
(
"fusedqknormropekernel",
"applyqknormrope",
"modelingqwen3py98applyqknormrope",
)
)
if lowered.startswith("memset"):
groups.append(("memset",))
return groups
def alias_kernel_entry_lookup(
kernels: Dict[str, dict], kernel_name: str
) -> Optional[dict]:
alias_groups = kernel_alias_token_groups(kernel_name)
if not alias_groups:
return None
best_key = None
best_score = -1
for candidate_key, entry in kernels.items():
candidate_text = kernel_entry_lookup_text(candidate_key, entry)
score = 0
for group_index, group in enumerate(alias_groups):
group_score = max(
(len(token) for token in group if token in candidate_text),
default=0,
)
if group_score:
score += 1000 * (group_index + 1) + group_score
if score <= 0:
continue
score += max(
source_location_priority(str(entry.get("best_location") or "")),
source_location_priority(best_site_summary(entry)[0]),
)
score += int(kernel_entry_total_us(entry) // 10)
if score > best_score:
best_key = candidate_key
best_score = score
return kernels.get(best_key) if best_key else None
def row_matches(row: KernelRow, *needles: str) -> bool:
lowered = " ".join([row.name, row.location, row.cpu_op]).lower()
lowered_compact = normalize_match_text(lowered)
for needle in needles:
needle_lowered = needle.lower()
if needle_lowered in lowered:
return True
needle_compact = normalize_match_text(needle)
if needle_compact and needle_compact in lowered_compact:
return True
return False
def summarize_text(values: Iterable[str], limit: int = 4) -> str:
items = ordered_unique(values, limit=limit)
return "<br>".join(items) if items else "-"
def summarize_locations(values: Iterable[str], limit: int = 4) -> str:
items = ordered_unique(
(format_location_for_fusion_display(value) for value in values),
limit=limit,
)
return "<br>".join(items) if items else "-"
def summarize_evidence(
rows: Sequence[KernelRow],
total_us: float,
limit: int = 3,
min_share_pct: float = 1.0,
) -> str:
items = []
for row in rows:
share = pct(row.total_us, total_us)
if share < min_share_pct:
continue
items.append(f"{row.name} ({share:.1f}%)")
if len(items) >= limit:
break
return "<br>".join(items) if items else "-"
def model_path_from_server_args(server_args: Optional[dict]) -> str:
if not isinstance(server_args, dict):
return ""
return str(server_args.get("model_path") or server_args.get("model") or "")
def fusion_framework_hints(spec: FusionPatternSpec) -> set[str]:
text = normalize_text(spec.candidate_path).lower()
hints: set[str] = set()
if "vllm/" in text:
hints.add("vllm")
if any(token in text for token in ("tokenspeed/", "tokenspeed-", "tokenspeed_")):
hints.add("tokenspeed")
if "tensorrt_llm/" in text:
hints.add("trtllm")
if any(token in text for token in ("python/sglang/", "sgl_kernel/")):
hints.add("sglang")
return hints
def pattern_supports_framework(
spec: FusionPatternSpec, framework: Optional[str]
) -> bool:
normalized = normalize_text(framework).lower()
if not normalized or normalized == "auto":
return True
hints = fusion_framework_hints(spec)
if not hints:
return True
return normalized in hints
def matching_rows_for_keywords(
kernel_rows: Sequence[KernelRow],
keywords: Sequence[str],
) -> List[KernelRow]:
if not keywords:
return []
return [row for row in kernel_rows if row_matches(row, *keywords)]
def row_identity(row: KernelRow) -> Tuple[str, str, str]:
return (row.name, row.location, row.cpu_op)
def merge_kernel_rows(*groups: Sequence[KernelRow]) -> List[KernelRow]:
output: List[KernelRow] = []
seen = set()
for group in groups:
for row in group:
row_key = row_identity(row)
if row_key in seen:
continue
seen.add(row_key)
output.append(row)
return output
def pattern_model_matches(spec: FusionPatternSpec, model_path: str) -> bool:
if spec.model_include and not any(
token in model_path for token in spec.model_include
):
return False
if spec.model_exclude and any(token in model_path for token in spec.model_exclude):
return False
return True
def pattern_status(spec: FusionPatternSpec, has_active_match: bool) -> str:
if spec.origin == "mainline":
return "mainline direct" if has_active_match else "mainline split"
if spec.origin == "upstream":
return "upstream direct" if has_active_match else "upstream split"
return "pending direct" if has_active_match else "pending split"
def build_pattern_rationale(
spec: FusionPatternSpec,
has_active_match: bool,
related_us: float,
total_us: float,
) -> str:
share = pct(related_us, total_us)
if spec.origin == "mainline":
if has_active_match:
return (
f"`{spec.pattern}` is present in this trace ({share:.1f}% related GPU time). "
f"{spec.rationale_hint}"
)
return (
f"Split kernels in this family take {share:.1f}% of GPU time. "
f"This tree already has a matching path. {spec.rationale_hint}"
)
if spec.origin == "upstream":
return (
f"Matches an upstream path ({share:.1f}% related GPU time). "
f"{spec.rationale_hint}"
)
return (
f"Matches an open upstream path ({share:.1f}% related GPU time). "
f"{spec.rationale_hint}"
)
def pattern_span(spec: FusionPatternSpec) -> int:
return max(len(spec.split_groups), 1 if spec.active_keywords else 0)
def fusion_priority_key(item: FusionOpportunity) -> Tuple[int, int, int, float]:
return (
item.priority,
item.pattern_span,
len(item.covered_row_keys),
item.related_us,
)
def detect_pattern_match(
spec: FusionPatternSpec,
kernel_rows: Sequence[KernelRow],
total_us: float,
model_path: str,
tp_size: int,
framework: Optional[str],
) -> Optional[FusionOpportunity]:
if total_us <= 0:
return None
if not pattern_supports_framework(spec, framework):
return None
if spec.require_tp and tp_size < spec.min_tp_size:
return None
if not pattern_model_matches(spec, model_path):
return None
active_rows = matching_rows_for_keywords(kernel_rows, spec.active_keywords)
split_groups = [
matching_rows_for_keywords(kernel_rows, keywords)
for keywords in spec.split_groups
]
has_active_match = bool(active_rows)
has_split_match = bool(split_groups) and all(split_groups)
if not has_active_match and not has_split_match:
return None
related_rows = merge_kernel_rows(active_rows, *split_groups)
related_us = sum(row.total_us for row in related_rows)
if related_us <= 0:
return None
if not has_active_match and pct(related_us, total_us) < spec.min_share:
return None
return FusionOpportunity(
pattern=spec.pattern,
status=pattern_status(spec, has_active_match),
confidence=(
"Confirmed"
if has_active_match or pct(related_us, total_us) >= spec.likely_share
else "Candidate"
),
related_us=related_us,
evidence=summarize_evidence(related_rows, total_us),
current_locations=summarize_locations(
location for row in related_rows for location in kernel_row_locations(row)
),
candidate_path=spec.candidate_path,
rationale=build_pattern_rationale(
spec=spec,
has_active_match=has_active_match,
related_us=related_us,
total_us=total_us,
),
covered_row_keys=tuple(row_identity(row) for row in related_rows),
pattern_span=pattern_span(spec),
has_active_match=has_active_match,
priority=spec.priority,
subsumes=spec.subsumes,
)
def detect_fusion_opportunities(
kernel_rows: Sequence[KernelRow],
total_us: float,
server_args: Optional[dict],
framework: Optional[str] = None,
) -> List[FusionOpportunity]:
opportunities: List[FusionOpportunity] = []
if total_us <= 0:
return opportunities
model_path = model_path_from_server_args(server_args).lower()
tp_size = 1
if isinstance(server_args, dict):
tp_size = int(server_args.get("tp_size") or 1)
raw_matches: List[FusionOpportunity] = []
for spec in FUSION_PATTERN_REGISTRY:
opportunity = detect_pattern_match(
spec=spec,
kernel_rows=kernel_rows,
total_us=total_us,
model_path=model_path,
tp_size=tp_size,
framework=framework,
)
if opportunity is not None:
raw_matches.append(opportunity)
raw_matches.sort(key=fusion_priority_key, reverse=True)
consumed_row_keys = set()
blocked_patterns = set()
for opportunity in raw_matches:
if opportunity.pattern in blocked_patterns:
continue
if any(
row_key in consumed_row_keys for row_key in opportunity.covered_row_keys
):
continue
opportunities.append(opportunity)
consumed_row_keys.update(opportunity.covered_row_keys)
blocked_patterns.update(opportunity.subsumes)
return opportunities
Reference in New Issue View Git Blame Copy Permalink