[CI][RFC] Replace black-jupyter with ruff-format (#37210)

Co-authored-by: Alison Shao <a.shao@wustl.edu>
This commit is contained in:
Alex Nails
2026-09-02 19:46:08 -07:00
committed by GitHub
co-authored by Alison Shao
parent 2641e427be
commit 28262c20df
1411 changed files with 7766 additions and 8176 deletions
@@ -687,8 +687,7 @@ def run_one_round(
rank_rows = fetch_rank_rows(base_url=context.base_url)
if len(rank_rows) != len(watermarks):
raise RuntimeError(
f"DP rank count changed mid-profile: {len(watermarks)} -> "
f"{len(rank_rows)}."
f"DP rank count changed mid-profile: {len(watermarks)} -> {len(rank_rows)}."
)
new_rank_rows = [
[row for row in rows if row.forward_ct > watermark]
@@ -136,13 +136,13 @@ class BenchArgs:
"--gsp-system-prompt-len",
type=int,
default=BenchArgs.gsp_system_prompt_len,
help="System prompt length, used" "only for generate-shared-prefix",
help="System prompt length, usedonly for generate-shared-prefix",
)
parser.add_argument(
"--gsp-question-len",
type=int,
default=BenchArgs.gsp_question_len,
help="Question length, used" "only for generate-shared-prefix",
help="Question length, usedonly for generate-shared-prefix",
)
parser.add_argument(
"--gsp-output-len",
@@ -259,9 +259,9 @@ def throughput_test_once(
]
if profile:
assert (
"SGLANG_TORCH_PROFILER_DIR" in os.environ
), "Please set SGLANG_TORCH_PROFILER_DIR."
assert "SGLANG_TORCH_PROFILER_DIR" in os.environ, (
"Please set SGLANG_TORCH_PROFILER_DIR."
)
os.makedirs(os.environ["SGLANG_TORCH_PROFILER_DIR"], exist_ok=True)
known_files = None
backend.start_profile(
+8 -8
View File
@@ -486,7 +486,7 @@ def _warmup_cache(
return
print(
f"Warming up cache with {cache_hit_rate*100:.1f}% hit rate "
f"Warming up cache with {cache_hit_rate * 100:.1f}% hit rate "
f"({cached_token_len} tokens per request)"
)
# Create prefix input_ids for cache warming
@@ -1024,7 +1024,7 @@ def get_report_summary(
f"\nInput lens: {bench_args.input_len}. Output lens: {bench_args.output_len}."
)
if bench_args.cache_hit_rate > 0.0:
summary += f" Cache hit rate: {bench_args.cache_hit_rate*100:.1f}%."
summary += f" Cache hit rate: {bench_args.cache_hit_rate * 100:.1f}%."
summary += "\n"
if is_blackwell():
@@ -1241,9 +1241,9 @@ def run_benchmark_internal(
skip_max_running_requests_threshold = float("inf")
skip_token_capacity_threshold = float("inf")
else:
assert (
max_running_requests_per_dp > 0
), f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}"
assert max_running_requests_per_dp > 0, (
f"effective_max_running_requests_per_dp is not set, {max_running_requests_per_dp=}"
)
skip_max_running_requests_threshold = max_running_requests_per_dp * dp_size
print(f"{max_running_requests_per_dp=}")
@@ -1288,9 +1288,9 @@ def run_benchmark_internal(
"--lora-request-distribution=distinct/skewed requires more than "
"one adapter via --lora-name."
)
assert (
bench_args.lora_zipf_alpha > 1
), f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}"
assert bench_args.lora_zipf_alpha > 1, (
f"--lora-zipf-alpha must be > 1, got {bench_args.lora_zipf_alpha}"
)
if bench_args.apply_chat_template and not (
bench_args.fixed_prompt_file or bench_args.dataset_name in REPLAY_TEXT_DATASETS
+26 -26
View File
@@ -261,9 +261,9 @@ async def async_request_openai_completions(
pbar: Optional[tqdm] = None,
) -> RequestFuncOutput:
api_url = request_func_input.api_url
assert api_url.endswith(
"completions"
), "OpenAI Completions API URL must end with 'completions'."
assert api_url.endswith("completions"), (
"OpenAI Completions API URL must end with 'completions'."
)
prompt = request_func_input.prompt
@@ -392,9 +392,9 @@ async def async_request_openai_chat_completions(
latency, TTFT, ITL, and success status.
"""
api_url = request_func_input.api_url
assert api_url.endswith(
"chat/completions"
), "OpenAI Chat Completions API URL must end with 'chat/completions'."
assert api_url.endswith("chat/completions"), (
"OpenAI Chat Completions API URL must end with 'chat/completions'."
)
# TODO put it to other functions when `pbar` logic is refactored
if getattr(args, "print_requests", False):
@@ -1296,9 +1296,9 @@ def _normalize_round_messages(turn: Any) -> Optional[List[Dict[str, str]]]:
def wrap_multi_turn_request_func(request_func: Callable, backend: str) -> Callable:
assert (
backend in MULTI_TURN_BACKENDS
), f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}"
assert backend in MULTI_TURN_BACKENDS, (
f"Multi-turn only supports chat backends: {MULTI_TURN_BACKENDS}, got {backend}"
)
async def f(
request_func_input: RequestFuncInput,
@@ -1534,9 +1534,9 @@ async def benchmark(
lora_name = lora_names[lora_idx]
lora_idx = (lora_idx + 1) % len(lora_names)
else:
assert (
lora_request_distribution == "skewed"
), f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'."
assert lora_request_distribution == "skewed", (
f"Unexpected lora_request_distribution: {lora_request_distribution}. Expected 'skewed'."
)
lora_name = np.random.choice(lora_names, p=lora_probs)
else:
@@ -2000,9 +2000,9 @@ def run_benchmark(args_: argparse.Namespace):
extra_request_body["bootstrap_room"] = 0
if args.tokenize_prompt:
assert (
args.backend == "sglang"
), "`--tokenize-prompt` only compatible with `--backend sglang` currently"
assert args.backend == "sglang", (
"`--tokenize-prompt` only compatible with `--backend sglang` currently"
)
# Set url
if args.port is None:
@@ -2079,18 +2079,18 @@ def run_benchmark(args_: argparse.Namespace):
if args.dataset_name in ["image", "mmmu"]:
args.apply_chat_template = True
assert (
not args.tokenize_prompt
), "`--tokenize-prompt` not compatible with image dataset"
assert not args.tokenize_prompt, (
"`--tokenize-prompt` not compatible with image dataset"
)
if args.lora_request_distribution in ["distinct", "skewed"]:
assert (
args.lora_name is not None and len(args.lora_name) > 1
), "More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution."
assert args.lora_name is not None and len(args.lora_name) > 1, (
"More than 1 LoRA adapter must be specified via --lora-name to use 'distinct' or 'skewed' request distribution."
)
assert (
args.lora_zipf_alpha > 1
), f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1."
assert args.lora_zipf_alpha > 1, (
f"Got invalid value for --lora-zipf-alpha of {args.lora_zipf_alpha}. It must be greater than 1."
)
print(f"{args}\n")
@@ -2364,13 +2364,13 @@ def cli_main():
"--image-format",
type=str,
default="jpeg",
help=("Format of images for image dataset. " "Supports jpeg and png."),
help=("Format of images for image dataset. Supports jpeg and png."),
)
parser.add_argument(
"--image-content",
type=str,
default="random",
help=("Content for images for image dataset. " "Supports random and blank."),
help=("Content for images for image dataset. Supports random and blank."),
)
parser.add_argument(
"--request-rate",
+1 -2
View File
@@ -315,8 +315,7 @@ def _print_diagnostics(unkillable_pids):
print(f" {line}")
else:
print(
"\n[killall] Diagnostic — no sglang/python/gpu processes "
"in this container"
"\n[killall] Diagnostic — no sglang/python/gpu processes in this container"
)
+1 -1
View File
@@ -195,7 +195,7 @@ def serve(args, extra_argv):
else:
registered = registry.get(backend_name)
logger.info(
"Dispatch override enabled: --model-type=%s " "(skip auto detection)",
"Dispatch override enabled: --model-type=%s (skip auto detection)",
backend_name,
)
+1 -2
View File
@@ -126,8 +126,7 @@ def get_model_path(extra_argv):
)
else:
raise Exception(
"Error: --model-path is required. "
"Please provide the path to the model."
"Error: --model-path is required. Please provide the path to the model."
)
return model_path
@@ -245,7 +245,7 @@ def worker(world_size, rank, port, results_queue):
if input_size_bytes > custom_ar.max_size:
if rank == 0:
print(
f" Deterministic kernel skipped: input size ({input_size_bytes/(1024*1024):.1f} MB) > buffer size ({custom_ar.max_size/(1024*1024):.1f} MB)"
f" Deterministic kernel skipped: input size ({input_size_bytes / (1024 * 1024):.1f} MB) > buffer size ({custom_ar.max_size / (1024 * 1024):.1f} MB)"
)
deterministic_kernel_available = False
else:
@@ -412,17 +412,17 @@ def worker(world_size, rank, port, results_queue):
}
print(
f" All-Reduce: {lat_ar_median*1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}"
f" All-Reduce: {lat_ar_median * 1000:.3f}ms, Deterministic: {ar_deterministic}, Max variance: {ar_max_variance:.6f}"
)
print(
f" RS+All-Gather: {lat_rs_ag_median*1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}"
f" RS+All-Gather: {lat_rs_ag_median * 1000:.3f}ms, Deterministic: {rs_ag_deterministic}, Max variance: {rs_ag_max_variance:.6f}"
)
if custom_ar is not None and lat_custom_ar_median is not None:
overhead_custom = (
(lat_custom_ar_median - lat_ar_median) / lat_ar_median
) * 100
print(
f" Custom AR: {lat_custom_ar_median*1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%"
f" Custom AR: {lat_custom_ar_median * 1000:.3f}ms, Deterministic: {custom_ar_deterministic}, Max variance: {custom_ar_max_variance:.6f}, Overhead: {overhead_custom:+.1f}%"
)
if lat_deterministic_kernel_median is not None:
overhead_kernel = (
@@ -433,7 +433,7 @@ def worker(world_size, rank, port, results_queue):
/ lat_rs_ag_median
) * 100
print(
f" Deterministic Kernel: {lat_deterministic_kernel_median*1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%"
f" Deterministic Kernel: {lat_deterministic_kernel_median * 1000:.3f}ms, Deterministic: {deterministic_kernel_deterministic}, Max variance: {deterministic_kernel_max_variance:.6f}, Overhead: {overhead_kernel:+.1f}%, Speedup vs RS+AG: {speedup_kernel_vs_rs_ag:+.1f}%"
)
if lat_optimized_rs_ag_median is not None:
overhead_opt = (
@@ -443,7 +443,7 @@ def worker(world_size, rank, port, results_queue):
(lat_rs_ag_median - lat_optimized_rs_ag_median) / lat_rs_ag_median
) * 100
print(
f" Optimized RS+AG: {lat_optimized_rs_ag_median*1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%"
f" Optimized RS+AG: {lat_optimized_rs_ag_median * 1000:.3f}ms, Deterministic: {optimized_rs_ag_deterministic}, Max variance: {optimized_rs_ag_max_variance:.6f}, Overhead: {overhead_opt:+.1f}%, Speedup vs RS+AG: {speedup_vs_rs_ag:+.1f}%"
)
print(f" RS+AG Overhead: {overhead_rs_ag:+.1f}%")
@@ -515,8 +515,8 @@ def main():
ar_det_str = "" if r["all_reduce"]["deterministic"] else ""
rs_ag_det_str = "" if r["rs_ag"]["deterministic"] else ""
line = (
f"{bs:<8} {r['all_reduce']['latency_median']*1000:<12.3f} {ar_det_str:<8} "
f"{r['rs_ag']['latency_median']*1000:<15.3f} {rs_ag_det_str:<10} "
f"{bs:<8} {r['all_reduce']['latency_median'] * 1000:<12.3f} {ar_det_str:<8} "
f"{r['rs_ag']['latency_median'] * 1000:<15.3f} {rs_ag_det_str:<10} "
f"{r['overhead_rs_ag_pct']:<12.1f}"
)
if r.get("custom_ar") is not None:
@@ -526,7 +526,7 @@ def main():
(custom_ar["latency_median"] - r["all_reduce"]["latency_median"])
/ r["all_reduce"]["latency_median"]
) * 100
line += f" {custom_ar['latency_median']*1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}"
line += f" {custom_ar['latency_median'] * 1000:<18.3f} {custom_ar_det_str:<15} {custom_ar_overhead:<15.1f}"
if r.get("deterministic_kernel") is not None:
det_kernel = r["deterministic_kernel"]
det_kernel_det_str = "" if det_kernel["deterministic"] else ""
@@ -538,7 +538,7 @@ def main():
(r["rs_ag"]["latency_median"] - det_kernel["latency_median"])
/ r["rs_ag"]["latency_median"]
) * 100
line += f" {det_kernel['latency_median']*1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}"
line += f" {det_kernel['latency_median'] * 1000:<18.3f} {det_kernel_det_str:<15} {det_kernel_overhead:<15.1f} {speedup_kernel:<10.1f}"
if r.get("optimized_rs_ag") is not None:
opt_rs_ag = r["optimized_rs_ag"]
opt_rs_ag_det_str = "" if opt_rs_ag["deterministic"] else ""
@@ -550,7 +550,7 @@ def main():
(r["rs_ag"]["latency_median"] - opt_rs_ag["latency_median"])
/ r["rs_ag"]["latency_median"]
) * 100
line += f" {opt_rs_ag['latency_median']*1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}"
line += f" {opt_rs_ag['latency_median'] * 1000:<18.3f} {opt_rs_ag_det_str:<15} {opt_rs_ag_overhead:<15.1f} {speedup:<10.1f}"
print(line)
print("=" * 80)
@@ -114,10 +114,8 @@ def benchmark(batch_size, seq_len, provider, block_size, num_kv_splits):
q_size = qn.numel() * qn.element_size() + qr.numel() * qr.element_size()
gbps = (
lambda ms: (
q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size()
)
gbps = lambda ms: (
(q_size + q_size * dv / d + kv_cache.numel() * kv_cache.element_size())
* 1e-9
/ (ms * 1e-3)
)
@@ -368,9 +368,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
res_fi,
backend="cudnn",
)
assert torch.allclose(
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
), "cudnn fp4 doesn't match cutlass fp4"
assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), (
"cudnn fp4 doesn't match cutlass fp4"
)
mm_fp4(
a_fp4,
b_fp4_T,
@@ -381,9 +381,9 @@ def benchmark(batch_size, provider, N, K, dtype, correctness, csv_file):
res_fi,
backend="trtllm",
)
assert torch.allclose(
res_fi, res_cutlass, atol=1e-3, rtol=1e-3
), "trtllm fp4 doesn't match cutlass fp4"
assert torch.allclose(res_fi, res_cutlass, atol=1e-3, rtol=1e-3), (
"trtllm fp4 doesn't match cutlass fp4"
)
if csv_file:
with open(csv_file, "a", newline="") as f:
@@ -127,8 +127,8 @@ def benchmark(batch_size, provider, N, K):
lambda: vllm_scaled_mm(a, b, scale_a, scale_b, torch.float16, bias),
quantiles=quantiles,
)
gbps = (
lambda ms: (
gbps = lambda ms: (
(
(2 * M * N * K - M * N) * a.element_size()
+ (3 * M * N) * scale_a.element_size()
)
@@ -38,9 +38,9 @@ def cutlass_mla_decode(
) -> torch.Tensor:
assert q_nope.ndim == 3, f"q_nope must be a 3D tensor, but got {q_nope.ndim}"
assert q_pe.ndim == 3, f"q_pe must be a 3D tensor, but got {q_pe.ndim}"
assert (
kv_c_and_k_pe_cache.ndim == 3
), f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}"
assert kv_c_and_k_pe_cache.ndim == 3, (
f"kv_c_and_k_pe_cache must be a 3D tensor, but got {kv_c_and_k_pe_cache.ndim}"
)
B_q, H, D_q_nope = q_nope.shape
B_q_2, H_2, D_q_pe = q_pe.shape
@@ -77,12 +77,12 @@ def cutlass_mla_decode(
torch.bfloat16,
), f"q_nope.dtype needs to be fp16 or bf16 but got {q_nope.dtype}."
assert q_nope.dtype == q_pe.dtype == kv_c_and_k_pe_cache.dtype
assert (
seq_lens.dtype == torch.int32
), f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}."
assert (
page_table.dtype == torch.int32
), f"page_table.dtype needs to be int32 but got {page_table.dtype}."
assert seq_lens.dtype == torch.int32, (
f"seq_lens.dtype needs to be int32 but got {seq_lens.dtype}."
)
assert page_table.dtype == torch.int32, (
f"page_table.dtype needs to be int32 but got {page_table.dtype}."
)
out = q_nope.new_empty((B_q, MAX_HEADS, D_latent))
@@ -247,12 +247,12 @@ def gemma_fused_add_rmsnorm(
def _check_shape(input: torch.Tensor, output: torch.Tensor) -> None:
assert input.ndim == output.ndim, f"{input.ndim} != {output.ndim}"
assert (
input.shape[:-1] == output.shape[:-1]
), f"{input.shape[:-1]} != {output.shape[:-1]}"
assert (
input.shape[-1] == 2 * output.shape[-1]
), f"{input.shape[-1]} != {2 * output.shape[-1]}"
assert input.shape[:-1] == output.shape[:-1], (
f"{input.shape[:-1]} != {output.shape[:-1]}"
)
assert input.shape[-1] == 2 * output.shape[-1], (
f"{input.shape[-1]} != {2 * output.shape[-1]}"
)
def silu_and_mul(input: torch.Tensor, out: torch.Tensor = None) -> torch.Tensor:
@@ -159,9 +159,9 @@ def flash_mla_with_kvcache(
assert extra_topk_length is None
if indices is not None:
assert causal == False, "causal must be `false` if sparse attention is enabled."
assert (descale_q is None) == (
descale_k is None
), "descale_q and descale_k should be both None or both not None"
assert (descale_q is None) == (descale_k is None), (
"descale_q and descale_k should be both None or both not None"
)
if indices is None and q.element_size() == 1:
out, softmax_lse = torch.ops.sgl_kernel.fwd_kvcache_mla_fp8.default(
@@ -257,9 +257,9 @@ def _flash_mla_with_kvcache_sched_meta(
assert sched_meta.config.causal == causal, helper_msg
assert sched_meta.config.is_fp8_kvcache == is_fp8_kvcache, helper_msg
assert sched_meta.config.topk == topk, helper_msg
assert (
sched_meta.config.extra_page_block_size == extra_page_block_size
), helper_msg
assert sched_meta.config.extra_page_block_size == extra_page_block_size, (
helper_msg
)
assert sched_meta.config.extra_topk == extra_topk, helper_msg
if topk is not None:
@@ -76,11 +76,11 @@ def rope_pool_fused(
if q_shape != (q_shape[0], num_qo_heads, head_dim):
raise ValueError(
"q shape must be [num_tokens, num_qo_heads, head_dim], " f"got {q.shape}"
f"q shape must be [num_tokens, num_qo_heads, head_dim], got {q.shape}"
)
if k_shape != (q_shape[0], num_kv_heads, head_dim):
raise ValueError(
"k shape must be [num_tokens, num_kv_heads, head_dim], " f"got {k.shape}"
f"k shape must be [num_tokens, num_kv_heads, head_dim], got {k.shape}"
)
if v_shape != k_shape:
raise ValueError(f"v shape must match k shape, got {v.shape} vs {k.shape}")
@@ -86,9 +86,9 @@ def musa_fused_gemv(
out_shape = x.shape[:-1] + (
qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
)
assert not (
use_swigelu and use_rms_norm
), "gemv only fused one activation (swigelu or rms_norm)!"
assert not (use_swigelu and use_rms_norm), (
"gemv only fused one activation (swigelu or rms_norm)!"
)
if use_rms_norm:
if gamma is None:
@@ -113,9 +113,9 @@ def musa_fused_gemv(
return output
# w4a16 gemv
elif qweight_scales is not None:
assert (
x.dtype == torch.bfloat16 or x.dtype == torch.float16
), "W4A16 gemv only support bfloat16 or float16!"
assert x.dtype == torch.bfloat16 or x.dtype == torch.float16, (
"W4A16 gemv only support bfloat16 or float16!"
)
use_int4_w4a16 = True
out_shape = x.shape[:-1] + (
qweight.shape[0] if not use_swigelu else qweight.shape[0] // 2,
@@ -70,9 +70,9 @@ class ScalarType:
"""
def _floating_point_max_int(self) -> int:
assert (
self.mantissa <= 52 and self.exponent <= 11
), f"Cannot represent max/min as a double for type {self.__str__()}"
assert self.mantissa <= 52 and self.exponent <= 11, (
f"Cannot represent max/min as a double for type {self.__str__()}"
)
max_mantissa = (1 << self.mantissa) - 1
if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN:
@@ -80,9 +80,9 @@ class ScalarType:
max_exponent = (1 << self.exponent) - 2
if self.nan_repr == NanRepr.EXTD_RANGE_MAX_MIN or self.nan_repr == NanRepr.NONE:
assert (
self.exponent < 11
), f"Cannot represent max/min as a double for type {self.__str__()}"
assert self.exponent < 11, (
f"Cannot represent max/min as a double for type {self.__str__()}"
)
max_exponent = max_exponent + 1
# adjust the exponent to match that of a double
@@ -109,25 +109,25 @@ class ScalarType:
if self.is_floating_point():
return self._floating_point_max()
else:
assert (
self.size_bits < 64 or self.size_bits == 64 and self.is_signed()
), "Cannot represent max as an int"
assert self.size_bits < 64 or self.size_bits == 64 and self.is_signed(), (
"Cannot represent max as an int"
)
return (1 << self.mantissa) - 1
def _raw_min(self) -> Union[int, float]:
if self.is_floating_point():
assert (
self.is_signed()
), "We currently assume all floating point types are signed"
assert self.is_signed(), (
"We currently assume all floating point types are signed"
)
sign_bit_double = 1 << 63
max_raw = self._floating_point_max_int()
min_raw = max_raw | sign_bit_double
return struct.unpack("!d", struct.pack("!Q", min_raw))[0]
else:
assert (
not self.is_signed() or self.size_bits <= 64
), "Cannot represent min as a int64_t"
assert not self.is_signed() or self.size_bits <= 64, (
"Cannot represent min as a int64_t"
)
if self.is_signed():
return -(1 << (self.size_bits - 1))
@@ -94,9 +94,9 @@ def _compute_imbalanced_split(
def assert_all_close_or_tiny_diff(a: torch.Tensor, b: torch.Tensor):
assert (a.shape == b.shape) and (
a.dtype == b.dtype
), f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}"
assert (a.shape == b.shape) and (a.dtype == b.dtype), (
f"{a.shape=} {b.shape=} {a.dtype=} {b.dtype=}"
)
numel = a.numel()
if a.dtype == torch.float8_e4m3fn:
@@ -112,9 +112,9 @@ class RotaryEmbedding(torch.nn.Module):
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
"""A PyTorch-native implementation of forward()."""
assert (
fused_set_kv_buffer_arg is None
), "fused_set_kv_buffer_arg is not supported for native implementation"
assert fused_set_kv_buffer_arg is None, (
"fused_set_kv_buffer_arg is not supported for native implementation"
)
if offsets is not None:
positions = positions + offsets
@@ -182,9 +182,9 @@ class SglKernelRotaryEmbedding(RotaryEmbedding):
offsets: Optional[torch.Tensor] = None,
fused_set_kv_buffer_arg: Optional[FusedSetKVBufferArg] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
assert (
fused_set_kv_buffer_arg is None
), "fused_set_kv_buffer_arg is not supported for sgl-kernel implementation"
assert fused_set_kv_buffer_arg is None, (
"fused_set_kv_buffer_arg is not supported for sgl-kernel implementation"
)
if self.cos_sin_cache.dtype != query.dtype:
self.cos_sin_cache = self.cos_sin_cache.to(query.dtype)
torch.ops.sgl_kernel.rotary_embedding(
@@ -33,9 +33,9 @@ def fast_topk_v2(
Returns:
The topk indices tensor of shape (B, topk)
"""
assert (
topk == 2048
), "fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048"
assert topk == 2048, (
"fast_topk_v2 is only optimized for deepseek v3.2 model, where topk=2048"
)
assert score.dim() == 2
topk_indices = score.new_empty((score.size(0), topk), dtype=torch.int32)
torch.ops.sgl_kernel.fast_topk(score, topk_indices, lengths, row_starts)
@@ -68,9 +68,9 @@ def fast_topk_transform_fused(
Returns:
The topk indices tensor of shape (B, topk)
"""
assert (
topk == 2048
), "fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048"
assert topk == 2048, (
"fast_topk_transform_fused is only optimized for deepseek v3.2 model, where topk=2048"
)
assert score.dim() == 2
src_page_table = page_table_size_1
dst_page_table = score.new_empty((score.shape[0], topk), dtype=torch.int32)
@@ -138,9 +138,9 @@ def fast_topk_transform_ragged_fused(
Returns:
The topk indices tensor of shape (B, topk)
"""
assert (
topk == 2048
), "fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048"
assert topk == 2048, (
"fast_topk_transform_ragged_fused is only optimized for deepseek v3.2 model, where topk=2048"
)
assert score.dim() == 2
topk_indices_ragged = score.new_empty((score.shape[0], topk), dtype=torch.int32)
torch.ops.sgl_kernel.fast_topk_transform_ragged_fused(
@@ -116,15 +116,15 @@ def test_tree_speculative_sampling_target_only(
deterministic=True,
)
assert (
predicts.tolist() == expected_predicts
), f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})"
assert (
accept_index.tolist() == expected_accept_index
), f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})"
assert (
accept_token_num.tolist() == expected_accept_token_num
), f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})"
assert predicts.tolist() == expected_predicts, (
f"Predicts mismatch for thresholds ({threshold_single}, {threshold_acc})"
)
assert accept_index.tolist() == expected_accept_index, (
f"Accept index mismatch for thresholds ({threshold_single}, {threshold_acc})"
)
assert accept_token_num.tolist() == expected_accept_token_num, (
f"Accept token num mismatch for thresholds ({threshold_single}, {threshold_acc})"
)
if __name__ == "__main__":
@@ -92,9 +92,9 @@ def multi_process_parallel(
for i in range(world_size):
procs[i].join()
assert (
procs[i].exitcode == 0
), f"Process {i} failed with exit code {procs[i].exitcode}"
assert procs[i].exitcode == 0, (
f"Process {i} failed with exit code {procs[i].exitcode}"
)
class TestCustomAllReduce(unittest.TestCase):
@@ -251,12 +251,14 @@ def test_sparse_attention(
ref_out, ref_lse = ref_attn(q, k, v)
torch.testing.assert_close(
out, ref_out, atol=2e-2, rtol=1e-2
), f"{torch.max(torch.abs(out - ref_out))}"
torch.testing.assert_close(
lse, ref_lse, atol=2e-2, rtol=1e-2
), f"{torch.max(torch.abs(lse - ref_lse))}"
(
torch.testing.assert_close(out, ref_out, atol=2e-2, rtol=1e-2),
f"{torch.max(torch.abs(out - ref_out))}",
)
(
torch.testing.assert_close(lse, ref_lse, atol=2e-2, rtol=1e-2),
f"{torch.max(torch.abs(lse - ref_lse))}",
)
# sparse attention utils
@@ -198,9 +198,7 @@ def reference_torch_prefill(
kvs = torch.index_select(
kvs, 0, indices.masked_fill(invalid_indices_mask, 0).flatten()
).view(
s_q, topk, 576
) # [s_q, topk, d_qk]
).view(s_q, topk, 576) # [s_q, topk, d_qk]
attn_score = qs @ kvs.transpose(1, 2) # [s_q, h_q, topk]
attn_score.masked_fill_(invalid_indices_mask.unsqueeze(1), float("-inf"))
attn_score *= sm_scale * math.log2(math.e)
@@ -76,9 +76,9 @@ def torch_ref_rms_norm_rope(
v_size = num_heads_v * head_dim
# Verify dimensions match
assert (
hidden_size == q_size + k_size + v_size
), f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}"
assert hidden_size == q_size + k_size + v_size, (
f"Hidden size {hidden_size} doesn't match Q+K+V dimensions {q_size + k_size + v_size}"
)
# Split the tensor into Q, K, V parts
q = qkv[:, :q_size]
@@ -44,13 +44,13 @@ def test_topk_sigmoid(num_tokens, num_experts, topk):
topk_weights_ref, topk_indices_ref = torch.topk(sigmoid_output, topk, dim=-1)
# Verify the top-k weights and indices match the torch native ones
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
)
assert torch.allclose(
topk_indices_ref.int(), topk_indices, atol=0, rtol=0
), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
)
@pytest.mark.parametrize(
@@ -87,13 +87,13 @@ def test_topk_sigmoid_dtype_regression(num_tokens, num_experts, topk, dtype):
gating_output.float(),
)
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: SGLang old interface={topk_weights_ref} vs SGLang new interface={topk_weights}"
)
assert torch.allclose(
topk_indices_ref.int(), topk_indices, atol=0, rtol=0
), f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}"
assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
f"Indices mismatch: SGLang old interface={topk_indices_ref}, SGLang new interface={topk_indices}"
)
@pytest.mark.parametrize(
@@ -136,13 +136,13 @@ def test_topk_sigmoid_renormalize(num_tokens, num_experts, topk):
)
topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: SGLang w/o fused renormalize={topk_weights_ref} vs SGLang w/ fused renormalize={topk_weights}"
)
assert torch.allclose(
topk_indices_ref.int(), topk_indices, atol=0, rtol=0
), f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}"
assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
f"Indices mismatch: SGLang w/o fused renormalize={topk_indices_ref}, SGLang w/ fused renormalize={topk_indices}"
)
@pytest.mark.parametrize(
@@ -180,13 +180,13 @@ def test_topk_sigmoid_renormalize_correction_bias(num_tokens, num_experts, topk)
topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
# Verify the top-k weights and indices match the torch native ones
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: torch={topk_weights_ref} vs SGLang={topk_weights}"
)
assert torch.allclose(
topk_indices_ref.int(), topk_indices, atol=0, rtol=0
), f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
assert torch.allclose(topk_indices_ref.int(), topk_indices, atol=0, rtol=0), (
f"Indices mismatch: torch={topk_indices_ref}, SGLang={topk_indices}"
)
if __name__ == "__main__":
@@ -41,13 +41,13 @@ def test_topkfast_softmax(num_tokens, num_experts, topk):
topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
# Verify the top-k weights and indices match the torch native ones
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
)
assert compare_topk_values(
gating_output, topk_indices_ref.int(), topk_indices
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
)
@pytest.mark.parametrize(
@@ -79,13 +79,13 @@ def test_topk_softmax(num_tokens, num_experts, topk):
topk_weights_ref, topk_indices_ref = torch.topk(softmax_output, topk, dim=-1)
# Verify the top-k weights and indices match the torch native ones
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: torch={topk_indices_ref} vs SGLang={topk_weights}"
)
assert compare_topk_values(
gating_output, topk_indices_ref.int(), topk_indices
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
)
@pytest.mark.parametrize(
@@ -122,13 +122,13 @@ def test_topk_softmax_dtype_regression(num_tokens, num_experts, topk, dtype):
gating_output.float(),
)
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: SGLang old interface={topk_indices_ref} vs SGLang new interface={topk_weights}"
)
assert compare_topk_values(
gating_output, topk_indices_ref.int(), topk_indices
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
)
@pytest.mark.parametrize(
@@ -171,13 +171,13 @@ def test_topk_softmax_renormalize(num_tokens, num_experts, topk):
)
topk_weights_ref = topk_weights_ref / topk_weights_ref.sum(dim=-1, keepdim=True)
assert torch.allclose(
topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3
), f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}"
assert torch.allclose(topk_weights_ref, topk_weights, atol=1e-3, rtol=1e-3), (
f"Weights mismatch: SGLang w/o fused renormalize={topk_indices_ref} vs SGLang w/ fused renormalize={topk_weights}"
)
assert compare_topk_values(
gating_output, topk_indices_ref.int(), topk_indices
), f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
assert compare_topk_values(gating_output, topk_indices_ref.int(), topk_indices), (
f"Values at the two indices are not equal: torch={topk_indices_ref}, SGLang={topk_indices}, values={gating_output}"
)
if __name__ == "__main__":
+3 -3
View File
@@ -72,9 +72,9 @@ def generate_clangd():
arch = make_jit_cuda_arch(int(major), int(minor))
else:
arch = get_jit_cuda_arch()
assert (
arch.major > 0
), "Cannot detect CUDA architecture, please specify --cuda-target explicitly."
assert arch.major > 0, (
"Cannot detect CUDA architecture, please specify --cuda-target explicitly."
)
compile_flags = [
"-xcuda",
+10 -11
View File
@@ -253,9 +253,9 @@ class Benchmark(Generic[F]):
f"parametrize name {name!r} is not a parameter of "
f"{self._fn.__name__}; available: {list(self._fn_params)}"
)
assert (
name not in self._seen_args
), f"parametrize name {name!r} is already used"
assert name not in self._seen_args, (
f"parametrize name {name!r} is already used"
)
self._seen_args.add(name)
self._configs.insert(0, (names, vals))
@@ -305,8 +305,7 @@ class Benchmark(Generic[F]):
if p.default is inspect.Parameter.empty and p.kind in kinds
} - (set(flat_names) | {self._line_arg})
assert not missing, (
f"parameters not parametrized for {self._fn.__name__}: "
f"{sorted(missing)}"
f"parameters not parametrized for {self._fn.__name__}: {sorted(missing)}"
)
results, bandwidths, should_log_bw = self._collect_results()
@@ -360,13 +359,13 @@ def parametrize(names: str, vals: List[Any], ci_vals: Optional[List[Any]] = None
return [(v,) for v in vs]
out: List[Tuple[Any, ...]] = []
for v in vs:
assert isinstance(
v, (tuple, list)
), f"parametrize: multi-name values must be tuples, got {v!r}"
assert isinstance(v, (tuple, list)), (
f"parametrize: multi-name values must be tuples, got {v!r}"
)
t = tuple(v)
assert (
len(t) == arity
), f"parametrize: each value must have length {arity}, got {t!r}"
assert len(t) == arity, (
f"parametrize: each value must have length {arity}, got {t!r}"
)
out.append(t)
return out
@@ -121,7 +121,7 @@ def load_jit(
# Also the benign case where a concurrent GC unlinked the leaf
# between the lookup and the load.
logger.warning(
"Cached JIT module %s failed to load; rebuilding. " "Got error: %s",
"Cached JIT module %s failed to load; rebuilding. Got error: %s",
spec.module_name,
e,
)
@@ -25,7 +25,7 @@ def _jit_causal_conv3d_cat_pad_module(dtype: torch.dtype) -> Module:
cuda_wrappers=[
(
"causal_conv3d_cat_pad",
"causal_conv3d_cat_pad::" f"CausalConv3dCatPadKernel<{args}>::run",
f"causal_conv3d_cat_pad::CausalConv3dCatPadKernel<{args}>::run",
)
],
)
@@ -353,9 +353,9 @@ class _Qwen3xNvfp4Sm120Kernel:
self.occupancy,
)
assert (
self.epi_stage > 0
), "epi_stage <= 0, not enough shared memory. This configuration will be skipped."
assert self.epi_stage > 0, (
"epi_stage <= 0, not enough shared memory. This configuration will be skipped."
)
(
self.a_smem_layout_staged,
@@ -34,11 +34,11 @@ def _jit_residual_gate_add_module(dtype: torch.dtype) -> Module:
cuda_wrappers=[
(
"residual_gate_add",
"residual_gate_add::" f"ResidualGateAddKernel<{args}>::run",
f"residual_gate_add::ResidualGateAddKernel<{args}>::run",
),
(
"residual_gate_add_transposed",
"residual_gate_add::" f"ResidualGateAddKernel<{args}>::run_transposed",
f"residual_gate_add::ResidualGateAddKernel<{args}>::run_transposed",
),
],
)
@@ -157,9 +157,9 @@ def run_unary_activation(
Unlike :func:`run_activation`, there is no gate/up split ``input`` and
``out`` share the same shape.
"""
assert (
op_name in SUPPORTED_UNARY_ACTIVATIONS
), f"Unsupported unary activation: {op_name}"
assert op_name in SUPPORTED_UNARY_ACTIVATIONS, (
f"Unsupported unary activation: {op_name}"
)
if out is None:
out = torch.empty_like(input)
_run_unary_activation_inplace(op_name, input, out)
@@ -101,9 +101,9 @@ def softcap_inplace_logits(full_logits, final_logit_softcapping):
row_stride = ncols
else:
assert full_logits.ndim == 2, "non-contiguous softcap requires 2D tensor"
assert (
full_logits.stride(1) == 1
), "non-contiguous softcap requires contiguous columns"
assert full_logits.stride(1) == 1, (
"non-contiguous softcap requires contiguous columns"
)
nrows, ncols = full_logits.shape
row_stride = full_logits.stride(0)
@@ -221,12 +221,12 @@ class FP8MQALogitsKernel:
self.block_kv = block_kv
self.phys_block_kv = phys_block_kv
self.num_blocks_per_mma = block_kv // phys_block_kv
assert (
block_kv % phys_block_kv == 0
), f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
assert (
self.num_blocks_per_mma <= 4
), f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4"
assert block_kv % phys_block_kv == 0, (
f"block_kv={block_kv} must be divisible by phys_block_kv={phys_block_kv}"
)
assert self.num_blocks_per_mma <= 4, (
f"num_blocks_per_mma={self.num_blocks_per_mma} exceeds max 4"
)
self.remove_kv_wait_in_epilogue = remove_kv_wait_in_epilogue
self.early_tmem_copy = early_tmem_copy
self.smem_subpartition_opt = smem_subpartition_opt
@@ -3080,9 +3080,9 @@ def gated_delta_rule_mtp_wide_vec(
assert K_val == 128 and V_val == 128
assert initial_state_source.dtype == torch.bfloat16
assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
assert (
V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0
), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout"
assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, (
f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout"
)
if cache_ring:
assert replayssm_rawv is not None and replayssm_rawk is not None
@@ -3194,13 +3194,13 @@ def gated_delta_rule_mtp_wide_vec(
)
# Validate recovery_steps for fused recovery+decode mode.
assert (
0 <= recovery_steps <= T_val
), f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}"
assert 0 <= recovery_steps <= T_val, (
f"recovery_steps must be in [0, T={T_val}], got {recovery_steps}"
)
if recovery_steps > 0:
assert (
not cache_intermediate_states
), "recovery_steps > 0 is incompatible with intermediate state caching"
assert not cache_intermediate_states, (
"recovery_steps > 0 is incompatible with intermediate state caching"
)
assert not disable_state_update, (
"recovery_steps > 0 requires state writeback "
"(disable_state_update=False); the boundary writeback at i_t=K-1 "
@@ -3220,12 +3220,12 @@ def gated_delta_rule_mtp_wide_vec(
# accepted_steps[i] is the per-request phase boundary.
per_request_accepted_steps = accepted_steps is not None
if per_request_accepted_steps:
assert accepted_steps.shape == (
B_val,
), f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}"
assert (
accepted_steps.dtype == torch.int32
), f"accepted_steps must be int32, got {accepted_steps.dtype}"
assert accepted_steps.shape == (B_val,), (
f"accepted_steps must have shape [B={B_val}], got {accepted_steps.shape}"
)
assert accepted_steps.dtype == torch.int32, (
f"accepted_steps must be int32, got {accepted_steps.dtype}"
)
assert accepted_steps.device == q.device
# FLA-style per-token pool scatter (vLLM API compat). When the public
@@ -3236,15 +3236,15 @@ def gated_delta_rule_mtp_wide_vec(
# this entry point hit the same fail-fast errors.
per_token_pool_scatter = ssm_state_indices is not None
if per_token_pool_scatter:
assert (
intermediate_states_buffer is None
), "ssm_state_indices and intermediate_states_buffer are mutually exclusive"
assert (
not disable_state_update
), "ssm_state_indices requires state writes; disable_state_update must be False"
assert (
recovery_steps == 0
), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
assert intermediate_states_buffer is None, (
"ssm_state_indices and intermediate_states_buffer are mutually exclusive"
)
assert not disable_state_update, (
"ssm_state_indices requires state writes; disable_state_update must be False"
)
assert recovery_steps == 0, (
"ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
)
assert T_val >= 2, (
f"ssm_state_indices requires T >= 2 (got T={T_val}); "
f"for T=1 use output_state_indices"
@@ -3253,9 +3253,9 @@ def gated_delta_rule_mtp_wide_vec(
f"ssm_state_indices must have shape [B={B_val}, T={T_val}], "
f"got {tuple(ssm_state_indices.shape)}"
)
assert (
ssm_state_indices.dtype == torch.int32
), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
assert ssm_state_indices.dtype == torch.int32, (
f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
)
assert ssm_state_indices.device == q.device
phase_b_unroll = _select_wide_vec_phase_b_unroll(
@@ -3517,9 +3517,9 @@ def gated_delta_rule_t1_wide_vec(
assert K_val == 128 and V_val == 128
assert initial_state_source.dtype == torch.bfloat16
assert tile_v in (32, 64, 128), f"tile_v must be 32/64/128, got {tile_v}"
assert (
V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0
), f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout"
assert V_val % tile_v == 0 and (tile_v // NUM_GROUPS) % ILP_ROWS == 0, (
f"tile_v={tile_v} incompatible with 8 groups × ILP=4 layout"
)
if scale is None:
scale = 1.0 / math.sqrt(K_val)
@@ -3827,9 +3827,9 @@ def gated_delta_rule_mtp(
f"intermediate_states_buffer dim 0 ({buffer_size}) must equal "
f"batch size B={B}; the buffer is batch-scoped, not pool-scoped"
)
assert (
cache_steps >= T
), f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}"
assert cache_steps >= T, (
f"intermediate_states_buffer dim 1 ({cache_steps}) must be >= T={T}"
)
assert intermediate_states_buffer.dtype == torch.bfloat16
intermediate_states = intermediate_states_buffer.reshape(
B * cache_steps * HV, V, K
@@ -3860,28 +3860,28 @@ def gated_delta_rule_mtp(
# results/2026-06-03/FLA_SCATTER_MODE_PLAN.md.
per_token_pool_scatter = ssm_state_indices is not None
if per_token_pool_scatter:
assert (
intermediate_states_buffer is None
), "ssm_state_indices and intermediate_states_buffer are mutually exclusive"
assert (
not disable_state_update
), "ssm_state_indices requires state writes; disable_state_update must be False"
assert (
recovery_steps == 0
), "ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
assert (
T >= 2
), f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices"
assert intermediate_states_buffer is None, (
"ssm_state_indices and intermediate_states_buffer are mutually exclusive"
)
assert not disable_state_update, (
"ssm_state_indices requires state writes; disable_state_update must be False"
)
assert recovery_steps == 0, (
"ssm_state_indices + recovery_steps>0 not yet supported (MVP exclusion)"
)
assert T >= 2, (
f"ssm_state_indices requires T >= 2 (got T={T}); for T=1 use output_state_indices"
)
assert ssm_state_indices.shape == (B, T), (
f"ssm_state_indices must have shape [B={B}, T={T}], "
f"got {tuple(ssm_state_indices.shape)}"
)
assert (
ssm_state_indices.dtype == torch.int32
), f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
assert (
ssm_state_indices.device == q.device
), f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}"
assert ssm_state_indices.dtype == torch.int32, (
f"ssm_state_indices must be int32, got {ssm_state_indices.dtype}"
)
assert ssm_state_indices.device == q.device, (
f"ssm_state_indices device {ssm_state_indices.device} != q device {q.device}"
)
# Dispatch to the wide_vec kernel when work_units (B*HV) amortizes its
# lower per-CTA parallelism. ``_select_wide_vec_tile_v`` picks tile_v
@@ -3960,12 +3960,12 @@ def gated_delta_rule_mtp(
# Per-request K opt-in (see gated_delta_rule_mtp_wide_vec for full rationale).
per_request_accepted_steps = accepted_steps is not None
if per_request_accepted_steps:
assert accepted_steps.shape == (
B,
), f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}"
assert (
accepted_steps.dtype == torch.int32
), f"accepted_steps must be int32, got {accepted_steps.dtype}"
assert accepted_steps.shape == (B,), (
f"accepted_steps must have shape [B={B}], got {accepted_steps.shape}"
)
assert accepted_steps.dtype == torch.int32, (
f"accepted_steps must be int32, got {accepted_steps.dtype}"
)
assert accepted_steps.device == q.device
# Contiguous pool -> sentinel keys + slot dim marked dynamic (pool-size
@@ -1427,12 +1427,12 @@ def cutedsl_fused_sigmoid_gating_kda_update(
N = initial_state_indices.shape[0]
assert K == TILE_K, f"Current CuTe DSL KDA kernel requires K={TILE_K}, got {K}"
assert (
V % TILE_V_SMALL == 0
), f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}"
assert (
V % TILE_V == 0
), f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}"
assert V % TILE_V_SMALL == 0, (
f"Current CuTe DSL KDA kernel requires V % {TILE_V_SMALL} == 0, got V={V}"
)
assert V % TILE_V == 0, (
f"Current CuTe DSL KDA kernel requires V % {TILE_V} == 0, got V={V}"
)
assert (V // TILE_V_SMALL) % NUM_BLOCKS_PER_STATE_SMALL == 0, (
"Small-batch KDA kernel requires num_v_tiles_small divisible by "
f"{NUM_BLOCKS_PER_STATE_SMALL}, got V={V}"
@@ -1483,7 +1483,6 @@ def _lean_attention_decode_kernel(
# Use a regular while loop instead of tl.static_range with a dynamic bound to avoid
# Triton compiler crashes in the Coalesce pass (max_output_tile_cnt is runtime-computed).
while iter < cta_end_tile_gid:
tile_row_idx = iter // tiles_per_khead
tile_idx = tile_row_idx * batch_size
tile_iter = tile_row_idx * tiles_per_khead
@@ -354,9 +354,9 @@ def apply_rotary_emb_triton(
grid = (batch_size, n_heads if is_3d else 1, num_blocks_dim)
if positions is not None:
assert positions.shape == (
batch_size,
), f"positions shape {positions.shape} != ({batch_size},)"
assert positions.shape == (batch_size,), (
f"positions shape {positions.shape} != ({batch_size},)"
)
apply_rotary_emb_triton_kernel[grid](
x,
@@ -374,9 +374,9 @@ def apply_rotary_emb_triton(
BLOCK_SIZE=BLOCK_SIZE,
)
else:
assert (
freqs_real.shape[0] == batch_size
), f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}"
assert freqs_real.shape[0] == batch_size, (
f"freqs_cis batch size {freqs_real.shape[0]} != x batch size {batch_size}"
)
apply_rotary_emb_triton_kernel[grid](
x,
@@ -621,9 +621,9 @@ def fused_norm_rope_inplace_triton(
if weight is not None:
assert weight.shape == (head_dim,)
if positions is None:
assert (
freqs_real.shape[0] == M
), f"freqs_cis row count {freqs_real.shape[0]} != M={M}"
assert freqs_real.shape[0] == M, (
f"freqs_cis row count {freqs_real.shape[0]} != M={M}"
)
else:
assert positions.shape == (M,) and positions.dim() == 1
@@ -181,9 +181,9 @@ def dequantize_k_cache_paged(
output: [num_tokens, 1, dim_nope + dim_rope], the de-quantized k-cache
"""
dim_quant = quant_k_cache.shape[-1]
assert (
dim_quant == 656
), f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
assert dim_quant == 656, (
f"dim_quant: {dim_quant} != 656 detected in dequantize_k_cache_paged"
)
quant_k_cache = quant_k_cache.view((-1, dim_quant))
# num_tokens can exceed kv_cache_size due to prefix sharing (multiple seqs share same KV slots)
@@ -308,9 +308,9 @@ def _set_k_and_s_triton(
assert scale_dim == 1
if _is_hip:
if _use_aiter_preshuffle:
assert (
page_size % 16 == 0
), f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
assert page_size % 16 == 0, (
f"HIP preshuffle requires page_size to be a multiple of 16, got {page_size}"
)
else:
assert page_size == 64
@@ -155,9 +155,9 @@ def act_quant(
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
assert x.size(-1) % block_size == 0, (
f"Last dimension size must be divisible by block_size (block_size={block_size})"
)
N = x.size(-1)
if _is_fp8_fnuz:
y = torch.empty_like(x, dtype=torch.float8_e4m3fnuz)
@@ -272,16 +272,16 @@ def sparse_attention_fwd_kernel_v1(
num_stages=2,
threads=256,
):
assert dim == tilelang.math.next_power_of_2(
dim
), f"haven't check padding correctness yet, dim={dim}"
assert tail_dim == tilelang.math.next_power_of_2(
tail_dim
), f"haven't check padding correctness yet, dim={tail_dim}"
assert dim == tilelang.math.next_power_of_2(dim), (
f"haven't check padding correctness yet, dim={dim}"
)
assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
f"haven't check padding correctness yet, dim={tail_dim}"
)
assert is_causal == True, "non-casual is not supported"
assert (
topk % block_I == 0
), "otherwise will load some index=0 thus causing wrong kv to be loaded"
assert topk % block_I == 0, (
"otherwise will load some index=0 thus causing wrong kv to be loaded"
)
if sm_scale is None:
sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e)
else:
@@ -361,7 +361,6 @@ def sparse_attention_fwd_kernel_v1(
T.copy(Q[b_i, s_i, H0:H1, D:], Q_tail_shared)
for i_i in T.Pipelined(NI, num_stages=num_stages):
for bi_i in T.Parallel(BI):
mask[bi_i] = Indices[b_i, s_i, g_i, i_i * BI + bi_i] >= 0
@@ -446,15 +445,15 @@ def sparse_attention_fwd_kernel_v2(
sm_scale: Optional[float] = None,
block_I: int = 64,
):
assert dim == tilelang.math.next_power_of_2(
dim
), f"haven't check padding correctness yet, dim={dim}"
assert tail_dim == tilelang.math.next_power_of_2(
tail_dim
), f"haven't check padding correctness yet, dim={tail_dim}"
assert (
topk % block_I == 0
), "otherwise will load some index=0 thus causing wrong kv to be loaded"
assert dim == tilelang.math.next_power_of_2(dim), (
f"haven't check padding correctness yet, dim={dim}"
)
assert tail_dim == tilelang.math.next_power_of_2(tail_dim), (
f"haven't check padding correctness yet, dim={tail_dim}"
)
assert topk % block_I == 0, (
"otherwise will load some index=0 thus causing wrong kv to be loaded"
)
if sm_scale is None:
sm_scale = (1.0 / (dim + tail_dim)) ** 0.5 * 1.44269504 # log2(e)
else:
@@ -1078,9 +1077,9 @@ def sparse_mla_fwd_decode_partial_fp8(
threads=256,
):
assert d_v == 512, f"only support d_v=512"
assert (
topk % block_I == 0
), "otherwise will load some index=0 thus causing wrong kv to be loaded"
assert topk % block_I == 0, (
"otherwise will load some index=0 thus causing wrong kv to be loaded"
)
# Softmax scores are in [0, 1]. We scale by fp8_max_val before FP8 cast
# to better utilize FP8 dynamic range, then apply the inverse scale after GEMM.
@@ -1104,9 +1103,9 @@ def sparse_mla_fwd_decode_partial_fp8(
h_per_block = 16
# Match bf16 partial behavior: keep fixed 16-head tiles and use
# sliced T.copy on H0:H1 for tail handling.
assert (
num_heads <= h_per_block or num_heads % h_per_block == 0
), "num_heads must be <=16 or divisible by 16"
assert num_heads <= h_per_block or num_heads % h_per_block == 0, (
"num_heads must be <=16 or divisible by 16"
)
head_blocks_per_seq = (num_heads + h_per_block - 1) // h_per_block
batch = 1
@@ -1594,9 +1593,7 @@ def dpsk_v4_fp8_partial_kernel(
sm_scale = sm_scale * log2e
assert dim == 448 and tail_dim == 64
assert topk_1 % block_I == 0
assert (
topk_1 // block_I
) % inner_iter_1 == 0, (
assert (topk_1 // block_I) % inner_iter_1 == 0, (
f"NI_1={topk_1 // block_I} must be divisible by inner_iter_1={inner_iter_1}"
)
assert block_size_kv_1 > 0 and (block_size_kv_1 & (block_size_kv_1 - 1)) == 0
@@ -1605,9 +1602,7 @@ def dpsk_v4_fp8_partial_kernel(
if is_dual:
assert inner_iter_2 > 0, "dual-cache call requires inner_iter_2 > 0"
assert topk_2 % block_I == 0
assert (
topk_2 // block_I
) % inner_iter_2 == 0, (
assert (topk_2 // block_I) % inner_iter_2 == 0, (
f"NI_2={topk_2 // block_I} must be divisible by inner_iter_2={inner_iter_2}"
)
assert block_size_kv_2 > 0 and (block_size_kv_2 & (block_size_kv_2 - 1)) == 0
@@ -2256,12 +2251,8 @@ def dpsk_v4_combine_kernel(
@T.prim_func
def main(
Partial_O: T.Tensor(
[batch, seq_len, n_groups, num_heads, DT], BF16
), # type: ignore
Partial_LSE: T.Tensor(
[batch, seq_len, n_groups, num_heads], accum_dtype
), # type: ignore
Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore
Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore
Topk_length_1: T.Tensor([batch], INT32), # type: ignore
Topk_length_2: T.Tensor([batch], INT32), # type: ignore
Attn_sink: T.Tensor([num_heads], FP32), # type: ignore
@@ -2369,12 +2360,8 @@ def dpsk_v4_combine_kernel(
@T.prim_func
def main(
Partial_O: T.Tensor(
[batch, seq_len, n_groups, num_heads, DT], BF16
), # type: ignore
Partial_LSE: T.Tensor(
[batch, seq_len, n_groups, num_heads], accum_dtype
), # type: ignore
Partial_O: T.Tensor([batch, seq_len, n_groups, num_heads, DT], BF16), # type: ignore
Partial_LSE: T.Tensor([batch, seq_len, n_groups, num_heads], accum_dtype), # type: ignore
Attn_sink: T.Tensor([num_heads], FP32), # type: ignore
Output: T.Tensor([batch, seq_len, num_heads, DT], BF16), # type: ignore
LSE: T.Tensor([batch, seq_len, num_heads], accum_dtype), # type: ignore
@@ -99,9 +99,9 @@ def act_quant(
- A tensor of scaling factors with dtype `torch.float32`.
"""
assert x.is_contiguous(), "Input tensor must be contiguous"
assert (
x.size(-1) % block_size == 0
), f"Last dimension size must be divisible by block_size (block_size={block_size})"
assert x.size(-1) % block_size == 0, (
f"Last dimension size must be divisible by block_size (block_size={block_size})"
)
# Flatten all dims except last
N = x.size(-1)
@@ -69,9 +69,7 @@ def _sparse_mla_fwd_kernel(
) # [H, D_V]
q_tail = tl.load(
q_rope_ptr + s_i * H * D_TAIL + h[:, None] * D_TAIL + dt[None, :]
).to(
q_nope_ptr.dtype.element_ty
) # [H, D_TAIL]
).to(q_nope_ptr.dtype.element_ty) # [H, D_TAIL]
m_i = tl.full([H], -float("inf"), tl.float32)
l_i = tl.zeros([H], tl.float32)
@@ -89,9 +87,7 @@ def _sparse_mla_fwd_kernel(
) # [BLOCK_N, D_V] -- reused as V
kv_tail = tl.load(
kbase + (D_V + dt)[None, :], mask=valid[:, None], other=0.0
).to(
q_nope_ptr.dtype.element_ty
) # [BLOCK_N, D_TAIL]
).to(q_nope_ptr.dtype.element_ty) # [BLOCK_N, D_TAIL]
qk = tl.dot(q_main, tl.trans(kv_main)).to(tl.float32)
qk += tl.dot(q_tail, tl.trans(kv_tail)).to(tl.float32)
@@ -209,7 +209,9 @@ def _set_k_and_s_torch(
== num_tokens_to_write_nope
== num_tokens_to_write_rope
== num_tokens_to_write_scale
), f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}"
), (
f"{num_tokens_to_write=} {num_tokens_to_write_nope=} {num_tokens_to_write_rope=} {num_tokens_to_write_scale=}"
)
assert buf.dtype == torch.uint8
assert loc.dtype in [
@@ -110,9 +110,9 @@ def _init_compressed_attn_metadata_triton(
# no cache-write locations. Keep the write buffers unpadded and mask those
# rows in the kernel.
num_write_tokens = raw_out_loc.shape[0]
assert (
num_write_tokens <= bs
), f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows"
assert num_write_tokens <= bs, (
f"raw_out_loc has {num_write_tokens} rows, expected at most {bs} metadata rows"
)
device = seq_lens.device
c4_out_loc = torch.empty(num_write_tokens, dtype=torch.int64, device=device)
@@ -126,12 +126,12 @@ def _init_compressed_attn_metadata_triton(
c128_seq_lens_clamp1 = torch.empty(bs, dtype=torch.int32, device=device)
if compute_page_indices:
assert (
page_table is not None
), "page_table required when compute_page_indices=True"
assert (
page_size >= 128 and page_size % 128 == 0
), "page_size must be a multiple of 128 when compute_page_indices=True"
assert page_table is not None, (
"page_table required when compute_page_indices=True"
)
assert page_size >= 128 and page_size % 128 == 0, (
"page_size must be a multiple of 128 when compute_page_indices=True"
)
max_pages = page_table.shape[1]
c128_page_size = page_size // 128
c128_cur_max_seq_len = c128_page_size * max_pages
@@ -1090,9 +1090,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase):
):
assert blocksparse_tensors is None, "Block sparsity is not supported on SM120"
assert (mBias is not None) == self.has_bias
assert (
mPageTable is None or self.paged_kv
), "SM120 paged KV requires the dedicated DMA-warp specialization"
assert mPageTable is None or self.paged_kv, (
"SM120 paged KV requires the dedicated DMA-warp specialization"
)
self._check_type(
*(
t.element_type if t is not None else None
@@ -1251,7 +1251,9 @@ class FlashAttentionForwardSm120(FlashAttentionForwardBase):
TileScheduler = (
Sm120UniformBatchScheduler
if is_varlen and self.direct_uniform_batch
else SingleTileVarlenScheduler if is_varlen else SingleTileScheduler
else SingleTileVarlenScheduler
if is_varlen
else SingleTileScheduler
)
tile_sched_args = TileSchedulerArguments(
num_block=cute.ceil_div(cute.size(mQ.shape[0]), self.tile_m),
@@ -75,9 +75,9 @@ class Sm120UniformBatchScheduler:
loc=None,
ip=None,
) -> Params:
assert (
scheduling_mode == SchedulingMode.STATIC
), f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}"
assert scheduling_mode == SchedulingMode.STATIC, (
f"SM120 uniform-batch scheduler only supports STATIC, got {scheduling_mode!r}"
)
return Sm120UniformBatchScheduler.Params.create(args, loc=loc, ip=ip)
@staticmethod
@@ -85,7 +85,6 @@ def chunk_gated_delta_rule_fwd(
class ChunkGatedDeltaRuleFunction(torch.autograd.Function):
@staticmethod
@input_guard
@autocast_custom_fwd
@@ -207,12 +206,12 @@ def chunk_gated_delta_rule(
)
"""
assert q.dtype == k.dtype == v.dtype
assert (
q.dtype != torch.float32
), "ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
assert (
len(beta.shape) == 3
), "beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
assert q.dtype != torch.float32, (
"ChunkGatedDeltaRuleFunction does not support float32. Please use bfloat16."
)
assert len(beta.shape) == 3, (
"beta must be of shape [B, T, H] if head_first=False, or [B, H, T] otherwise."
)
if head_first:
raise DeprecationWarning(
@@ -82,9 +82,10 @@ def chunk_gated_delta_rule_fwd_kernel_h_blockdim64(
i_v, i_nh = tl.program_id(0), tl.program_id(1)
i_n, i_h = i_nh // H, i_nh % H
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
NT = tl.cdiv(T, BT)
boh = tl.load(chunk_offsets + i_n).to(tl.int32)
@@ -326,9 +327,9 @@ def chunk_gated_delta_rule_fwd_h(
chunk_indices: Optional[torch.LongTensor] = None,
use_exp2: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
assert not (
use_exp2 and g is not None
), "use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
assert not (use_exp2 and g is not None), (
"use_exp2 covers only the per-channel gk path; scalar g stays natural-exp"
)
B, T, Hg, K, V = *k.shape, u.shape[-1]
H = u.shape[-2]
BT = CHUNK_SIZE
@@ -71,12 +71,14 @@ def chunk_gated_delta_rule_fwd_kkt_solve_kernel(
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -88,12 +88,14 @@ def chunk_kda_fwd_kernel_inter_solve_fused(
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -819,12 +821,14 @@ def chunk_kda_fwd_kernel_intra_sub_chunk(
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -62,9 +62,10 @@ def chunk_kda_fwd_kernel_intra_token_parallel(
left = mid + 1
i_n = left
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
i_t = i_tg - bos
else:
@@ -53,12 +53,14 @@ def chunk_fwd_kernel_o(
if IS_VARLEN:
i_tg = i_t
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
NT = tl.cdiv(T, BT)
else:
@@ -37,12 +37,14 @@ def chunk_local_cumsum_scalar_kernel(
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -97,12 +99,14 @@ def chunk_local_cumsum_vector_kernel(
i_s, i_t, i_bh = tl.program_id(0), tl.program_id(1), tl.program_id(2)
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -169,9 +173,9 @@ def chunk_local_cumsum_scalar(
B, H, T = g.shape
else:
B, T, H = g.shape
assert chunk_size == 2 ** (
chunk_size.bit_length() - 1
), "chunk_size must be a power of 2"
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
"chunk_size must be a power of 2"
)
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
@@ -216,9 +220,9 @@ def chunk_local_cumsum_vector(
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = triton.cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
assert chunk_size == 2 ** (
chunk_size.bit_length() - 1
), "chunk_size must be a power of 2"
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
"chunk_size must be a power of 2"
)
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
@@ -260,9 +264,9 @@ def chunk_local_cumsum(
**kwargs,
) -> torch.Tensor:
if cu_seqlens is not None:
assert (
g.shape[0] == 1
), "Only batch size 1 is supported when cu_seqlens are provided"
assert g.shape[0] == 1, (
"Only batch size 1 is supported when cu_seqlens are provided"
)
if len(g.shape) == 3:
return chunk_local_cumsum_scalar(
g=g,
@@ -390,9 +390,9 @@ class FusedRMSNormGated(nn.Module):
residual_in_fp32: bool = False,
) -> torch.Tensor:
if _use_cpu:
assert (
self.activation == "silu"
), "CPU rmsnorm_gated currently only supports activation silu"
assert self.activation == "silu", (
"CPU rmsnorm_gated currently only supports activation silu"
)
return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
x, self.weight, g, self.eps
)
@@ -43,9 +43,10 @@ def fused_recurrent_gated_delta_rule_fwd_kernel(
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(
cu_seqlens + i_n + 1
).to(tl.int64)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int64),
tl.load(cu_seqlens + i_n + 1).to(tl.int64),
)
all = T
T = eos - bos
else:
@@ -708,7 +709,6 @@ def fused_recurrent_kda_packed_decode(
class FusedRecurrentFunction(torch.autograd.Function):
@staticmethod
@input_guard
def forward(
@@ -907,9 +907,10 @@ def fused_recurrent_gated_delta_rule_update_fwd_kernel(
i_n, i_hv = i_nh // HV, i_nh % HV
i_h = i_hv // (HV // H)
if IS_VARLEN:
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int64), tl.load(
cu_seqlens + i_n + 1
).to(tl.int64)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int64),
tl.load(cu_seqlens + i_n + 1).to(tl.int64),
)
all = T
T = eos - bos
else:
@@ -1144,7 +1145,6 @@ def fused_recurrent_gated_delta_rule_update_fwd(
class FusedRecurrentUpdateFunction(torch.autograd.Function):
@staticmethod
@input_guard
def forward(
@@ -678,9 +678,9 @@ def _launch_gdn_spec(
num_slots, HV, V, K = checkpoint_state.shape
H = k.shape[1]
B = query_start_loc.shape[0] - 1
assert (
max_cache_len & (max_cache_len - 1) == 0
), "circular cache requires power-of-two max_cache_len"
assert max_cache_len & (max_cache_len - 1) == 0, (
"circular cache requires power-of-two max_cache_len"
)
assert d_cache.shape[2] == max_cache_len
BK = triton.next_power_of_2(K)
@@ -1046,18 +1046,18 @@ def kda_gate_chunk_cumsum(
Cumulative-summed gated tensor of shape [B, T, H, K].
"""
if cu_seqlens is not None:
assert (
g.shape[0] == 1
), "Only batch size 1 is supported when cu_seqlens are provided"
assert g.shape[0] == 1, (
"Only batch size 1 is supported when cu_seqlens are provided"
)
assert len(g.shape) == 4
B, T, H, S = g.shape
BT = chunk_size
if chunk_indices is None and cu_seqlens is not None:
chunk_indices = prepare_chunk_indices(cu_seqlens, BT)
NT = cdiv(T, BT) if cu_seqlens is None else len(chunk_indices)
assert chunk_size == 2 ** (
chunk_size.bit_length() - 1
), "chunk_size must be a power of 2"
assert chunk_size == 2 ** (chunk_size.bit_length() - 1), (
"chunk_size must be a power of 2"
)
g_org, g = g, torch.empty_like(g, dtype=output_dtype or g.dtype)
@@ -120,7 +120,6 @@ def l2norm_fwd(
class L2NormFunction(torch.autograd.Function):
@staticmethod
@input_guard
def forward(ctx, x, eps=1e-6, output_dtype=None):
@@ -137,7 +136,6 @@ l2_norm = l2norm
class L2Norm(nn.Module):
def __init__(self, eps: float = 1e-6, output_dtype: Optional[torch.dtype] = None):
super().__init__()
self.eps = eps
@@ -345,7 +345,6 @@ def rms_norm_gated(
class LayerNormFn(torch.autograd.Function):
@staticmethod
def forward(
ctx,
@@ -389,7 +388,6 @@ def layernorm_fn(
class LayerNorm(torch.nn.Module):
def __init__(
self,
hidden_size,
@@ -431,7 +429,6 @@ class LayerNorm(torch.nn.Module):
class RMSNorm(torch.nn.Module):
def __init__(
self,
hidden_size,
@@ -465,7 +462,9 @@ class RMSNorm(torch.nn.Module):
self.norm_before_gate
and self.group_size is None
and self.activation == "swish"
), "CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish"
), (
"CPU rmsnorm_gated currently only supports norm before gate without group size or activation other than swish"
)
return torch.ops.sgl_kernel.fused_rmsnorm_gated_cpu(
x, self.weight, z, self.eps
)
@@ -326,9 +326,9 @@ if torch_release >= (2, 4):
return device_torch_lib.device(index)
else:
assert (
device == "cuda"
), "Only cuda device is supported for PyTorch version < 2.4.0."
assert device == "cuda", (
"Only cuda device is supported for PyTorch version < 2.4.0."
)
autocast_custom_fwd = device_torch_lib.amp.custom_fwd
autocast_custom_bwd = device_torch_lib.amp.custom_bwd
@@ -43,12 +43,14 @@ def recompute_w_u_fwd_kernel(
i_t, i_bh = tl.program_id(0), tl.program_id(1)
i_b, i_h = i_bh // H, i_bh % H
if IS_VARLEN:
i_n, i_t = tl.load(chunk_indices + i_t * 2).to(tl.int32), tl.load(
chunk_indices + i_t * 2 + 1
).to(tl.int32)
bos, eos = tl.load(cu_seqlens + i_n).to(tl.int32), tl.load(
cu_seqlens + i_n + 1
).to(tl.int32)
i_n, i_t = (
tl.load(chunk_indices + i_t * 2).to(tl.int32),
tl.load(chunk_indices + i_t * 2 + 1).to(tl.int32),
)
bos, eos = (
tl.load(cu_seqlens + i_n).to(tl.int32),
tl.load(cu_seqlens + i_n + 1).to(tl.int32),
)
T = eos - bos
else:
bos, eos = i_b * T, i_b * T + T
@@ -77,9 +77,7 @@ def _gather_and_dequant(k_cache, indices, page_size):
raw_pages = k_cache.as_strided(
(num_pages, page_bytes),
(page_bytes, 1),
).view(
torch.uint8
) # (num_pages, page_bytes) uint8
).view(torch.uint8) # (num_pages, page_bytes) uint8
# Note: float8_e4m3fn and uint8 are both 1 byte, view is safe
# Compute byte offsets within each page
@@ -922,8 +922,7 @@ def chunk_kda_fwd(
if needs_eqlen_pad:
if B != 1 and T % BT != 0:
raise NotImplementedError(
f"eqlen with B>1 and T % {BT} != 0 not supported "
f"(got B={B}, T={T})."
f"eqlen with B>1 and T % {BT} != 0 not supported (got B={B}, T={T})."
)
T_padded = ((T + CPB_BT - 1) // CPB_BT) * CPB_BT
# Pre-allocated padded scratch buffers (per (B,T_padded,H,K,dtype) cache
@@ -407,7 +407,6 @@ def k4_persistent_kernel(
# ==== WG1: State readout + decay ====
if warpgroup_idx == STATE_WG:
cId_128 = cute.make_identity_tensor((M6, N6))
tCtState_mn = transform_partitioned_tensor_layout(tCtState)
@@ -698,7 +697,6 @@ def k4_persistent_kernel(
# ==== TMA warp (warp 2) ====
elif warp_idx == TMA_WARP:
cta_layout = cute.make_layout(1)
scheduler = GDNTileScheduler.create(
@@ -995,7 +993,6 @@ def k4_persistent_kernel(
# ==== WG2: W/NV/O readout (SS-mode, no Phase 2) ====
elif warpgroup_idx == READOUT_WG:
tCtW_mn = transform_partitioned_tensor_layout(tCtW)
tCtNV_mn = transform_partitioned_tensor_layout(tCtNV)
tCtO_mn = transform_partitioned_tensor_layout(tCtO)
@@ -107,9 +107,9 @@ def chunk_kda_fwd(
Returns the fla-shaped 12-tuple: (o [B,T,H,128] bf16, final_state
[N,H,128,128] fp32 or None, then Nones, ..., h, initial_state).
"""
assert (
chunk_size == CHUNK
), f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}"
assert chunk_size == CHUNK, (
f"kda_prefill supports chunk_size={CHUNK} only, got {chunk_size}"
)
if cp_context is not None or disable_recompute:
raise NotImplementedError(
"kda_prefill is the inference forward path: cp_context, "
@@ -124,9 +124,9 @@ def chunk_kda_fwd(
if state_v_first and initial_state is not None:
# [V,K]-layout state: pure transpose (K==V==128), exact, ~us/call
initial_state = initial_state.transpose(-1, -2).contiguous()
assert (
q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K
), f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}"
assert q.dim() == 4 and q.shape[-1] == K and v.shape[-1] == K, (
f"expected [B,T,H,{K}] q/k/v, got q={tuple(q.shape)} v={tuple(v.shape)}"
)
B, T, H, _ = q.shape
cu_cpu = None
@@ -145,9 +145,9 @@ def chunk_kda_fwd(
betaf = beta.reshape(Tt, H).contiguous()
if use_gate_in_kernel:
assert (
A_log is not None and dt_bias is not None
), "use_gate_in_kernel=True requires A_log and dt_bias"
assert A_log is not None and dt_bias is not None, (
"use_gate_in_kernel=True requires A_log and dt_bias"
)
assert g.dtype == torch.bfloat16, f"raw gate input must be bf16, got {g.dtype}"
gf = g.reshape(Tt, H, K).contiguous()
sg = lower_bound is not None # fla: lb presence selects safe-gate
@@ -392,7 +392,6 @@ def _fwd_none_diag_kernel(
class _attention(torch.autograd.Function):
@staticmethod
def forward(ctx, q, k, v, s, kv_history):
# Forward pass of the lightning attention algorithm
@@ -158,7 +158,6 @@ def seg_la_kernel(
state = state * block_decay + tl.dot(k, v)
else:
qk = tl.dot(q, k) * softmax_scale
decays = tl.exp(decay_scale * (offs_b[:, None] - offs_b[None, :]))
decays = tl.where(offs_b[None, :] <= offs_b[:, None], decays, 0.0)
@@ -496,9 +496,9 @@ def draft_extend_set_metadata(
row tails keep stale values that attention kernels never read past
cache_seqlens, matching the eager replay path's bounded writes.
"""
assert (
page_size > 0 and (page_size & (page_size - 1)) == 0
), f"page_size must be a power of two, got {page_size}"
assert page_size > 0 and (page_size & (page_size - 1)) == 0, (
f"page_size must be a power of two, got {page_size}"
)
batch_size = cache_seqlens_int32.shape[0]
max_seq_pages = page_table.shape[1]
@@ -588,9 +588,9 @@ def normal_decode_set_metadata(
page_table / swa_page_table row is (re)written; the tail keeps stale values
across CUDA-graph replays, so consumers must bound reads by cache_seqlens.
"""
assert (
page_size > 0 and (page_size & (page_size - 1)) == 0
), f"page_size must be a power of two, got {page_size}"
assert page_size > 0 and (page_size & (page_size - 1)) == 0, (
f"page_size must be a power of two, got {page_size}"
)
batch_size = cache_seqlens_int32.shape[0]
device = seq_lens.device
@@ -113,9 +113,9 @@ def minimax_qknorm_rope_grouped(
"""
groups = [(w, off, cnt) for (w, off, cnt) in groups if cnt > 0]
num_groups = len(groups)
assert (
1 <= num_groups <= _MAX_GROUPS
), f"need 1..{_MAX_GROUPS} groups, got {num_groups}"
assert 1 <= num_groups <= _MAX_GROUPS, (
f"need 1..{_MAX_GROUPS} groups, got {num_groups}"
)
weights: List[torch.Tensor] = [g[0] for g in groups]
offsets: List[int] = [int(g[1]) for g in groups]
@@ -331,9 +331,9 @@ def flash_decode_with_gqa_share_sparse(
max_slots, num_kv_heads, _ = k_cache.shape
assert slot_ids.shape[0] == batch_size and seq_lens.shape[0] == batch_size
assert topk_idx.shape[0] == num_kv_heads
assert (
triton.next_power_of_2(block_size) == block_size
), f"block_size must be a power of 2, but got {block_size}"
assert triton.next_power_of_2(block_size) == block_size, (
f"block_size must be a power of 2, but got {block_size}"
)
# assert slot_ids.max() < max_slots, f"get slot_ids {slot_ids}, but kv_cache shape is {kv_cache.shape}"
max_kv_len = req_to_token.shape[1]
# gqa
@@ -495,9 +495,9 @@ def flash_prefill_with_topk_index(
assert qk_head_dim <= 256 and v_head_dim <= 256, "head_dim must be less than 256"
if sink is not None:
assert sink.shape[0] == num_heads and sink.shape[1] == qk_head_dim
assert (
init_blocks + local_blocks <= topk
), "init_blocks + local_blocks must be less than topk"
assert init_blocks + local_blocks <= topk, (
"init_blocks + local_blocks must be less than topk"
)
if sm_scale is None:
sm_scale = qk_head_dim**-0.5
# q_scale multiplies every Q-side logit (QK dot and sink), so it folds into
@@ -189,21 +189,21 @@ def mla_kv_pack_quantize_fp8(
torch.bfloat16,
torch.float16,
), f"k_nope must be bf16/fp16, got {k_nope.dtype}"
assert (
k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype
), "k_nope, k_pe, v must share dtype"
assert k_pe.dtype == k_nope.dtype and v.dtype == k_nope.dtype, (
"k_nope, k_pe, v must share dtype"
)
assert fp8_dtype in (torch.float8_e4m3fn, torch.float8_e5m2)
s, num_heads, qk_nope = k_nope.shape
qk_rope = k_pe.shape[-1]
v_head = v.shape[-1]
assert (
v.shape[0] == s and v.shape[1] == num_heads
), f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}"
assert (
k_pe.shape[0] == s
), f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}"
assert v.shape[0] == s and v.shape[1] == num_heads, (
f"v shape {tuple(v.shape)} mismatches k_nope {tuple(k_nope.shape)}"
)
assert k_pe.shape[0] == s, (
f"k_pe first dim {k_pe.shape[0]} mismatches k_nope first dim {s}"
)
assert k_nope.stride(-1) == 1, "k_nope must have stride-1 inner dim"
assert v.stride(-1) == 1, "v must have stride-1 inner dim"
assert k_pe.stride(-1) == 1, "k_pe must have stride-1 inner dim"
@@ -326,9 +326,9 @@ def _decode_grouped_att_m_fwd_rope(
is_neox_style=True,
):
if use_rope:
assert (
k_pe_tokens_out is not None
), "We must output the k_pe tokens with rope applied if rope fusion enabled."
assert k_pe_tokens_out is not None, (
"We must output the k_pe tokens with rope applied if rope fusion enabled."
)
BLOCK = 32
@@ -30,9 +30,9 @@ import triton.language as tl
def unpack_aux_tensors(score_mod, aux_tensors):
if score_mod is None:
return None, 0, 0, 0
assert (
aux_tensors is not None and len(aux_tensors) == 1
), "Triton score_mod currently requires exactly one aux tensor"
assert aux_tensors is not None and len(aux_tensors) == 1, (
"Triton score_mod currently requires exactly one aux tensor"
)
aux0 = aux_tensors[0]
assert aux0.dim() == 3 and aux0.stride(2) == 1, (
f"aux_tensors[0] must be 3D with a contiguous last dim, "
@@ -303,7 +303,7 @@ def sparse_mla_q8kv8_prefill_fwd(
)
if indices.ndim != 3:
raise ValueError(
"indices must have shape (s_q, h_kv, topk), " f"got {tuple(indices.shape)}"
f"indices must have shape (s_q, h_kv, topk), got {tuple(indices.shape)}"
)
s_q, h_q, d_qk = q.shape
@@ -362,8 +362,7 @@ def sparse_mla_q8kv8_prefill_fwd(
if indices.shape[:2] != (s_q, h_kv):
raise ValueError(
"indices must have shape "
f"({s_q}, {h_kv}, topk), got {tuple(indices.shape)}"
f"indices must have shape ({s_q}, {h_kv}, topk), got {tuple(indices.shape)}"
)
if indices.dtype != torch.int32:
@@ -385,14 +384,13 @@ def sparse_mla_q8kv8_prefill_fwd(
raise ValueError("topk_length must be a CUDA tensor")
if topk_length.device != device:
raise ValueError(
"topk_length must be on q's device "
f"{device}, got {topk_length.device}"
f"topk_length must be on q's device {device}, got {topk_length.device}"
)
if not topk_length.is_contiguous():
raise ValueError("topk_length must be contiguous")
if torch.any(topk_length < 0).item() or torch.any(topk_length > topk).item():
raise ValueError(
"topk_length values must satisfy " f"0 <= topk_length <= topk ({topk})"
f"topk_length values must satisfy 0 <= topk_length <= topk ({topk})"
)
if d_v != 512:
@@ -121,9 +121,9 @@ _AR_TUNED_TP8 = {
}
_AR_TUNED = {4: _AR_TUNED_TP4, 8: _AR_TUNED_TP8}
_AR_TUNED_TOKENS = sorted(_AR_TUNED_TP4) # same token grid for every table
assert all(
set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values()
), "all tuned tables must share the same token grid"
assert all(set(t) == set(_AR_TUNED_TP4) for t in _AR_TUNED.values()), (
"all tuned tables must share the same token grid"
)
def select_ar_config(num_tokens: int, world_size: int = 4):
@@ -161,8 +161,7 @@ def multigpu_launch(
for N in num_gpus:
if N <= 1 or N > num_devices:
raise ValueError(
f"Invalid number of GPUs requested: {N} "
f"(available: {num_devices})"
f"Invalid number of GPUs requested: {N} (available: {num_devices})"
)
os.environ[env_key] = "1"
os.environ[pid_key] = str(os.getpid())
@@ -1666,14 +1666,14 @@ def cam_scan_bidi_chunkwise(
q, k, v: camera-prepared ``(B, H, D, N)`` fp32; beta: ``(B, H, F, S)`` fp32;
decay: ``(B, H, F)`` fp32. Returns ``(B, H, D, N)`` fp32.
"""
assert (
q.shape == k.shape == v.shape
), f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}"
assert q.shape == k.shape == v.shape, (
f"q/k/v shape mismatch: {q.shape} {k.shape} {v.shape}"
)
assert q.is_contiguous() and k.is_contiguous() and v.is_contiguous()
assert beta.is_contiguous() and decay.is_contiguous()
assert (
q.dtype == torch.float32
), f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}"
assert q.dtype == torch.float32, (
f"cam_scan_bidi_chunkwise requires fp32 q/k/v, got {q.dtype}"
)
B, H, D, N = q.shape
F = beta.shape[2]
@@ -175,9 +175,9 @@ def fused_qk_inv_rms(
qkv: (B, N, 3, H, D) contiguous. Returns (q_inv_rms, k_inv_rms), each (B, N) float32.
"""
assert qkv.is_contiguous(), "qkv must be contiguous (B, N, 3, H, D)"
assert (
qkv.dim() == 5 and qkv.shape[2] == 3
), f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}"
assert qkv.dim() == 5 and qkv.shape[2] == 3, (
f"expected (B, N, 3, H, D), got {tuple(qkv.shape)}"
)
B, N, _, H, D = qkv.shape
C = H * D
q_inv_rms = torch.empty((B, N), dtype=torch.float32, device=qkv.device)
@@ -26,7 +26,7 @@ def _jit_usp_relayout_module(dtype: torch.dtype) -> Module:
cuda_wrappers=[
(
"usp_merge_heads",
"usp_relayout::" f"UspMergeHeadsKernel<{args}>::run",
f"usp_relayout::UspMergeHeadsKernel<{args}>::run",
),
],
)
@@ -413,9 +413,9 @@ def fuse_scale_shift_kernel(
num_warps = 2 if block_n == 64 else 4
grid = (rows, triton.cdiv(C, block_n))
num_frames = scale.shape[1]
assert (
L % num_frames == 0
), "seq_len must be divisible by num_frames for 4D scale/shift"
assert L % num_frames == 0, (
"seq_len must be divisible by num_frames for 4D scale/shift"
)
frame_seqlen = L // num_frames
# Compact scale [B, F, 1, C] -> [B*F, C] (per-frame)
@@ -51,9 +51,9 @@ def _build_fused_norm_module(D: int, is_rms: bool, has_gate: bool, has_weight: b
VEC = _VEC
NUM_WAVES = _NUM_WAVES
BLOCK = NUM_WAVES * WARP_SIZE
assert (
D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0
), f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, (
f"FlyDSL fused_residual_norm requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
)
NUM_ITERS = D // (BLOCK * VEC)
@flyc.kernel(known_block_size=[BLOCK, 1, 1])
@@ -543,9 +543,9 @@ def _build_norm_scale_shift_module(D: int, is_rms: bool, has_weight: bool):
VEC = _VEC
NUM_WAVES = _NUM_WAVES
BLOCK = NUM_WAVES * WARP_SIZE
assert (
D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0
), f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
assert D % FLYDSL_NORM_MIN_ALIGNED_DIM == 0, (
f"FlyDSL norm_scale_shift requires D % {FLYDSL_NORM_MIN_ALIGNED_DIM} == 0, got D={D}"
)
NUM_ITERS = D // (BLOCK * VEC)
@flyc.kernel(known_block_size=[BLOCK, 1, 1])
@@ -103,9 +103,9 @@ fused_dual_residual_rmsnorm_kernel_autotune = rmsnorm_autotune(
def fused_dual_residual_rmsnorm(x, residual, weight1, weight2, eps, autotune=False):
assert len(x.shape) == 2
assert (
x.shape == residual.shape and x.dtype == residual.dtype
), f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}"
assert x.shape == residual.shape and x.dtype == residual.dtype, (
f"{x.shape=} {residual.shape=} {x.dtype=} {residual.dtype=}"
)
output, mid = torch.empty_like(x), torch.empty_like(x)
bs, hidden_dim = x.shape
if autotune:
@@ -434,9 +434,9 @@ def fused_sigmoid_mul(
gate_stride_head = gate.stride(1)
else:
# Flat path: both tensors have the same shape
assert (
attn_output.shape == gate.shape
), "attn_output and gate must have the same shape"
assert attn_output.shape == gate.shape, (
"attn_output and gate must have the same shape"
)
hidden_dim = attn_output.shape[-1]
num_tokens = attn_output.numel() // hidden_dim
head_dim = hidden_dim
@@ -338,9 +338,9 @@ def _dsv3_fused_a_gemm_run(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Te
assert mat_a.dtype == torch.bfloat16 and mat_b.dtype == torch.bfloat16
assert K % 1024 == 0, f"K must be a multiple of 1024, got {K}"
assert N % TILE_M == 0, f"N must be a multiple of {TILE_M}, got {N}"
assert (
tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1
), "mat_b must be [K, N] column-major"
assert tuple(mat_b.shape) == (K, N) and mat_b.stride(0) == 1, (
"mat_b must be [K, N] column-major"
)
assert 1 <= M <= 16, "num_tokens must be in [1, 16]"
assert mat_a.stride(1) == 1, "mat_a must be row-major [M, K]"
@@ -229,7 +229,9 @@ def default_tactic(m: int, n: int, k: int) -> SplitKTactic:
ab_stages=(
_MIN_AB_STAGES
if k <= 2 * _CTA_K and m > 8
else min(max_stages, 6) if k <= 4 * _CTA_K else max_stages
else min(max_stages, 6)
if k <= 4 * _CTA_K
else max_stages
),
)
validate_tactic(tactic, m, n, k)
@@ -114,9 +114,9 @@ def apply_token_bitmask_inplace_triton(
indices = torch.tensor(indices, dtype=torch.int32, device=logits.device)
num_rows = indices.shape[0]
else:
assert (
logits_shape[0] == bitmask_shape[0]
), f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}"
assert logits_shape[0] == bitmask_shape[0], (
f"batch size mismatch: logits {logits_shape[0]} vs bitmask {bitmask_shape[0]}"
)
num_rows = logits_shape[0]
if NUM_SMS > 0:
@@ -51,10 +51,7 @@ def _device_name(device: torch.device) -> str:
def _table(world_size: int, hidden_size: int, device: torch.device) -> Optional[dict]:
path = os.path.join(
_CONFIG_DIR,
(
f"world={world_size},H={hidden_size},"
f"device_name={_device_name(device)}.json"
),
(f"world={world_size},H={hidden_size},device_name={_device_name(device)}.json"),
)
if path not in _TABLES:
if os.path.exists(path):
+24 -24
View File
@@ -100,18 +100,18 @@ def concat_and_cast_mha_k_triton(
k_rope: torch.Tensor,
):
# The source data type will be implicitly converted to the target data type.
assert (
len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3
), f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
assert (
k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0]
), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
assert (
k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1]
), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
assert (
k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1]
), f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
assert len(k.shape) == 3 and len(k_nope.shape) == 3 and len(k_rope.shape) == 3, (
f"shape should be 3d, but got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
)
assert k.shape[0] == k_nope.shape[0] and k.shape[0] == k_rope.shape[0], (
f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
)
assert k.shape[1] == k_nope.shape[1] and 1 == k_rope.shape[1], (
f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
)
assert k.shape[-1] == k_nope.shape[-1] + k_rope.shape[-1], (
f"invalid shape, got {k.shape=}, {k_nope.shape=}, {k_rope.shape=}"
)
nope_dim = k_nope.shape[-1]
rope_dim = k_rope.shape[-1]
@@ -638,9 +638,9 @@ def absorbed_bmm_concat_cast_q_fp8(
assert q_fp8_pad.shape[0] >= num_tokens and q_fp8_pad.shape[1] >= num_heads
assert q_fp8_pad.shape[2] == n_dim + rope_dim
# tl.arange / tl.dot constraints
assert (
k_dim % 16 == 0 and 16 <= k_dim <= 256
), "K must be a multiple of 16 in [16, 256]"
assert k_dim % 16 == 0 and 16 <= k_dim <= 256, (
"K must be a multiple of 16 in [16, 256]"
)
assert (rope_dim & (rope_dim - 1)) == 0, "ROPE must be a power of two"
assert n_dim % block_n == 0, "N must be a multiple of block_n"
assert q_nope.stride(2) == 1 and q_rope.stride(2) == 1
@@ -680,22 +680,22 @@ def absorbed_bmm_concat_cast_q_fp8(
# Largest power-of-2 divisor of K, capped at 128 (K % 16 == 0
# makes this >= 16), unless the caller pinned block_k.
blk_k = block_k or min(k_dim & -k_dim, 128)
assert (
k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16
), "loop needs BLOCK_K a power-of-2 divisor of K >= 16"
assert k_dim % blk_k == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, (
"loop needs BLOCK_K a power-of-2 divisor of K >= 16"
)
k_mode = 1
elif v == "two_dot":
blk_k = 1 << (k_dim.bit_length() - 1) # largest power of 2 < K
k1 = k_dim - blk_k
assert (
k1 & (k1 - 1) == 0 and k1 >= 16
), "two_dot needs K = pow2 + pow2 with both halves >= 16"
assert k1 & (k1 - 1) == 0 and k1 >= 16, (
"two_dot needs K = pow2 + pow2 with both halves >= 16"
)
k_mode = 2
elif v == "three_dot":
blk_k = k_dim // 3
assert (
k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16
), "three_dot needs K = 3 * pow2 with pow2 >= 16"
assert k_dim % 3 == 0 and blk_k & (blk_k - 1) == 0 and blk_k >= 16, (
"three_dot needs K = 3 * pow2 with pow2 >= 16"
)
k_mode = 3
elif v == "pad":
blk_k = 1 << k_dim.bit_length() # next power of 2 above K
@@ -292,9 +292,9 @@ def _load_cache_to_device_buffer_mla(
miss_count: torch.Tensor | None,
skip_io: bool,
) -> None:
assert (
hot_buffer_size >= num_top_k
), f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})"
assert hot_buffer_size >= num_top_k, (
f"hot_buffer_size ({hot_buffer_size}) must be >= num_top_k ({num_top_k})"
)
record_miss_plan = miss_src is not None
module = _jit_sparse_module(
@@ -192,9 +192,9 @@ def build_kv_read_table(
region's live prefix is written -- never rebound, never tail-cleared.
"""
bs = int(req_pool_indices.numel())
assert (
out.dtype == torch.int32
), f"build_kv_read_table: out must be int32, got {out.dtype}"
assert out.dtype == torch.int32, (
f"build_kv_read_table: out must be int32, got {out.dtype}"
)
assert out.dim() == 2 and out.shape[0] >= bs and out.shape[1] >= max_pages, (
f"build_kv_read_table: out {tuple(out.shape)} cannot hold "
f"(bs={bs}, max_pages={max_pages})"
@@ -262,7 +262,7 @@ def build_kv_read_table_packed(
"""
bs = int(req_pool_indices.numel())
assert out.dtype in (torch.int32, torch.int64), (
f"build_kv_read_table_packed: out must be int32 or int64, got " f"{out.dtype}"
f"build_kv_read_table_packed: out must be int32 or int64, got {out.dtype}"
)
assert out.dim() == 1 and out.numel() >= max_tokens, (
f"build_kv_read_table_packed: out {tuple(out.shape)} cannot hold "
+27 -27
View File
@@ -613,37 +613,37 @@ def fused_qk_rope_reshape_and_cache(
value_shuffle_layout = False
(t_slot,) = slot_mapping.shape
assert (
t == tk == tv and t_slot <= tk
), f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
assert (
block_size == block_size_v
), f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
assert (
kh == vh == kh_cache == vh_cache
), "KV head should be identical for k, v, key_cache, and value_cache"
assert (
t_cache == t_cache_v
), "Number of tokens should be identical for key_cache, and value_cache"
assert t == tk == tv and t_slot <= tk, (
f"Number of tokens should be identical for q, kand v. The number of tokens of slot_mapping should no more than that of q, k and v, {t=} {tk=} {tv=} {t_slot=}"
)
assert block_size == block_size_v, (
f"block size should be identical for key_cache, and value_cache {block_size} {block_size_v}"
)
assert kh == vh == kh_cache == vh_cache, (
"KV head should be identical for k, v, key_cache, and value_cache"
)
assert t_cache == t_cache_v, (
"Number of tokens should be identical for key_cache, and value_cache"
)
if flash_layout:
assert (
d == dk == dv == dk_cache == dv_cache
), "D dimension should be identical for q, k, and v"
assert d == dk == dv == dk_cache == dv_cache, (
"D dimension should be identical for q, k, and v"
)
else:
assert (
d == dk == dv == dkx_cache * x_cache == dv_cache
), "D dimension should be identical for q, k, and v"
assert d == dk == dv == dkx_cache * x_cache == dv_cache, (
"D dimension should be identical for q, k, and v"
)
assert x_cache == triton.next_power_of_2(x_cache), "x_size should be power of 2"
assert d == triton.next_power_of_2(d), "D dimension should be power of 2"
assert block_size == triton.next_power_of_2(
block_size
), "block_size should be power of 2"
assert block_size == triton.next_power_of_2(block_size), (
"block_size should be power of 2"
)
assert qh % kh == 0, "Q heads must be multiple of H heads"
d_freq = cos_sin.shape[-1] // 2
assert (d_freq == d // 2) or (
d_freq == d
), "cos/sin last dim should be the same or half of the qk last dim"
assert (d_freq == d // 2) or (d_freq == d), (
"cos/sin last dim should be the same or half of the qk last dim"
)
reuse_freqs_front_part = d_freq == d // 2
if q_out is None:
@@ -654,9 +654,9 @@ def fused_qk_rope_reshape_and_cache(
if zeros_out is not None:
tz, qhz, dz = zeros_out.shape
assert (
t == tz and qh == qhz and d == dz
), f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}"
assert t == tz and qh == qhz and d == dz, (
f"q and zeros shape mismatch {q.shape=} {zeros_out.shape=}"
)
output_zeros = True
elif output_zeros:
zeros_out = torch.empty((t, qh, d), dtype=q.dtype, device=q.device)
@@ -113,12 +113,12 @@ def build_trtllm_mha_page_table(
``full_to_swa`` is provided, which then also requires ``swa_page_table``.
"""
has_swa = full_to_swa is not None
assert has_swa == (
swa_page_table is not None
), "full_to_swa and swa_page_table must be provided together"
assert (
_MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0
), f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
assert has_swa == (swa_page_table is not None), (
"full_to_swa and swa_page_table must be provided together"
)
assert _MHA_KV_INDEX_BLOCK_TOKENS % page_size == 0, (
f"page_size={page_size} must divide _MHA_KV_INDEX_BLOCK_TOKENS={_MHA_KV_INDEX_BLOCK_TOKENS}"
)
bs, num_pages = page_table.shape
full_to_swa_numel = full_to_swa.numel() if has_swa else 0
create_trtllm_mha_kv_indices_triton[
@@ -385,8 +385,7 @@ class GemmaFusedAddRMSNormOp(BaseFusedOp):
"(rocm-triton, sglang.kernels.jit)."
),
KernelBackend.TORCH: (
"Gemma-style fused residual-add + RMS normalization "
"(pure-torch reference)."
"Gemma-style fused residual-add + RMS normalization (pure-torch reference)."
),
}
+6 -6
View File
@@ -1103,9 +1103,9 @@ def mhc_pre(
gemm_out_sqrsum = torch.empty(
n_splits, num_tokens, dtype=torch.float32, device=residual.device
)
assert (
n_splits == 1
), "The simple TileLang version gemm_sqrsum doesn't support split-k"
assert n_splits == 1, (
"The simple TileLang version gemm_sqrsum doesn't support split-k"
)
_mhc_pre_gemm_sqrsum_dispatch()(
residual_flat.view(num_tokens, hc_mult * hidden_size),
fn_flat,
@@ -1119,9 +1119,9 @@ def mhc_pre(
if norm_weight is not None:
assert norm_eps is not None, "norm_eps required when norm_weight is provided"
assert norm_weight.shape == (
hidden_size,
), f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)"
assert norm_weight.shape == (hidden_size,), (
f"norm_weight shape {tuple(norm_weight.shape)} != (hidden_size={hidden_size},)"
)
norm_weight_bf = (
norm_weight.bfloat16()
if norm_weight.dtype != torch.bfloat16
@@ -250,9 +250,9 @@ def dispatch_probability(
if random_vals is None:
random_vals = torch.rand(n, dtype=torch.float32, device=topk_ids.device)
else:
assert random_vals.shape == (
n,
), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
assert random_vals.shape == (n,), (
f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
)
module = _dispatch_module(max_copies, DISPATCH_BLOCK_DIM)
module.dispatch_probability(out, flat_ids, log2phy_prob, map32, random_vals)
return out.view(original_shape).to(topk_ids.dtype)
@@ -299,9 +299,9 @@ def dispatch_probability_torch_reference(
n = flat_ids.shape[0]
num_logical, max_copies = log2phy_prob.shape
assert log2phy_map.shape == (num_logical, max_copies)
assert random_vals.shape == (
n,
), f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
assert random_vals.shape == (n,), (
f"random_vals must be shape ({n},), got {tuple(random_vals.shape)}"
)
# Gather per-row probabilities and physical maps.
probs = log2phy_prob[flat_ids] # (N, max_copies), float32
@@ -111,8 +111,8 @@ def assert_fits(nc: int, nv: int, gpu: str = "h100") -> None:
cap = gpu_budget_bytes(gpu)
if used > cap:
raise ValueError(
f"fused IPM kernel needs {used/1024:.1f} KiB of shared memory for "
f"NC={nc}, NV={nv}, but {gpu} allows {cap/1024:.1f} KiB/block. "
f"fused IPM kernel needs {used / 1024:.1f} KiB of shared memory for "
f"NC={nc}, NV={nv}, but {gpu} allows {cap / 1024:.1f} KiB/block. "
f"Either reduce problem size or switch to a tiled design."
)
@@ -145,8 +145,8 @@ def report(nc: int, nv: int, gpu: str = "h100") -> str:
status = "FITS" if bd.total_bytes <= cap else "OVER BUDGET"
return (
f"[shmem] NC={nc} NV={nv} gpu={gpu} | "
f"A={bd.a_bytes/1024:.1f}K "
f"ata={bd.ata_bytes/1024:.1f}K "
f"rest={(bd.c_bytes+bd.x_bytes+bd.rhs_bytes+bd.d_bytes)/1024:.1f}K | "
f"total={bd.total_bytes/1024:.1f}K / {cap/1024:.1f}K {status}"
f"A={bd.a_bytes / 1024:.1f}K "
f"ata={bd.ata_bytes / 1024:.1f}K "
f"rest={(bd.c_bytes + bd.x_bytes + bd.rhs_bytes + bd.d_bytes) / 1024:.1f}K | "
f"total={bd.total_bytes / 1024:.1f}K / {cap / 1024:.1f}K {status}"
)
@@ -332,7 +332,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching
matrix_w = w_col0
matrix_x = col0
for j in tl.static_range(KERNEL_WIDTH):
if KERNEL_WIDTH == 2:
if j == 1: # KERNEL_WIDTH-1:
matrix_w = w_col1
@@ -502,9 +501,9 @@ def causal_conv1d_fn(
assert padded_batch == cache_indices.size(0)
if has_initial_state is not None:
assert has_initial_state.size() == (padded_batch,)
assert (
conv_states is not None
), "ERROR: `has_initial_state` is used, which needs also `conv_states`"
assert conv_states is not None, (
"ERROR: `has_initial_state` is used, which needs also `conv_states`"
)
assert weight.stride(1) == 1
assert (dim, width) == weight.shape
assert is_channel_last, "Need to run in channel-last layout"
@@ -1053,9 +1052,9 @@ def causal_conv1d_update(
if validate_data:
assert dim == weight.size(0)
assert (
conv_state.stride(-2) == 1
), f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})"
assert conv_state.stride(-2) == 1, (
f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})"
)
assert state_len >= width - 1
# when above happens, we don't shift-left to keep any records in conv_state
assert dim == conv_state.size(1)

Some files were not shown because too many files have changed in this diff Show More