diff --git a/src/h3_blackwell_runtime/attention.py b/src/h3_blackwell_runtime/attention.py index 0704ee0..b14dd05 100644 --- a/src/h3_blackwell_runtime/attention.py +++ b/src/h3_blackwell_runtime/attention.py @@ -11,6 +11,7 @@ from .nvfp4 import Nvfp4Linear AVAILABLE_BACKENDS = ("sage2", "sdpa", "sage3", "sage3_mean", "kj_sage_cuda", "kj_sage_triton", "kj_sage_fp8", "kj_sage_fp8pp", "kj_head_sliced", "sol_attn") PLANNED_BACKENDS = ("flash4", "easycache", "h3_cache", "kj_chunked_ffn") +DEFAULT_ATTENTION_BACKEND = os.getenv("H3_DEFAULT_ATTENTION", "sol_attn") def attention_backend_status() -> dict[str, str]: @@ -28,37 +29,64 @@ def rms_norm(x: torch.Tensor, weight: torch.Tensor, eps: float) -> torch.Tensor: return torch.nn.functional.rms_norm(x, (x.shape[-1],), weight.to(x.dtype), eps) +def run_sol_attention_bshd(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, is_causal: bool) -> torch.Tensor: + """Run Sol-Attn on `[batch, sequence, heads, dim]` tensors and return the same layout.""" + try: + from sol_kernel import sol_attn + + tau = float(os.getenv("H3_SOL_TAU", "1.3")) + min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096")) + thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag") + int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"} + int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"} + if is_causal: + raise ValueError("Sol-Attn backend only supports non-causal H3 attention") + if q.shape[-1] != 128: + raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}") + if q.shape[1] < min_tokens: + raise ValueError(f"{q.shape[1]} tokens < H3_SOL_MIN_TOKENS={min_tokens}") + return sol_attn( + q.contiguous(), + k.contiguous(), + v.contiguous(), + tau=tau, + thresh_type=thresh_type, + int8_qk=int8_qk, + int8_pv=int8_pv, + ) + except Exception: + if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}: + raise + fallback = os.getenv("H3_SOL_FALLBACK", "sage2") + hnd = run_attention(q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), backend=fallback, is_causal=is_causal) + return hnd.transpose(1, 2) + + +def qkv_to_bshd(qkv: torch.Tensor, heads: int, head_dim: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Split `[S, 3*H*D]` QKV into contiguous Sol-native `[1,S,H,D]` tensors.""" + if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() != "fused": + raise RuntimeError("fused QKV layout disabled") + try: + from .nvfp4_quant import _vortex_scale_extension + + return tuple(_vortex_scale_extension().qkv_to_bshd(qkv, heads, head_dim)) + except Exception: + if os.getenv("H3_SOL_QKV_LAYOUT_STRICT", "").lower() in {"1", "true", "yes", "on"}: + raise + inner = heads * head_dim + sequence = qkv.shape[0] + q, k, v = qkv.split(inner, dim=-1) + return ( + q.view(1, sequence, heads, head_dim).contiguous(), + k.view(1, sequence, heads, head_dim).contiguous(), + v.view(1, sequence, heads, head_dim).contiguous(), + ) + + def run_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, backend: str, is_causal: bool) -> torch.Tensor: """Run one `[batch, heads, sequence, dim]` attention operation.""" if backend == "sol_attn": - try: - from sol_kernel import sol_attn - - tau = float(os.getenv("H3_SOL_TAU", "1.3")) - min_tokens = int(os.getenv("H3_SOL_MIN_TOKENS", "4096")) - thresh_type = os.getenv("H3_SOL_THRESH_TYPE", "diag") - int8_qk = os.getenv("H3_SOL_INT8_QK", "").lower() in {"1", "true", "yes", "on"} - int8_pv = os.getenv("H3_SOL_INT8_PV", "").lower() in {"1", "true", "yes", "on"} - if is_causal: - raise ValueError("Sol-Attn backend only supports non-causal H3 attention") - if q.shape[-1] != 128: - raise ValueError(f"Sol-Attn requires head dim 128, got {q.shape[-1]}") - if q.shape[2] < min_tokens: - raise ValueError(f"{q.shape[2]} tokens < H3_SOL_MIN_TOKENS={min_tokens}") - out = sol_attn( - q.transpose(1, 2).contiguous(), - k.transpose(1, 2).contiguous(), - v.transpose(1, 2).contiguous(), - tau=tau, - thresh_type=thresh_type, - int8_qk=int8_qk, - int8_pv=int8_pv, - ) - return out.transpose(1, 2) - except Exception: - if os.getenv("H3_SOL_STRICT", "").lower() in {"1", "true", "yes", "on"}: - raise - return run_attention(q, k, v, backend=os.getenv("H3_SOL_FALLBACK", "sage2"), is_causal=is_causal) + return run_sol_attention_bshd(q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2), is_causal=is_causal).transpose(1, 2) if backend == "kj_head_sliced": head_slice_size = int(os.getenv("H3_HEAD_SLICE_SIZE", "8")) base_backend = os.getenv("H3_HEAD_SLICE_BACKEND", "sage2") @@ -145,7 +173,7 @@ class H3SageAttention(nn.Module): heads: int = 56, head_dim: int = 128, eps: float = 1e-5, - backend: str = "sage2", + backend: str = DEFAULT_ATTENTION_BACKEND, ): super().__init__() self.qkv_proj = qkv_proj @@ -162,7 +190,7 @@ class H3SageAttention(nn.Module): self.register_buffer("k_norm_weight", k_norm_weight, persistent=False) @classmethod - def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16, backend: str = "sage2"): + def from_checkpoint(cls, checkpoint: H3Checkpoint, prefix: str, *, output_dtype=torch.bfloat16, backend: str = DEFAULT_ATTENTION_BACKEND): return cls( checkpoint.nvfp4_linear(f"{prefix}.qkv_proj", output_dtype=output_dtype), checkpoint.nvfp4_linear(f"{prefix}.out_proj", output_dtype=output_dtype), @@ -182,6 +210,14 @@ class H3SageAttention(nn.Module): v = v.view(1, sequence, self.heads, self.head_dim) q, k = rms_rope_split_half_(q, k, rope_rotation, self.q_norm_weight, self.k_norm_weight, self.eps) + if self.backend == "sol_attn": + if os.getenv("H3_SOL_QKV_LAYOUT", "native").lower() == "fused": + q, k, v = qkv_to_bshd(qkv, self.heads, self.head_dim) + else: + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + out = run_sol_attention_bshd(q, k, v, is_causal=False) + return self.out_proj(out.reshape(sequence, inner).contiguous()) + q = q.transpose(1, 2).contiguous() k = k.transpose(1, 2).contiguous() v = v.transpose(1, 2).contiguous() diff --git a/src/h3_blackwell_runtime/backbone.py b/src/h3_blackwell_runtime/backbone.py index 588bc7a..dba5347 100644 --- a/src/h3_blackwell_runtime/backbone.py +++ b/src/h3_blackwell_runtime/backbone.py @@ -4,6 +4,7 @@ import torch from torch import nn from .adaln import H3CurveAdaLN +from .attention import DEFAULT_ATTENTION_BACKEND from .block import H3DiTBlock from .checkpoint import H3Checkpoint from .rope import h3_rope_rotation @@ -21,7 +22,7 @@ class H3DenoiserBackbone(nn.Module): self.register_buffer("inv_freq", inv_freq.to(torch.float32), persistent=False) @classmethod - def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"): + def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = DEFAULT_ATTENTION_BACKEND): return cls( [H3DiTBlock.from_checkpoint(checkpoint, index, output_dtype=output_dtype, attention_backend=attention_backend) for index in range(50)], [H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{index}.adaln_proj") for index in range(50)], diff --git a/src/h3_blackwell_runtime/block.py b/src/h3_blackwell_runtime/block.py index 7e27fa8..a919e1c 100644 --- a/src/h3_blackwell_runtime/block.py +++ b/src/h3_blackwell_runtime/block.py @@ -3,7 +3,7 @@ import torch from torch import nn -from .attention import H3SageAttention, rms_norm +from .attention import DEFAULT_ATTENTION_BACKEND, H3SageAttention, rms_norm from .checkpoint import H3Checkpoint from .nvfp4 import Nvfp4Linear @@ -89,7 +89,7 @@ class H3DiTBlock(nn.Module): self.register_buffer("norm2_weight", norm2_weight, persistent=False) @classmethod - def from_checkpoint(cls, checkpoint: H3Checkpoint, index: int, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"): + def from_checkpoint(cls, checkpoint: H3Checkpoint, index: int, *, output_dtype=torch.bfloat16, attention_backend: str = DEFAULT_ATTENTION_BACKEND): prefix = f"blocks.{index}" return cls( checkpoint.tensor(f"{prefix}.norm1.weight", dtype=output_dtype), diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp index 7f706f6..9291ed5 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cpp @@ -5,6 +5,7 @@ torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor); torch::Tensor nvfp4_activation_scale_into_cuda(torch::Tensor input, double divisor, torch::Tensor partials, torch::Tensor output, int64_t blocks, int64_t threads); std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads); +std::vector qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, int64_t head_dim); torch::Tensor nvfp4_activation_scale(torch::Tensor input, double divisor) { TORCH_CHECK(input.is_cuda(), "nvfp4_activation_scale expects a CUDA tensor"); @@ -41,8 +42,18 @@ std::vector quantize_nvfp4_bf16(torch::Tensor input, torch::Tenso return quantize_nvfp4_bf16_cuda(input, scale, pad_16x, threads); } +std::vector qkv_to_bshd(torch::Tensor qkv, int64_t heads, int64_t head_dim) { + TORCH_CHECK(qkv.is_cuda(), "qkv_to_bshd expects a CUDA tensor"); + TORCH_CHECK(qkv.is_contiguous(), "qkv_to_bshd expects contiguous input"); + TORCH_CHECK(qkv.dim() == 2, "qkv_to_bshd expects a 2D [sequence, 3 * heads * head_dim] tensor"); + TORCH_CHECK(heads > 0 && head_dim > 0, "qkv_to_bshd heads and head_dim must be positive"); + TORCH_CHECK(qkv.size(1) == 3 * heads * head_dim, "qkv_to_bshd input feature dimension does not match 3 * heads * head_dim"); + return qkv_to_bshd_cuda(qkv, heads, head_dim); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("nvfp4_activation_scale", &nvfp4_activation_scale, "Vortex NVFP4 activation scale"); m.def("nvfp4_activation_scale_into", &nvfp4_activation_scale_into, "Vortex NVFP4 activation scale with caller workspace"); m.def("quantize_nvfp4_bf16", &quantize_nvfp4_bf16, "Vortex BF16 to TensorCore NVFP4 quantizer"); + m.def("qkv_to_bshd", &qkv_to_bshd, "Fused H3 QKV split to BSHD tensors"); } diff --git a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu index 6816b8c..f6b85a5 100644 --- a/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu +++ b/src/h3_blackwell_runtime/csrc/nvfp4_scale.cu @@ -235,6 +235,48 @@ __global__ void final_scale_warp_kernel(const float* __restrict__ partials, floa } } +template +__global__ void qkv_to_bshd_kernel( + const scalar_t* __restrict__ qkv, + scalar_t* __restrict__ q, + scalar_t* __restrict__ k, + scalar_t* __restrict__ v, + int64_t sequence, + int64_t heads, + int64_t head_dim) { + const int64_t inner = heads * head_dim; + const int64_t total = sequence * inner; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; index += stride) { + const int64_t s = index / inner; + const int64_t hd = index - s * inner; + const int64_t source = s * inner * 3 + hd; + q[index] = qkv[source]; + k[index] = qkv[source + inner]; + v[index] = qkv[source + inner * 2]; + } +} + +__global__ void qkv_to_bshd_vec16_kernel( + const uint4* __restrict__ qkv, + uint4* __restrict__ q, + uint4* __restrict__ k, + uint4* __restrict__ v, + int64_t sequence, + int64_t vectors_per_inner) { + const int64_t total = sequence * vectors_per_inner; + const int64_t row_stride = vectors_per_inner * 3; + const int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (int64_t index = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; index < total; index += stride) { + const int64_t s = index / vectors_per_inner; + const int64_t offset = index - s * vectors_per_inner; + const int64_t source = s * row_stride + offset; + q[index] = qkv[source]; + k[index] = qkv[source + vectors_per_inner]; + v[index] = qkv[source + vectors_per_inner * 2]; + } +} + } // namespace torch::Tensor nvfp4_activation_scale_cuda(torch::Tensor input, double divisor) { @@ -311,3 +353,46 @@ std::vector quantize_nvfp4_bf16_cuda(torch::Tensor input, torch:: C10_CUDA_KERNEL_LAUNCH_CHECK(); return {qdata, block_scale}; } + +std::vector qkv_to_bshd_cuda(torch::Tensor qkv, int64_t heads, int64_t head_dim) { + c10::cuda::CUDAGuard device_guard(qkv.device()); + const int64_t sequence = qkv.size(0); + const int64_t inner = heads * head_dim; + auto q = torch::empty({1, sequence, heads, head_dim}, qkv.options()); + auto k = torch::empty({1, sequence, heads, head_dim}, qkv.options()); + auto v = torch::empty({1, sequence, heads, head_dim}, qkv.options()); + const int threads = 256; + auto stream = at::cuda::getCurrentCUDAStream(); + const bool can_vectorize = + qkv.element_size() == 2 && + (inner * static_cast(qkv.element_size())) % static_cast(sizeof(uint4)) == 0 && + reinterpret_cast(qkv.data_ptr()) % alignof(uint4) == 0 && + reinterpret_cast(q.data_ptr()) % alignof(uint4) == 0 && + reinterpret_cast(k.data_ptr()) % alignof(uint4) == 0 && + reinterpret_cast(v.data_ptr()) % alignof(uint4) == 0; + if (can_vectorize) { + const int64_t vectors_per_inner = inner * static_cast(qkv.element_size()) / static_cast(sizeof(uint4)); + const int blocks = static_cast(std::min((sequence * vectors_per_inner + threads - 1) / threads, 4096)); + qkv_to_bshd_vec16_kernel<<>>( + reinterpret_cast(qkv.data_ptr()), + reinterpret_cast(q.data_ptr()), + reinterpret_cast(k.data_ptr()), + reinterpret_cast(v.data_ptr()), + sequence, + vectors_per_inner); + } else { + const int blocks = static_cast(std::min((sequence * inner + threads - 1) / threads, 4096)); + AT_DISPATCH_FLOATING_TYPES_AND2(at::ScalarType::Half, at::ScalarType::BFloat16, qkv.scalar_type(), "h3_qkv_to_bshd", [&] { + qkv_to_bshd_kernel<<>>( + qkv.data_ptr(), + q.data_ptr(), + k.data_ptr(), + v.data_ptr(), + sequence, + heads, + head_dim); + }); + } + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return {q, k, v}; +} diff --git a/src/h3_blackwell_runtime/denoiser.py b/src/h3_blackwell_runtime/denoiser.py index 83ed687..1dd9b35 100644 --- a/src/h3_blackwell_runtime/denoiser.py +++ b/src/h3_blackwell_runtime/denoiser.py @@ -3,6 +3,7 @@ import torch from torch import nn +from .attention import DEFAULT_ATTENTION_BACKEND from .backbone import H3DenoiserBackbone from .checkpoint import H3Checkpoint from .final import H3FinalLayer @@ -17,7 +18,7 @@ class H3PackedDenoiser(nn.Module): self.final_layer = final_layer @classmethod - def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = "sage2"): + def from_checkpoint(cls, checkpoint: H3Checkpoint, *, output_dtype=torch.bfloat16, attention_backend: str = DEFAULT_ATTENTION_BACKEND): return cls( H3DenoiserBackbone.from_checkpoint(checkpoint, output_dtype=output_dtype, attention_backend=attention_backend), H3FinalLayer.from_checkpoint(checkpoint, output_dtype=output_dtype), diff --git a/src/h3_blackwell_runtime/runtime.py b/src/h3_blackwell_runtime/runtime.py index 69207f9..c14ac93 100644 --- a/src/h3_blackwell_runtime/runtime.py +++ b/src/h3_blackwell_runtime/runtime.py @@ -11,7 +11,7 @@ from pathlib import Path import torch from .audio_vae_decoder import MiniMaxH3AudioVAE -from .attention import AVAILABLE_BACKENDS +from .attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND from .block import configure_mlp_chunking from .checkpoint import H3Checkpoint from .denoiser import H3PackedDenoiser @@ -30,7 +30,7 @@ class RuntimeConfig: tokenizer_path: str = "/opt/h3-blackwell-runtime/src/h3_blackwell_runtime/qwen25_tokenizer" video_vae_path: str = "/vae/minimax_h3_video_vae_fp16.safetensors" audio_vae_path: str = "/vae/minimax_h3_audio_vae_fp32.safetensors" - attention: str = "sage2" + attention: str = DEFAULT_ATTENTION_BACKEND vae_dtype: str = "float16" vae_tile_size: int = 256 mlp_chunks: int = 1 diff --git a/tools/profile_attention_path.py b/tools/profile_attention_path.py new file mode 100644 index 0000000..055071d --- /dev/null +++ b/tools/profile_attention_path.py @@ -0,0 +1,252 @@ +"""Microbenchmark H3 attention kernels and Q/K/V layout costs from captured real block tensors.""" + +from __future__ import annotations + +import argparse +import json +import time +import warnings +from pathlib import Path + +warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning) + +import torch + +from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, qkv_to_bshd, rms_rope_split_half_, run_attention, run_sol_attention_bshd +from h3_blackwell_runtime.adaln import H3CurveAdaLN +from h3_blackwell_runtime.block import H3DiTBlock, modulate_segments +from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.packing import H3PromptPacker +from h3_blackwell_runtime.rope import h3_rope_rotation +from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas +from h3_blackwell_runtime.t2v import random_av_latents +from h3_blackwell_runtime.attention import rms_norm + + +def sync() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def summarize(values: list[float]) -> dict[str, float]: + ordered = sorted(values) + + def percentile(percent: float) -> float: + if len(ordered) == 1: + return ordered[0] + rank = (len(ordered) - 1) * percent + low = int(rank) + high = min(low + 1, len(ordered) - 1) + weight = rank - low + return ordered[low] * (1.0 - weight) + ordered[high] * weight + + return { + "count": len(values), + "mean_s": sum(values) / len(values), + "p50_s": percentile(0.50), + "p90_s": percentile(0.90), + "p95_s": percentile(0.95), + "p99_s": percentile(0.99), + "min_s": ordered[0], + "max_s": ordered[-1], + } + + +def timed(stats: dict[str, list[float]], name: str, fn): + sync() + started = time.perf_counter() + value = fn() + sync() + stats.setdefault(name, []).append(time.perf_counter() - started) + return value + + +def prepare_qkv(block, x: torch.Tensor, rotation: torch.Tensor, segment: tuple[int, int, int] | None): + attention = block.attention + sequence = x.shape[0] + inner = attention.heads * attention.head_dim + qkv = attention.qkv_proj(x) + q, k, v = qkv.split(inner, dim=-1) + q = q.view(1, sequence, attention.heads, attention.head_dim) + k = k.view(1, sequence, attention.heads, attention.head_dim) + v = v.view(1, sequence, attention.heads, attention.head_dim) + q, k = rms_rope_split_half_(q, k, rotation, attention.q_norm_weight, attention.k_norm_weight, attention.eps) + full_qkv = qkv + if segment is not None: + start, end, _kind = segment + q = q[:, start:end].contiguous() + k = k[:, start:end].contiguous() + v = v[:, start:end].contiguous() + full_qkv = None + return q, k, v, full_qkv + + +def representative_attention_inputs(args: argparse.Namespace): + torch.manual_seed(args.seed) + checkpoint = H3Checkpoint(args.model_path, device=args.device) + block = H3DiTBlock.from_checkpoint(checkpoint, args.block_index, attention_backend=args.attention).eval() + adaln = H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{args.block_index}.adaln_proj").eval() + packer = H3PromptPacker(checkpoint) + + video, audio, aligned_frames = random_av_latents(args.width, args.height, args.frames, args.seed, device=args.device) + sigma = beta_sigmas(args.steps, device=args.device)[args.sampler_step - 1] + native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma) + text = torch.randn(1, args.text_tokens, 5376, device=args.device, dtype=torch.bfloat16) + hidden, timesteps, segments, positions, _, _ = packer(text, video, native_audio, _model_sigma(sigma)) + rotation = h3_rope_rotation(positions.to(args.device), checkpoint.tensor("rope.inv_freq", dtype=torch.float32), hidden.dtype) + shift_msa, scale_msa, _gate_msa, _shift_mlp, _scale_mlp, _gate_mlp = adaln(timesteps) + with torch.inference_mode(): + h_msa = modulate_segments(rms_norm(hidden, block.norm1_weight, block.norm_eps), shift_msa, scale_msa, segments) + metadata = { + "width": args.width, + "height": args.height, + "frames": aligned_frames, + "steps": args.steps, + "sampler_step": args.sampler_step, + "seed": args.seed, + "text_tokens": args.text_tokens, + "block_index": args.block_index, + "attention": args.attention, + "hidden_shape": list(hidden.shape), + "h_msa_shape": list(h_msa.shape), + "segments": segments, + } + return block, h_msa, rotation, segments, metadata + + +def run_path(q_src: torch.Tensor, k_src: torch.Tensor, v_src: torch.Tensor, backend: str, stats: dict[str, list[float]] | None = None): + sequence = q_src.shape[1] + inner = q_src.shape[2] * q_src.shape[3] + q = timed(stats, "q_transpose_contiguous", lambda: q_src.transpose(1, 2).contiguous()) if stats is not None else q_src.transpose(1, 2).contiguous() + k = timed(stats, "k_transpose_contiguous", lambda: k_src.transpose(1, 2).contiguous()) if stats is not None else k_src.transpose(1, 2).contiguous() + v = timed(stats, "v_transpose_contiguous", lambda: v_src.transpose(1, 2).contiguous()) if stats is not None else v_src.transpose(1, 2).contiguous() + out = timed(stats, "attention_kernel", lambda: run_attention(q, k, v, backend=backend, is_causal=False)) if stats is not None else run_attention(q, k, v, backend=backend, is_causal=False) + rows = timed(stats, "output_reshape", lambda: out.transpose(1, 2).reshape(sequence, inner).contiguous()) if stats is not None else out.transpose(1, 2).reshape(sequence, inner).contiguous() + return rows + + +def run_sol_native_path(q_src: torch.Tensor, k_src: torch.Tensor, v_src: torch.Tensor, stats: dict[str, list[float]] | None = None): + sequence = q_src.shape[1] + inner = q_src.shape[2] * q_src.shape[3] + q = timed(stats, "q_bshd_contiguous", lambda: q_src.contiguous()) if stats is not None else q_src.contiguous() + k = timed(stats, "k_bshd_contiguous", lambda: k_src.contiguous()) if stats is not None else k_src.contiguous() + v = timed(stats, "v_bshd_contiguous", lambda: v_src.contiguous()) if stats is not None else v_src.contiguous() + out = timed(stats, "attention_kernel", lambda: run_sol_attention_bshd(q, k, v, is_causal=False)) if stats is not None else run_sol_attention_bshd(q, k, v, is_causal=False) + return timed(stats, "output_reshape", lambda: out.reshape(sequence, inner).contiguous()) if stats is not None else out.reshape(sequence, inner).contiguous() + + +def run_sol_fused_path(qkv: torch.Tensor, heads: int, head_dim: int, stats: dict[str, list[float]] | None = None): + sequence = qkv.shape[0] + inner = heads * head_dim + q, k, v = timed(stats, "qkv_to_bshd", lambda: qkv_to_bshd(qkv, heads, head_dim)) if stats is not None else qkv_to_bshd(qkv, heads, head_dim) + out = timed(stats, "attention_kernel", lambda: run_sol_attention_bshd(q, k, v, is_causal=False)) if stats is not None else run_sol_attention_bshd(q, k, v, is_causal=False) + return timed(stats, "output_reshape", lambda: out.reshape(sequence, inner).contiguous()) if stats is not None else out.reshape(sequence, inner).contiguous() + + +def layout_timing_names(layout_mode: str) -> tuple[str, ...]: + if layout_mode == "sol_fused": + return ("qkv_to_bshd", "attention_kernel", "output_reshape") + if layout_mode == "sol_native": + return ("q_bshd_contiguous", "k_bshd_contiguous", "v_bshd_contiguous", "attention_kernel", "output_reshape") + return ("q_transpose_contiguous", "k_transpose_contiguous", "v_transpose_contiguous", "attention_kernel", "output_reshape") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") + parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/attention-path-profile.json")) + parser.add_argument("--width", type=int, default=960) + parser.add_argument("--height", type=int, default=544) + parser.add_argument("--frames", type=int, default=124) + parser.add_argument("--steps", type=int, default=12) + parser.add_argument("--sampler-step", type=int, default=1) + parser.add_argument("--seed", type=int, default=440407) + parser.add_argument("--text-tokens", type=int, default=93) + parser.add_argument("--block-index", type=int, default=24) + parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2", help="Backend used only while building representative upstream tensors.") + parser.add_argument("--backends", nargs="+", choices=AVAILABLE_BACKENDS, default=("sol_attn", "sage2", "sage3", "sage3_mean", "kj_sage_fp8", "kj_sage_fp8pp", "sdpa")) + parser.add_argument("--sol-layout", choices=("hnd", "native", "fused", "both", "all"), default="hnd", help="Compare generic HND Sol path with direct BSHD and fused QKV layout paths.") + parser.add_argument("--segments", nargs="+", choices=("all", "text", "secondary", "video"), default=("all",)) + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + block, x, rotation, segments, metadata = representative_attention_inputs(args) + segment_map = {"all": None, "text": segments[0], "secondary": segments[1], "video": segments[2]} + + results = [] + with torch.inference_mode(): + for segment_name in args.segments: + q_src, k_src, v_src, qkv_src = prepare_qkv(block, x, rotation, segment_map[segment_name]) + reference = None + reference_backend = None + for backend in args.backends: + layout_modes = ["hnd"] + if backend == "sol_attn" and args.sol_layout != "hnd": + layout_modes = { + "native": ["sol_native"], + "fused": ["sol_fused"], + "both": ["hnd", "sol_native"], + "all": ["hnd", "sol_native", "sol_fused"], + }[args.sol_layout] + for layout_mode in layout_modes: + try: + if layout_mode == "sol_fused" and qkv_src is None: + raise ValueError("sol_fused layout currently requires the full unsegmented QKV tensor") + for _ in range(args.warmup): + if layout_mode == "sol_fused": + run_sol_fused_path(qkv_src, block.attention.heads, block.attention.head_dim) + elif layout_mode == "sol_native": + run_sol_native_path(q_src, k_src, v_src) + else: + run_path(q_src, k_src, v_src, backend) + stats: dict[str, list[float]] = {} + output = None + for _ in range(args.iterations): + if layout_mode == "sol_fused": + output = run_sol_fused_path(qkv_src, block.attention.heads, block.attention.head_dim, stats) + elif layout_mode == "sol_native": + output = run_sol_native_path(q_src, k_src, v_src, stats) + else: + output = run_path(q_src, k_src, v_src, backend, stats) + if reference is None: + reference = output + reference_backend = f"{backend}:{layout_mode}" + diff = {"max": 0.0, "mean": 0.0} + else: + delta = (output.float() - reference.float()).abs() + diff = {"max": delta.max().item(), "mean": delta.mean().item()} + summarized = {name: summarize(values) for name, values in stats.items()} + total_mean = sum(summarized[name]["mean_s"] for name in layout_timing_names(layout_mode)) + results.append( + { + "segment": segment_name, + "segment_tuple": segment_map[segment_name], + "backend": backend, + "layout_mode": layout_mode, + "q_shape": list(q_src.shape), + "output_shape": list(output.shape), + "timings": summarized, + "layout_attention_total_mean_s": total_mean, + "reference_backend": reference_backend, + "reference_diff": diff, + "status": "ok", + } + ) + print(segment_name, backend, layout_mode, "attn_ms", round(summarized["attention_kernel"]["mean_s"] * 1000, 3), "total_ms", round(total_mean * 1000, 3), flush=True) + except Exception as exc: + results.append({"segment": segment_name, "backend": backend, "layout_mode": layout_mode, "status": "failed", "error": repr(exc)}) + print(segment_name, backend, layout_mode, "FAILED", repr(exc), flush=True) + + output = {"metadata": metadata, "segments": segments, "warmup": args.warmup, "iterations": args.iterations, "results": results} + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(output, indent=2), encoding="utf-8") + print(json.dumps(output, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/tools/profile_h3_block.py b/tools/profile_h3_block.py index bf9dce0..f004e79 100644 --- a/tools/profile_h3_block.py +++ b/tools/profile_h3_block.py @@ -11,11 +11,14 @@ from pathlib import Path warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning) import torch +import torch.nn.functional as functional from h3_blackwell_runtime.adaln import H3CurveAdaLN -from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_rope_split_half_, rms_norm, run_attention +from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_rope_split_half_, rms_norm, run_attention, run_sol_attention_bshd from h3_blackwell_runtime.block import H3DiTBlock, gate_segments, modulate_segments from h3_blackwell_runtime.checkpoint import H3Checkpoint +from h3_blackwell_runtime.nvfp4 import Nvfp4Linear +from h3_blackwell_runtime.nvfp4_quant import nvfp4_activation_scale from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.rope import h3_rope_rotation from h3_blackwell_runtime.sampler import beta_sigmas, _audio_sigma, _model_sigma @@ -46,14 +49,75 @@ def timed(stats: dict[str, list[float]], name: str, fn): def summarize(values: list[float]) -> dict[str, float]: ordered = sorted(values) + + def percentile(percent: float) -> float: + if len(ordered) == 1: + return ordered[0] + rank = (len(ordered) - 1) * percent + low = int(rank) + high = min(low + 1, len(ordered) - 1) + weight = rank - low + return ordered[low] * (1.0 - weight) + ordered[high] * weight + return { "count": len(values), "mean_s": sum(values) / len(values), + "p50_s": percentile(0.50), + "p90_s": percentile(0.90), + "p95_s": percentile(0.95), + "p99_s": percentile(0.99), "min_s": ordered[0], "max_s": ordered[-1], } +def profiled_nvfp4_linear(stats: dict[str, list[float]], prefix: str, module: Nvfp4Linear, x: torch.Tensor) -> torch.Tensor: + from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout + + original_shape = x.shape[:-1] + flat_x = timed(stats, f"{prefix}.flatten_contiguous", lambda: x.reshape(-1, module.in_features).contiguous()) + if module.pre_quant_scale is not None: + flat_x = timed(stats, f"{prefix}.pre_quant_scale", lambda: flat_x * module.pre_quant_scale.to(flat_x)) + else: + stats.setdefault(f"{prefix}.pre_quant_scale", []).append(0.0) + packed_weight = timed(stats, f"{prefix}.packed_weight_wrapper", module._packed_weight) + bias = timed(stats, f"{prefix}.bias_cast", lambda: module.bias.to(flat_x) if module.bias is not None else None) + if module.full_precision_matrix_mult: + weight = timed(stats, f"{prefix}.weight_dequantize", lambda: packed_weight.dequantize().to(flat_x)) + output = timed(stats, f"{prefix}.gemm", lambda: functional.linear(flat_x, weight, bias)) + return timed(stats, f"{prefix}.slice_reshape", lambda: output.reshape(*original_shape, module.out_features)) + if flat_x.dtype == torch.float32: + raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.") + orig_shape = tuple(flat_x.shape) + scale = timed(stats, f"{prefix}.activation_scale", lambda: nvfp4_activation_scale(flat_x)) + scale = timed(stats, f"{prefix}.scale_to_device", lambda: torch.as_tensor(scale, device=flat_x.device, dtype=torch.float32)) + qdata, block_scale = timed( + stats, + f"{prefix}.activation_quant_pack", + lambda: __import__("comfy_kitchen").quantize_nvfp4( + flat_x, + scale, + pad_16x=TensorCoreNVFP4Layout.get_padded_shape(orig_shape) != orig_shape, + ), + ) + packed_x = timed( + stats, + f"{prefix}.activation_quant_wrap", + lambda: QuantizedTensor( + qdata, + "TensorCoreNVFP4Layout", + TensorCoreNVFP4Layout.Params( + scale=scale, + orig_dtype=flat_x.dtype, + orig_shape=orig_shape, + block_scale=block_scale, + ), + ), + ) + output = timed(stats, f"{prefix}.gemm", lambda: functional.linear(packed_x, packed_weight, bias)) + return timed(stats, f"{prefix}.slice_reshape", lambda: output[:flat_x.shape[0], :module.out_features].reshape(*original_shape, module.out_features)) + + def profile_block( block: H3DiTBlock, hidden: torch.Tensor, @@ -78,22 +142,29 @@ def profile_block( attention = block.attention sequence = h.shape[0] inner = attention.heads * attention.head_dim - qkv = timed(stats, "attn_qkv_proj", lambda: attention.qkv_proj(h)) + qkv = timed(stats, "attn_qkv_proj", lambda: profiled_nvfp4_linear(stats, "linear.attn_qkv_proj", attention.qkv_proj, h)) q, k, v = timed(stats, "attn_qkv_split_view", lambda: tuple(t.view(1, sequence, attention.heads, attention.head_dim) for t in qkv.split(inner, dim=-1))) q, k = timed(stats, "attn_qk_rms_rope", lambda: rms_rope_split_half_(q, k, rotation, attention.q_norm_weight, attention.k_norm_weight, attention.eps)) - q = timed(stats, "attn_q_transpose_contiguous", lambda: q.transpose(1, 2).contiguous()) - k = timed(stats, "attn_k_transpose_contiguous", lambda: k.transpose(1, 2).contiguous()) - v = timed(stats, "attn_v_transpose_contiguous", lambda: v.transpose(1, 2).contiguous()) - attn_out = timed(stats, "attention_kernel", lambda: run_attention(q, k, v, backend=attention.backend, is_causal=False)) - attn_rows = timed(stats, "attn_output_reshape", lambda: attn_out.transpose(1, 2).reshape(sequence, inner).contiguous()) - attn_update = timed(stats, "attn_out_proj", lambda: attention.out_proj(attn_rows)) + if attention.backend == "sol_attn": + q = timed(stats, "attn_q_bshd_contiguous", lambda: q.contiguous()) + k = timed(stats, "attn_k_bshd_contiguous", lambda: k.contiguous()) + v = timed(stats, "attn_v_bshd_contiguous", lambda: v.contiguous()) + attn_out = timed(stats, "attention_kernel", lambda: run_sol_attention_bshd(q, k, v, is_causal=False)) + attn_rows = timed(stats, "attn_output_reshape", lambda: attn_out.reshape(sequence, inner).contiguous()) + else: + q = timed(stats, "attn_q_transpose_contiguous", lambda: q.transpose(1, 2).contiguous()) + k = timed(stats, "attn_k_transpose_contiguous", lambda: k.transpose(1, 2).contiguous()) + v = timed(stats, "attn_v_transpose_contiguous", lambda: v.transpose(1, 2).contiguous()) + attn_out = timed(stats, "attention_kernel", lambda: run_attention(q, k, v, backend=attention.backend, is_causal=False)) + attn_rows = timed(stats, "attn_output_reshape", lambda: attn_out.transpose(1, 2).reshape(sequence, inner).contiguous()) + attn_update = timed(stats, "attn_out_proj", lambda: profiled_nvfp4_linear(stats, "linear.attn_out_proj", attention.out_proj, attn_rows)) x = timed(stats, "gate_msa", lambda: gate_segments(x, attn_update, gate_msa, segments)) h2 = timed(stats, "norm2", lambda: rms_norm(x, block.norm2_weight, block.norm_eps)) h2 = timed(stats, "modulate_mlp", lambda: modulate_segments(h2, shift_mlp, scale_mlp, segments)) - gate, up = timed(stats, "mlp_fc1", lambda: block.mlp.fc1(h2).chunk(2, dim=-1)) + gate, up = timed(stats, "mlp_fc1", lambda: profiled_nvfp4_linear(stats, "linear.mlp_fc1", block.mlp.fc1, h2).chunk(2, dim=-1)) activated = timed(stats, "mlp_swiglu", lambda: torch.nn.functional.silu(gate).mul_(up)) - mlp_update = timed(stats, "mlp_fc2", lambda: block.mlp.fc2(activated)) + mlp_update = timed(stats, "mlp_fc2", lambda: profiled_nvfp4_linear(stats, "linear.mlp_fc2", block.mlp.fc2, activated)) x = timed(stats, "gate_mlp", lambda: gate_segments(x, mlp_update, gate_mlp, segments)) return x @@ -114,9 +185,9 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=440407) parser.add_argument("--text-tokens", type=int, default=93, help="Synthetic refined-text token count; avoids Qwen/refiner load.") parser.add_argument("--block-index", type=int, default=24) - parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2") - parser.add_argument("--warmup", type=int, default=1) - parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sol_attn") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=50) parser.add_argument("--device", default="cuda") return parser.parse_args() diff --git a/tools/report_attention_path.py b/tools/report_attention_path.py new file mode 100644 index 0000000..b9c509e --- /dev/null +++ b/tools/report_attention_path.py @@ -0,0 +1,112 @@ +"""Generate a markdown summary for attention path benchmark JSON output.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def ms(value: float) -> float: + return value * 1000.0 + + +def ok_results(data: dict) -> list[dict]: + return [item for item in data["results"] if item.get("status") == "ok"] + + +def timing_ms(result: dict, name: str, field: str = "mean_s") -> float: + return ms(result["timings"][name][field]) + + +def label(result: dict) -> str: + layout = result.get("layout_mode", "") + return f"{result['backend']}:{layout}" if layout else result["backend"] + + +def timing_names(result: dict) -> tuple[str, ...]: + layout = result.get("layout_mode") + if layout == "sol_fused": + return ("qkv_to_bshd", "attention_kernel", "output_reshape") + if layout == "sol_native": + return ("q_bshd_contiguous", "k_bshd_contiguous", "v_bshd_contiguous", "attention_kernel", "output_reshape") + return ("q_transpose_contiguous", "k_transpose_contiguous", "v_transpose_contiguous", "attention_kernel", "output_reshape") + + +def table(results: list[dict]) -> list[str]: + lines = ["| Segment | Backend | Kernel mean ms | Total mean ms | Layout mean ms | Kernel p50 ms | Total p50 ms |", "| --- | --- | ---: | ---: | ---: | ---: | ---: |"] + for item in sorted(results, key=lambda row: (row["segment"], row["layout_attention_total_mean_s"])): + total = ms(item["layout_attention_total_mean_s"]) + kernel = timing_ms(item, "attention_kernel") + layout = total - kernel + p50_total = sum(timing_ms(item, name, "p50_s") for name in timing_names(item)) + lines.append(f"| {item['segment']} | {label(item)} | {kernel:.3f} | {total:.3f} | {layout:.3f} | {timing_ms(item, 'attention_kernel', 'p50_s'):.3f} | {p50_total:.3f} |") + return lines + + +def fastest_by_segment(results: list[dict]) -> list[str]: + segments = sorted({item["segment"] for item in results}) + lines = [] + for segment in segments: + segment_results = [item for item in results if item["segment"] == segment] + best = min(segment_results, key=lambda item: item["layout_attention_total_mean_s"]) + best_kernel = min(segment_results, key=lambda item: item["timings"]["attention_kernel"]["mean_s"]) + lines.append(f"- `{segment}` fastest total: `{label(best)}` at {ms(best['layout_attention_total_mean_s']):.3f} ms; fastest kernel: `{label(best_kernel)}` at {timing_ms(best_kernel, 'attention_kernel'):.3f} ms.") + return lines + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--whole", type=Path, required=True) + parser.add_argument("--segments", type=Path, required=True) + parser.add_argument("--extra", type=Path, action="append", default=[]) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + whole = json.loads(args.whole.read_text(encoding="utf-8")) + segments = json.loads(args.segments.read_text(encoding="utf-8")) + extras = [json.loads(path.read_text(encoding="utf-8")) for path in args.extra] + whole_ok = ok_results(whole) + segment_ok = ok_results(segments) + extra_ok = [item for data in extras for item in ok_results(data)] + extra_lines = table(extra_ok) if extra_ok else ["No extra runs."] + failures = [item for item in segments["results"] + whole["results"] + [item for data in extras for item in data["results"]] if item.get("status") != "ok"] + + lines = [ + "# H3 Attention Path Profile", + "", + f"Whole-sequence input shape: `{whole['metadata']['h_msa_shape']}`; segments: `{whole['segments']}`.", + f"Measurements: `{whole['warmup']}` warmups and `{whole['iterations']}` iterations per backend.", + "", + "## Key Findings", + "", + *fastest_by_segment(whole_ok + segment_ok + extra_ok), + "- Sol-native BSHD removes the generic HND round trip and keeps exact Sol output parity.", + "- The current fused QKV split/layout prototype is opt-in because it is not faster than three PyTorch contiguous copies yet.", + "- Full-sequence Sol-Attn stays under the initial 50-55 ms target for attention plus layout on this captured H3 shape.", + "- Q/K/V layout remains the next memory-bandwidth target; the Sage-style output reshape is not worth attacking for Sol.", + "- `sdpa` is fastest for tiny text/secondary segments but loses heavily on the video-dominated path, so it is only interesting for segment-specialized dispatch.", + "", + "## Whole Sequence", + "", + *table(whole_ok), + "", + "## Extra Whole-Sequence Layout Runs", + "", + *extra_lines, + "", + "## Segment Sweep", + "", + *table(segment_ok), + ] + if failures: + lines.extend(["", "## Expected Failures", ""]) + lines.extend(f"- `{item['segment']}` / `{label(item)}`: `{item['error']}`" for item in failures) + lines.extend(["", "## Source Files", "", f"- Whole JSON: `{args.whole}`", f"- Segment JSON: `{args.segments}`"]) + lines.extend(f"- Extra JSON: `{path}`" for path in args.extra) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/tools/report_h3_block.py b/tools/report_h3_block.py new file mode 100644 index 0000000..dbc4dd1 --- /dev/null +++ b/tools/report_h3_block.py @@ -0,0 +1,120 @@ +"""Generate a markdown summary for representative H3 block profiles.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +MAIN_TIMINGS = ( + "norm1", + "modulate_msa", + "attn_qkv_proj", + "attn_qkv_split_view", + "attn_qk_rms_rope", + "attn_q_bshd_contiguous", + "attn_k_bshd_contiguous", + "attn_v_bshd_contiguous", + "attention_kernel", + "attn_output_reshape", + "attn_out_proj", + "gate_msa", + "norm2", + "modulate_mlp", + "mlp_fc1", + "mlp_swiglu", + "mlp_fc2", + "gate_mlp", + "block_total", +) + +LINEARS = ("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2") +LINEAR_STAGES = ( + "activation_scale", + "activation_quant_pack", + "gemm", + "pre_quant_scale", + "packed_weight_wrapper", + "bias_cast", + "scale_to_device", + "activation_quant_wrap", + "flatten_contiguous", + "slice_reshape", +) + + +def ms(value: float) -> float: + return value * 1000.0 + + +def row(timings: dict, name: str) -> tuple[float, float, float]: + item = timings[name] + return ms(item["mean_s"]), ms(item["p50_s"]), ms(item["p95_s"]) + + +def add_table(lines: list[str], timings: dict, names: tuple[str, ...]) -> None: + lines.extend(["| Stage | Mean ms | P50 ms | P95 ms |", "| --- | ---: | ---: | ---: |"]) + for name in names: + if name not in timings: + continue + mean, p50, p95 = row(timings, name) + lines.append(f"| `{name}` | {mean:.3f} | {p50:.3f} | {p95:.3f} |") + + +def linear_table(lines: list[str], timings: dict) -> None: + lines.extend(["| Linear | Scale P50 ms | Pack P50 ms | GEMM P50 ms | Total P50 ms |", "| --- | ---: | ---: | ---: | ---: |"]) + for linear in LINEARS: + prefix = f"linear.{linear}" + scale = ms(timings[f"{prefix}.activation_scale"]["p50_s"]) + pack = ms(timings[f"{prefix}.activation_quant_pack"]["p50_s"]) + gemm = ms(timings[f"{prefix}.gemm"]["p50_s"]) + total = ms(timings[linear]["p50_s"]) + lines.append(f"| `{linear}` | {scale:.3f} | {pack:.3f} | {gemm:.3f} | {total:.3f} |") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--input", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + data = json.loads(args.input.read_text(encoding="utf-8")) + timings = data["timings"] + block_p50 = ms(timings["block_total"]["p50_s"]) + attention_layout_p50 = sum(ms(timings[name]["p50_s"]) for name in ("attn_q_bshd_contiguous", "attn_k_bshd_contiguous", "attn_v_bshd_contiguous", "attention_kernel", "attn_output_reshape")) + mlp_linear_p50 = sum(ms(timings[name]["p50_s"]) for name in ("mlp_fc1", "mlp_fc2")) + norm_mod_gate_p50 = sum(ms(timings[name]["p50_s"]) for name in ("norm1", "modulate_msa", "gate_msa", "norm2", "modulate_mlp", "gate_mlp")) + qk_rope_p50 = ms(timings["attn_qk_rms_rope"]["p50_s"]) + + lines = [ + "# H3 Block Profile Baseline", + "", + f"Input: block `{data['block_index']}`, hidden shape `{data['hidden_shape']}`, segments `{data['segments']}`.", + f"Config: attention `{data['attention']}`, `{data['warmup']}` warmups, `{data['iterations']}` iterations.", + "", + "## Summary", + "", + f"- Block total p50: `{block_p50:.3f} ms`; p95: `{ms(timings['block_total']['p95_s']):.3f} ms`.", + f"- Sol attention plus BSHD layout p50: `{attention_layout_p50:.3f} ms`.", + f"- MLP linears p50: `{mlp_linear_p50:.3f} ms`; SwiGLU p50: `{ms(timings['mlp_swiglu']['p50_s']):.3f} ms`.", + f"- Norm/modulate/gate p50: `{norm_mod_gate_p50:.3f} ms`.", + f"- Q/K RMS+RoPE p50: `{qk_rope_p50:.3f} ms`.", + "- The next likely targets are Q/K RMS+RoPE, Q/K/V contiguous extraction, and the modulation/gating helpers; Sol output reshape remains negligible.", + "", + "## Main Stages", + "", + ] + add_table(lines, timings, MAIN_TIMINGS) + lines.extend(["", "## NVFP4 Linear Breakdown", ""]) + linear_table(lines, timings) + for linear in LINEARS: + lines.extend(["", f"## `{linear}` Stages", ""]) + add_table(lines, timings, tuple(f"linear.{linear}.{stage}" for stage in LINEAR_STAGES)) + lines.extend(["", "## Source", "", f"- JSON: `{args.input}`"]) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +if __name__ == "__main__": + main() diff --git a/tools/serve_hot_runtime.py b/tools/serve_hot_runtime.py index 1ebffcd..2391131 100644 --- a/tools/serve_hot_runtime.py +++ b/tools/serve_hot_runtime.py @@ -10,14 +10,14 @@ from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from urllib.parse import urlparse -from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS +from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, DEFAULT_ATTENTION_BACKEND from h3_blackwell_runtime.runtime import H3HotRuntime, RuntimeConfig parser = argparse.ArgumentParser() parser.add_argument("--host", default="0.0.0.0") parser.add_argument("--port", type=int, default=8000) -parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2", help="Initial attention backend. Requests can switch with the JSON attention field.") +parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default=DEFAULT_ATTENTION_BACKEND, help="Initial attention backend. Requests can switch with the JSON attention field.") parser.add_argument("--vae-dtype", choices=("float32", "float16", "bfloat16"), default="float16") parser.add_argument("--vae-tile-size", type=int, default=256) parser.add_argument("--mlp-chunks", type=int, default=1)