Add Sol native attention baseline
This commit is contained in:
Daniel Maddern
parent
c2fb5c1e3f
commit
731a2813fd
12 changed files with 740 additions and 51 deletions
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)],
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::Tensor scale, bool pad_16x, int64_t threads);
|
||||
std::vector<torch::Tensor> 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<torch::Tensor> quantize_nvfp4_bf16(torch::Tensor input, torch::Tenso
|
|||
return quantize_nvfp4_bf16_cuda(input, scale, pad_16x, threads);
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> 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");
|
||||
}
|
||||
|
|
|
|||
|
|
@ -235,6 +235,48 @@ __global__ void final_scale_warp_kernel(const float* __restrict__ partials, floa
|
|||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__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<int64_t>(blockDim.x) * gridDim.x;
|
||||
for (int64_t index = static_cast<int64_t>(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<int64_t>(blockDim.x) * gridDim.x;
|
||||
for (int64_t index = static_cast<int64_t>(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<torch::Tensor> quantize_nvfp4_bf16_cuda(torch::Tensor input, torch::
|
|||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return {qdata, block_scale};
|
||||
}
|
||||
|
||||
std::vector<torch::Tensor> 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<int64_t>(qkv.element_size())) % static_cast<int64_t>(sizeof(uint4)) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(qkv.data_ptr()) % alignof(uint4) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(q.data_ptr()) % alignof(uint4) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(k.data_ptr()) % alignof(uint4) == 0 &&
|
||||
reinterpret_cast<uintptr_t>(v.data_ptr()) % alignof(uint4) == 0;
|
||||
if (can_vectorize) {
|
||||
const int64_t vectors_per_inner = inner * static_cast<int64_t>(qkv.element_size()) / static_cast<int64_t>(sizeof(uint4));
|
||||
const int blocks = static_cast<int>(std::min<int64_t>((sequence * vectors_per_inner + threads - 1) / threads, 4096));
|
||||
qkv_to_bshd_vec16_kernel<<<blocks, threads, 0, stream>>>(
|
||||
reinterpret_cast<const uint4*>(qkv.data_ptr()),
|
||||
reinterpret_cast<uint4*>(q.data_ptr()),
|
||||
reinterpret_cast<uint4*>(k.data_ptr()),
|
||||
reinterpret_cast<uint4*>(v.data_ptr()),
|
||||
sequence,
|
||||
vectors_per_inner);
|
||||
} else {
|
||||
const int blocks = static_cast<int>(std::min<int64_t>((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<scalar_t><<<blocks, threads, 0, stream>>>(
|
||||
qkv.data_ptr<scalar_t>(),
|
||||
q.data_ptr<scalar_t>(),
|
||||
k.data_ptr<scalar_t>(),
|
||||
v.data_ptr<scalar_t>(),
|
||||
sequence,
|
||||
heads,
|
||||
head_dim);
|
||||
});
|
||||
}
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();
|
||||
return {q, k, v};
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
252
tools/profile_attention_path.py
Normal file
252
tools/profile_attention_path.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
112
tools/report_attention_path.py
Normal file
112
tools/report_attention_path.py
Normal file
|
|
@ -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()
|
||||
120
tools/report_h3_block.py
Normal file
120
tools/report_h3_block.py
Normal file
|
|
@ -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()
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue