From fd54735de6f8fee8db93ad5ea7d23f4084dc765c Mon Sep 17 00:00:00 2001 From: maxiaosong1124 Date: Sat, 12 Sep 2026 12:49:23 +0800 Subject: [PATCH] style: format Python sources and benchmark results Signed-off-by: maxiaosong1124 --- benchmarks/benchmark_rocm_det_gemm_leaf.py | 16 +- .../benchmark_rocm_det_gemm_root_output.py | 4 +- benchmarks/benchmark_rocm_ffn.py | 258 ++++------ benchmarks/profile_rocm_ffn.py | 41 +- .../results/ws2_rocm_mi300x/results.json | 2 +- .../analyze_performance.py | 118 ++--- .../collect_results.py | 31 +- .../vime_qwen3_8b_tp4_cp2_200/plot_results.py | 24 +- .../prepare_dapo_data.py | 10 +- examples/vime_qwen3_8b_tp4_cp2_200/run.py | 60 +-- .../validate_artifacts.py | 27 +- .../vime_qwen3_8b_tp4_cp2_200/validate_run.py | 95 ++-- examples/vime_rocm_attention_ablation/run.py | 35 +- .../run_full_pp_200.py | 3 +- .../tis_metrics.py | 3 +- .../validate_artifacts.py | 68 +-- rl_engine/distributed/collectives.py | 12 +- rl_engine/distributed/rocm_collectives.py | 12 +- rl_engine/integrations/framework_operators.py | 203 ++++---- rl_engine/integrations/linear_logp.py | 3 +- rl_engine/integrations/megatron_runtime.py | 15 +- rl_engine/integrations/rocm_ablation.py | 24 +- rl_engine/integrations/runtime.py | 11 +- .../integrations/vime/linear_logp_provider.py | 3 +- rl_engine/integrations/vllm_runtime.py | 41 +- rl_engine/kernels/ops/pytorch/ffn/ffn.py | 24 +- .../ops/rocm/attention/paged_gather.py | 4 +- .../ops/rocm/attention/strict_runtime.py | 23 +- .../ops/rocm/loss/vocab_parallel_logp.py | 51 +- rl_engine/kernels/ops/rocm/matmul/det_gemm.py | 4 +- .../kernels/ops/rocm/rotary_embedding/rope.py | 4 +- .../kernels/ops/triton/activation/swiglu.py | 4 +- rl_engine/kernels/ops/triton/ffn/ffn.py | 48 +- .../kernels/ops/triton/matmul/det_gemm.py | 94 ++-- tests/distributed/test_qwen_ffn_topology.py | 13 +- tests/test_det_gemm.py | 26 +- tests/test_framework_runtime_adapters.py | 471 +++++++++--------- tests/test_qwen_ffn.py | 4 +- tests/test_rocm_e2e_ablation.py | 8 +- tests/test_rocm_packed_ffn.py | 8 +- tests/test_vime_linear_logp_provider.py | 12 +- 41 files changed, 713 insertions(+), 1204 deletions(-) diff --git a/benchmarks/benchmark_rocm_det_gemm_leaf.py b/benchmarks/benchmark_rocm_det_gemm_leaf.py index eb00a05b..76f22134 100644 --- a/benchmarks/benchmark_rocm_det_gemm_leaf.py +++ b/benchmarks/benchmark_rocm_det_gemm_leaf.py @@ -56,8 +56,7 @@ class LeafConfig: @property def slug(self) -> str: return ( - f"{self.block_m}x{self.block_n}x{self.num_warps}x" - f"{self.order}xw{self.waves_per_eu}" + f"{self.block_m}x{self.block_n}x{self.num_warps}x" f"{self.order}xw{self.waves_per_eu}" ) @@ -153,8 +152,7 @@ def _parse_config(value: str) -> LeafConfig: waves_per_eu = int(parts[4]) if len(parts) == 5 else 0 except (TypeError, ValueError) as error: raise argparse.ArgumentTypeError( - "config must be BLOCK_MxBLOCK_NxNUM_WARPS[xORDER[xWAVES_PER_EU]], " - f"got {value!r}" + "config must be BLOCK_MxBLOCK_NxNUM_WARPS[xORDER[xWAVES_PER_EU]], " f"got {value!r}" ) from error if ( block_m <= 0 @@ -384,17 +382,14 @@ def _run_case( samples: int, ) -> dict[str, object]: print( - f"{case.name}: A=({case.m_size}, {case.k_size}), " - f"B=({case.k_size}, {case.n_size})", + f"{case.name}: A=({case.m_size}, {case.k_size}), " f"B=({case.k_size}, {case.n_size})", flush=True, ) a, b = _inputs(case, device) plan = _device_tree_plan(case.k_size, device) workspace_shape = (plan.host.node_count, case.m_size, case.n_size) output_shape = ( - (case.n_size, case.m_size) - if case.transpose_output - else (case.m_size, case.n_size) + (case.n_size, case.m_size) if case.transpose_output else (case.m_size, case.n_size) ) reference_workspace = torch.empty(workspace_shape, dtype=torch.bfloat16, device=device) reference_output = torch.empty(output_shape, dtype=torch.bfloat16, device=device) @@ -413,8 +408,7 @@ def _run_case( reference_fingerprints = { "leaf_workspace_sha256_raw_bytes": _tensor_sha256(reference_leaves), "root_sha256_raw_bytes": _tensor_sha256(reference_output), - "leaf_workspace_nbytes": reference_leaves.numel() - * reference_leaves.element_size(), + "leaf_workspace_nbytes": reference_leaves.numel() * reference_leaves.element_size(), "root_nbytes": reference_output.numel() * reference_output.element_size(), } diff --git a/benchmarks/benchmark_rocm_det_gemm_root_output.py b/benchmarks/benchmark_rocm_det_gemm_root_output.py index 44ada2d6..ed58e152 100644 --- a/benchmarks/benchmark_rocm_det_gemm_root_output.py +++ b/benchmarks/benchmark_rocm_det_gemm_root_output.py @@ -131,9 +131,7 @@ def _launch(state: CaseState, *, direct_root_output: bool) -> None: ) if not direct_root_output: - det_gemm._copy_tree_root_kernel[ - (triton.cdiv(state.output.numel(), reduction_block),) - ]( + det_gemm._copy_tree_root_kernel[(triton.cdiv(state.output.numel(), reduction_block),)]( state.workspace, state.output, plan.host.root, diff --git a/benchmarks/benchmark_rocm_ffn.py b/benchmarks/benchmark_rocm_ffn.py index 8b0b342b..ecb8c983 100644 --- a/benchmarks/benchmark_rocm_ffn.py +++ b/benchmarks/benchmark_rocm_ffn.py @@ -117,9 +117,7 @@ def _accuracy(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: "max_abs": float(difference.abs().max().item()), "mean_abs": float(difference.abs().mean().item()), "relative_l2": _relative_l2(actual, expected), - "exact_fraction": float( - (actual.detach() == expected.detach()).float().mean().item() - ), + "exact_fraction": float((actual.detach() == expected.detach()).float().mean().item()), } @@ -317,9 +315,7 @@ def _official_distributed_ffn( if sequence_parallel else _NativeCopyToTensorParallel.apply(hidden, tp_group) ) - activated = F.silu(F.linear(full_hidden, gate_weight)) * F.linear( - full_hidden, up_weight - ) + activated = F.silu(F.linear(full_hidden, gate_weight)) * F.linear(full_hidden, up_weight) partial_output = F.linear(activated, down_weight) if sequence_parallel: return _NativeReduceScatter.apply(partial_output, tp_group) @@ -369,9 +365,7 @@ def _single_gpu_benchmarks( official = _official_qwen3_mlp(gate_weight, up_weight, down_weight) for index, tokens in enumerate((1, 8, 32)): hidden = _randn((tokens, 4096), seed=3010 + index * 2, device=device) - grad_output = _randn( - (tokens, 4096), seed=3011 + index * 2, device=device - ) + grad_output = _randn((tokens, 4096), seed=3011 + index * 2, device=device) weights = (gate_weight, up_weight, down_weight) official_timing = _summary_ms( _gpu_event_samples( @@ -410,12 +404,9 @@ def _single_gpu_benchmarks( official_hidden = hidden.detach().clone().requires_grad_(True) triton_inputs = [ - value.detach().clone().requires_grad_(True) - for value in (hidden, *weights) + value.detach().clone().requires_grad_(True) for value in (hidden, *weights) ] - triton_forward_weights = pack_qwen3_ffn_forward_weights( - *triton_inputs[1:] - ) + triton_forward_weights = pack_qwen3_ffn_forward_weights(*triton_inputs[1:]) def triton_training_step() -> torch.Tensor: return _training_step( @@ -460,8 +451,7 @@ def official_training_step() -> torch.Tensor: "official_tp1": official_train_timing, "triton": triton_train_timing, "latency_ratio_vs_official_tp1": ( - triton_train_timing["median_ms"] - / official_train_timing["median_ms"] + triton_train_timing["median_ms"] / official_train_timing["median_ms"] ), } ) @@ -478,26 +468,16 @@ def official_training_step() -> torch.Tensor: del official, forward_weights, gate_weight, up_weight, down_weight torch.cuda.empty_cache() tokens = 8 - fp32_hidden = _randn( - (tokens, 4096), seed=3100, device=device, dtype=torch.float32 - ) - fp32_gate = _randn( - (12288, 4096), seed=3101, device=device, dtype=torch.float32 - ) - fp32_up = _randn( - (12288, 4096), seed=3102, device=device, dtype=torch.float32 - ) - fp32_down = _randn( - (4096, 12288), seed=3103, device=device, dtype=torch.float32 - ) + fp32_hidden = _randn((tokens, 4096), seed=3100, device=device, dtype=torch.float32) + fp32_gate = _randn((12288, 4096), seed=3101, device=device, dtype=torch.float32) + fp32_up = _randn((12288, 4096), seed=3102, device=device, dtype=torch.float32) + fp32_down = _randn((4096, 12288), seed=3103, device=device, dtype=torch.float32) official_fp32 = _official_qwen3_mlp(fp32_gate, fp32_up, fp32_down) with torch.no_grad(): fp32_output = official_fp32(fp32_hidden) del official_fp32 torch.cuda.empty_cache() - official_fp16 = _official_qwen3_mlp( - fp32_gate.half(), fp32_up.half(), fp32_down.half() - ) + official_fp16 = _official_qwen3_mlp(fp32_gate.half(), fp32_up.half(), fp32_down.half()) with torch.no_grad(): fp16_output = official_fp16(fp32_hidden.half()) results["dtype_accuracy"].append( @@ -514,8 +494,6 @@ def official_training_step() -> torch.Tensor: return results - - def _mesh_groups( world_size: int, tp_size: int, @@ -583,9 +561,7 @@ def _distributed_wall_samples( return timings -def _slowest_rank_summary( - local_timings: list[float], group: Any -) -> dict[str, float]: +def _slowest_rank_summary(local_timings: list[float], group: Any) -> dict[str, float]: world_size = dist.get_world_size(group=group) gathered: list[list[float] | None] = [None] * world_size dist.all_gather_object(gathered, local_timings, group=group) @@ -596,8 +572,6 @@ def _slowest_rank_summary( return _summary_ms(slowest) - - def _distributed_ffn_benchmark( rank: int, world_size: int, @@ -612,37 +586,37 @@ def _distributed_ffn_benchmark( hidden_size = 4096 intermediate_size = 12288 hidden_full = _randn((token_count, hidden_size), seed=5000, device=device) - gate_full = _randn( - (intermediate_size, hidden_size), seed=5001, device=device - ) + gate_full = _randn((intermediate_size, hidden_size), seed=5001, device=device) up_full = _randn((intermediate_size, hidden_size), seed=5002, device=device) - down_full = _randn( - (hidden_size, intermediate_size), seed=5003, device=device - ) - grad_output_full = _randn( - (token_count, hidden_size), seed=5004, device=device - ) + down_full = _randn((hidden_size, intermediate_size), seed=5003, device=device) + grad_output_full = _randn((token_count, hidden_size), seed=5004, device=device) # Exactness reference: the same deterministic Triton implementation at TP=1. full_values = (hidden_full, gate_full, up_full, down_full) tp1_forward_weights = pack_qwen3_ffn_forward_weights(*full_values[1:]) with torch.no_grad(): - tp1_forward = qwen3_ffn( - *full_values, - forward_weights=tp1_forward_weights, - ).detach().clone() - tp1_inputs = [ - value.detach().clone().requires_grad_(True) for value in full_values - ] + tp1_forward = ( + qwen3_ffn( + *full_values, + forward_weights=tp1_forward_weights, + ) + .detach() + .clone() + ) + tp1_inputs = [value.detach().clone().requires_grad_(True) for value in full_values] tp1_training_weights = pack_qwen3_ffn_forward_weights(*tp1_inputs[1:]) - tp1_train = _training_step( - lambda *values: qwen3_ffn( - *values, - forward_weights=tp1_training_weights, - ), - tp1_inputs, - grad_output_full, - ).detach().clone() + tp1_train = ( + _training_step( + lambda *values: qwen3_ffn( + *values, + forward_weights=tp1_training_weights, + ), + tp1_inputs, + grad_output_full, + ) + .detach() + .clone() + ) tp1_grads = [value.grad.detach().clone() for value in tp1_inputs] del tp1_inputs, tp1_forward_weights, tp1_training_weights torch.cuda.empty_cache() @@ -693,9 +667,7 @@ def official_distributed_forward(): ), dist.group.WORLD, ) - official_inputs = [ - value.detach().clone().requires_grad_(True) for value in shard - ] + official_inputs = [value.detach().clone().requires_grad_(True) for value in shard] official_train_summary = _slowest_rank_summary( _distributed_wall_samples( lambda: _official_distributed_training_step( @@ -734,12 +706,8 @@ def triton_forward(): dist.group.WORLD, ) - triton_inputs = [ - value.detach().clone().requires_grad_(True) for value in shard - ] - repeat_inputs = [ - value.detach().clone().requires_grad_(True) for value in shard - ] + triton_inputs = [value.detach().clone().requires_grad_(True) for value in shard] + repeat_inputs = [value.detach().clone().requires_grad_(True) for value in shard] triton_forward_weights = pack_qwen3_ffn_forward_weights(*triton_inputs[1:]) repeat_forward_weights = pack_qwen3_ffn_forward_weights(*repeat_inputs[1:]) @@ -756,15 +724,23 @@ def triton_training_step(inputs, packed_weights): output.backward(local_grad_output) return output - triton_train = triton_training_step( - triton_inputs, - triton_forward_weights, - ).detach().clone() + triton_train = ( + triton_training_step( + triton_inputs, + triton_forward_weights, + ) + .detach() + .clone() + ) triton_grads = [value.grad.detach().clone() for value in triton_inputs] - repeat_train = triton_training_step( - repeat_inputs, - repeat_forward_weights, - ).detach().clone() + repeat_train = ( + triton_training_step( + repeat_inputs, + repeat_forward_weights, + ) + .detach() + .clone() + ) repeat_grads = [value.grad.detach().clone() for value in repeat_inputs] triton_train_summary = _slowest_rank_summary( _distributed_wall_samples( @@ -790,14 +766,10 @@ def triton_training_step(inputs, packed_weights): local_exactness = { "tp1_forward_output": _mismatches(triton_output, expected_forward), "tp1_training_output": _mismatches(triton_train, expected_train), - "tp1_hidden_gradient": _mismatches( - triton_grads[0], expected_grads[0] - ), + "tp1_hidden_gradient": _mismatches(triton_grads[0], expected_grads[0]), "tp1_weight_gradient": sum( _mismatches(actual, expected) - for actual, expected in zip( - triton_grads[1:], expected_grads[1:], strict=True - ) + for actual, expected in zip(triton_grads[1:], expected_grads[1:], strict=True) ), "repeat_forward": _mismatches(triton_output, triton_repeat), "train_infer_mismatch_count": _mismatches(triton_output, triton_train), @@ -834,13 +806,9 @@ def triton_training_step(inputs, packed_weights): / official_forward_summary["median_ms"] ), "tp1_mismatch": { - "forward_output": sum( - value["tp1_forward_output"] for value in valid - ), + "forward_output": sum(value["tp1_forward_output"] for value in valid), }, - "repeat_mismatch_count": sum( - value["repeat_forward"] for value in valid - ), + "repeat_mismatch_count": sum(value["repeat_forward"] for value in valid), "train_infer_mismatch_count": sum( value["train_infer_mismatch_count"] for value in valid ), @@ -851,23 +819,14 @@ def triton_training_step(inputs, packed_weights): "official_distributed": official_train_summary, "triton": triton_train_summary, "latency_ratio_triton_vs_official_distributed": ( - triton_train_summary["median_ms"] - / official_train_summary["median_ms"] + triton_train_summary["median_ms"] / official_train_summary["median_ms"] ), "tp1_mismatch": { - "training_output": sum( - value["tp1_training_output"] for value in valid - ), - "hidden_gradient": sum( - value["tp1_hidden_gradient"] for value in valid - ), - "weight_gradient": sum( - value["tp1_weight_gradient"] for value in valid - ), + "training_output": sum(value["tp1_training_output"] for value in valid), + "hidden_gradient": sum(value["tp1_hidden_gradient"] for value in valid), + "weight_gradient": sum(value["tp1_weight_gradient"] for value in valid), }, - "repeat_mismatch_count": sum( - value["repeat_training"] for value in valid - ), + "repeat_mismatch_count": sum(value["repeat_training"] for value in valid), "train_infer_mismatch_count": sum( value["train_infer_mismatch_count"] for value in valid ), @@ -997,9 +956,7 @@ def _run_distributed_world( for process in processes: if process.is_alive(): process.terminate() - raise RuntimeError( - f"timed out waiting for world_size={world_size} benchmark" - ) from exc + raise RuntimeError(f"timed out waiting for world_size={world_size} benchmark") from exc finally: for process in processes: process.join(timeout=60) @@ -1014,14 +971,10 @@ def _run_distributed_world( raise RuntimeError(result.get("traceback", str(result))) for process in processes: if process.exitcode != 0: - raise RuntimeError( - f"world_size={world_size} worker exited with {process.exitcode}" - ) + raise RuntimeError(f"world_size={world_size} worker exited with {process.exitcode}") return result - - def _topology_exactness_rows( distributed_rows: list[dict[str, Any]], ) -> list[dict[str, Any]]: @@ -1067,8 +1020,7 @@ def _distributed_platform_comparison_rows( if comparison_payload is None: return [] h100_lookup = { - (row["name"], row["direction"]): row - for row in comparison_payload["distributed"] + (row["name"], row["direction"]): row for row in comparison_payload["distributed"] } rows: list[dict[str, Any]] = [] for current in current_rows: @@ -1078,9 +1030,7 @@ def _distributed_platform_comparison_rows( continue h100_official_ms = float(h100["official_h100_ms"]) h100_deterministic_ms = float(h100["cuda_h100_ms"]) - mi300x_official_ms = float( - current["official_distributed"]["median_ms"] - ) + mi300x_official_ms = float(current["official_distributed"]["median_ms"]) mi300x_deterministic_ms = float(current["triton"]["median_ms"]) rows.append( { @@ -1143,20 +1093,11 @@ def _write_report( platform_comparison = _distributed_platform_comparison_rows( distributed_speed, comparison_payload ) - previous_comparison = _previous_deterministic_comparison_rows( - distributed_speed - ) - previous_reductions = [ - row["latency_reduction_ratio"] for row in previous_comparison - ] + previous_comparison = _previous_deterministic_comparison_rows(distributed_speed) + previous_reductions = [row["latency_reduction_ratio"] for row in previous_comparison] exactness_rows = _topology_exactness_rows(distributed_speed) - single_ratios = [ - row["latency_ratio_vs_official_tp1"] for row in single_speed - ] - platform_ratios = [ - row["deterministic_mi300x_over_h100_ratio"] - for row in platform_comparison - ] + single_ratios = [row["latency_ratio_vs_official_tp1"] for row in single_speed] + platform_ratios = [row["deterministic_mi300x_over_h100_ratio"] for row in platform_comparison] total_tp1_mismatch = sum( row[key] for row in exactness_rows @@ -1281,8 +1222,7 @@ def _write_report( "", "## Single-GPU FFN speed", "", - "Performance only; no official-versus-Triton accuracy metric is " - "reported here.", + "Performance only; no official-versus-Triton accuracy metric is " "reported here.", "", "| Shape / direction | Official Qwen3MLP TP=1 (ms) | Deterministic " "Triton, packed (ms) | Triton / official TP=1 |", @@ -1364,9 +1304,7 @@ def _write_report( if comparison_payload is not None: comparison_environment = comparison_payload["environment"] comparison_source = comparison_payload["source"] - local_single = { - (row["tokens"], row["direction"]): row for row in single_speed - } + local_single = {(row["tokens"], row["direction"]): row for row in single_speed} lines.extend( ( "", @@ -1402,11 +1340,7 @@ def _write_report( ) for row in comparison_payload["single_gpu"]: local = local_single[(row["tokens"], row["direction"])] - direction = ( - "forward" - if row["direction"] == "forward" - else "forward+backward" - ) + direction = "forward" if row["direction"] == "forward" else "forward+backward" lines.append( f"| M={row['tokens']}, {direction} | " f"{row['official_cpu_ms']:.4f} | " @@ -1480,8 +1414,7 @@ def _write_report( "", "## Figures", "", - "![Single-GPU CUDA, packed Triton, and CPU latency]" - "(single_gpu_overhead.png)", + "![Single-GPU CUDA, packed Triton, and CPU latency]" "(single_gpu_overhead.png)", "", "![Topology mismatch versus Triton TP=1](collective_overhead.png)", "", @@ -1490,9 +1423,7 @@ def _write_report( "", ) ) - (output_directory / "report.md").write_text( - "\n".join(lines), encoding="utf-8" - ) + (output_directory / "report.md").write_text("\n".join(lines), encoding="utf-8") def _write_figures( @@ -1521,16 +1452,13 @@ def _write_figures( single_rows = payload["single_gpu"]["speed"] if comparison_payload is None: single_labels = [ - f"M={row['tokens']}\n" - f"{'FWD' if row['direction'] == 'forward' else 'FWD+BWD'}" + f"M={row['tokens']}\n" f"{'FWD' if row['direction'] == 'forward' else 'FWD+BWD'}" for row in single_rows ] positions = np.arange(len(single_rows)) width = 0.37 figure, axis = plt.subplots(figsize=(17, 10)) - official_values = [ - row["official_tp1"]["median_ms"] for row in single_rows - ] + official_values = [row["official_tp1"]["median_ms"] for row in single_rows] triton_values = [row["triton"]["median_ms"] for row in single_rows] official_bars = axis.bar( positions - width / 2, @@ -1558,21 +1486,16 @@ def _write_figures( rotation=90, ) axis.set_yscale("log") - axis.set_xlabel( - "Token count M and measured direction\nH=4096, I=12288, BF16" - ) + axis.set_xlabel("Token count M and measured direction\nH=4096, I=12288, BF16") axis.set_ylabel("Median latency (ms, log scale)") axis.set_title("MI300X single-GPU FFN speed: official TP=1 vs Triton") axis.set_xticks(positions, single_labels) axis.legend(loc="upper left") figure.tight_layout() else: - local_lookup = { - (row["tokens"], row["direction"]): row for row in single_rows - } + local_lookup = {(row["tokens"], row["direction"]): row for row in single_rows} comparison_lookup = { - (row["tokens"], row["direction"]): row - for row in comparison_payload["single_gpu"] + (row["tokens"], row["direction"]): row for row in comparison_payload["single_gpu"] } figure, axes = plt.subplots(1, 2, figsize=(24, 10), sharey=True) series = ( @@ -1707,15 +1630,10 @@ def _write_figures( plt.close(figure) rows = payload["distributed_ffn"] - platform_comparison = _distributed_platform_comparison_rows( - rows, comparison_payload - ) + platform_comparison = _distributed_platform_comparison_rows(rows, comparison_payload) figure, axes = plt.subplots(1, 2, figsize=(26, 10), sharey=True) if platform_comparison: - comparison_lookup = { - (row["name"], row["direction"]): row - for row in platform_comparison - } + comparison_lookup = {(row["name"], row["direction"]): row for row in platform_comparison} distributed_series = ( ("H100 official distributed", "h100_official_distributed_ms", "#60a5fa"), ("H100 deterministic CUDA", "h100_deterministic_cuda_ms", "#dc2626"), @@ -1755,9 +1673,7 @@ def _write_figures( combined_rows.append(combined) for series_index, (label, key, color) in enumerate(distributed_series): values = [row[key] for row in combined_rows] - offset = ( - series_index - (len(distributed_series) - 1) / 2 - ) * width + offset = (series_index - (len(distributed_series) - 1) / 2) * width bars = axis.bar( positions + offset, values, @@ -1905,9 +1821,7 @@ def main() -> None: "contract": "same M=32 workload, direction, and TP/CP/SP topology", "rows": platform_comparison, } - previous_comparison = _previous_deterministic_comparison_rows( - payload["distributed_ffn"] - ) + previous_comparison = _previous_deterministic_comparison_rows(payload["distributed_ffn"]) if previous_comparison: payload["previous_deterministic_comparison"] = { "source": "previous checked MI300X benchmark before PR #357", diff --git a/benchmarks/profile_rocm_ffn.py b/benchmarks/profile_rocm_ffn.py index 099ffe0f..9b9413f2 100644 --- a/benchmarks/profile_rocm_ffn.py +++ b/benchmarks/profile_rocm_ffn.py @@ -145,9 +145,7 @@ def _build_case(args: argparse.Namespace, direction: str) -> FFNCase: device = torch.device("cuda", args.device) training = direction == "forward-backward" input_seed = ( - args.input_seed - if args.input_seed is not None - else _INPUT_SEEDS.get(args.tokens, 3010) + args.input_seed if args.input_seed is not None else _INPUT_SEEDS.get(args.tokens, 3010) ) hidden = _randn( (args.tokens, args.hidden_size), @@ -311,8 +309,7 @@ def _run_and_fingerprint( output = case.run(use_forward_weights=use_forward_weights) torch.cuda.synchronize() fingerprints = { - name: _tensor_fingerprint(tensor) - for name, tensor in case.result_tensors(output).items() + name: _tensor_fingerprint(tensor) for name, tensor in case.result_tensors(output).items() } del output case.clear_gradients() @@ -406,9 +403,7 @@ def _kernel_breakdown(rows: list[dict[str, Any]]) -> dict[str, Any]: else: category = "other_device_kernels" categories[category]["count"] += int(row["count"]) - categories[category]["self_device_time_us"] += float( - row["self_device_time_us"] - ) + categories[category]["self_device_time_us"] += float(row["self_device_time_us"]) total_device_us = sum( float(category["self_device_time_us"]) for category in categories.values() @@ -486,13 +481,9 @@ def _profile_case( # passes no forward-weight cache. In packed mode this exercises the original # per-call transpose path independently of the profiled candidate. standard_reference_fingerprints = ( - _run_and_fingerprint(case, use_forward_weights=False) - if not args.skip_bitwise_hash - else {} - ) - before_profile_fingerprints = ( - _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} + _run_and_fingerprint(case, use_forward_weights=False) if not args.skip_bitwise_hash else {} ) + before_profile_fingerprints = _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} torch.cuda.synchronize() device = torch.device("cuda", args.device) @@ -537,9 +528,7 @@ def _profile_case( kernel_breakdown = _kernel_breakdown(rows) _write_json(output_dir / f"{case.slug}.kernel_breakdown.json", kernel_breakdown) - key_averages = profiler.key_averages( - group_by_input_shape=effective_record_shapes - ) + key_averages = profiler.key_averages(group_by_input_shape=effective_record_shapes) summary = "\n\n".join( ( "Sorted by self device time\n" @@ -566,9 +555,7 @@ def _profile_case( except (AssertionError, RuntimeError, ValueError) as error: memory_timeline_error = f"{type(error).__name__}: {error}" - after_profile_fingerprints = ( - _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} - ) + after_profile_fingerprints = _run_and_fingerprint(case) if not args.skip_bitwise_hash else {} before_matches_reference = { name: before_profile_fingerprints.get(name) == fingerprint for name, fingerprint in standard_reference_fingerprints.items() @@ -621,9 +608,7 @@ def _profile_case( "kernel_breakdown": f"{case.slug}.kernel_breakdown.json", "latency": f"{case.slug}.latency.json", "correctness": f"{case.slug}.correctness.json", - "memory_timeline": ( - f"{case.slug}.memory.raw.json.gz" if args.profile_memory else "" - ), + "memory_timeline": (f"{case.slug}.memory.raw.json.gz" if args.profile_memory else ""), }, "profiler": { "active_steps": args.active_steps, @@ -712,11 +697,7 @@ def main() -> None: torch.backends.cuda.matmul.allow_tf32 = False args.output_dir.mkdir(parents=True, exist_ok=True) - directions = ( - ("forward", "forward-backward") - if args.direction == "both" - else (args.direction,) - ) + directions = ("forward", "forward-backward") if args.direction == "both" else (args.direction,) manifest: dict[str, Any] = { "environment": _environment(args), "workload": { @@ -745,9 +726,7 @@ def main() -> None: "profiler_use": "attribution and launch analysis only", "latency_use": "uninstrumented GPU events", "jit_and_tree_plan": "warmed before profiler starts", - "weight_packing": ( - "performed once during case construction outside all timing" - ), + "weight_packing": ("performed once during case construction outside all timing"), "bitwise_fingerprint": "SHA256 over raw BF16 bytes", "packed_bitwise_reference": ( "uncached standard qwen3_ffn path using the same canonical tensors" diff --git a/benchmarks/results/ws2_rocm_mi300x/results.json b/benchmarks/results/ws2_rocm_mi300x/results.json index 0f5b2176..36c2bc10 100644 --- a/benchmarks/results/ws2_rocm_mi300x/results.json +++ b/benchmarks/results/ws2_rocm_mi300x/results.json @@ -1907,4 +1907,4 @@ } } ] -} \ No newline at end of file +} diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/analyze_performance.py b/examples/vime_qwen3_8b_tp4_cp2_200/analyze_performance.py index 7063ea28..105cb431 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/analyze_performance.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/analyze_performance.py @@ -66,9 +66,7 @@ def parse_log(path: Path) -> tuple[dict[str, dict[int, dict[str, Any]]], np.ndar expected = list(range(200)) for kind, values in records.items(): if sorted(values) != expected: - raise RuntimeError( - f"{path}: {kind} has {len(values)} records; expected steps 0..199" - ) + raise RuntimeError(f"{path}: {kind} has {len(values)} records; expected steps 0..199") if len(progress_seconds) != 400: raise RuntimeError( f"{path}: expected 400 duplicated rollout progress events, got " @@ -102,9 +100,7 @@ def describe(values: np.ndarray) -> dict[str, float]: def moving_average(values: np.ndarray, window: int = 10) -> np.ndarray: - totals = np.convolve(values, np.ones(window, dtype=float), mode="full")[ - : len(values) - ] + totals = np.convolve(values, np.ones(window, dtype=float), mode="full")[: len(values)] counts = np.minimum(np.arange(1, len(values) + 1), window) return totals / counts @@ -164,15 +160,9 @@ def build_rows( "response_tokens_total": 8 * response_length[step], "rollout_time_s": rollout_time[step], "rollout_progress_time_s_rounded": progress_seconds[step], - "rollout_tok_per_gpu_s": float( - rollout_perf[step]["perf/tokens_per_gpu_per_sec"] - ), - "rollout_aggregate_tok_s": 8 - * response_length[step] - / rollout_time[step], - "rollout_truncated_ratio": float( - rollout_perf[step]["rollout/truncated_ratio"] - ), + "rollout_tok_per_gpu_s": float(rollout_perf[step]["perf/tokens_per_gpu_per_sec"]), + "rollout_aggregate_tok_s": 8 * response_length[step] / rollout_time[step], + "rollout_truncated_ratio": float(rollout_perf[step]["rollout/truncated_ratio"]), "update_weights_time_s": update_weights[step], "wait_residual_time_s": ( train_wait[step] - rollout_time[step] - update_weights[step] @@ -181,16 +171,10 @@ def build_rows( "actor_train_time_s": actor_train[step], "train_residual_time_s": train_time[step] - actor_train[step], "train_time_s": train_time[step], - "data_preprocess_time_s": float( - train_perf[step]["perf/data_preprocess_time"] - ), + "data_preprocess_time_s": float(train_perf[step]["perf/data_preprocess_time"]), "step_time_s": float(train_perf[step]["perf/step_time"]), - "actor_train_tok_s": float( - train_perf[step]["perf/actor_train_tok_per_s"] - ), - "actor_train_tflops": float( - train_perf[step]["perf/actor_train_tflops"] - ), + "actor_train_tok_s": float(train_perf[step]["perf/actor_train_tok_per_s"]), + "actor_train_tflops": float(train_perf[step]["perf/actor_train_tflops"]), } ) return rows @@ -214,15 +198,11 @@ def style_axis(axis: plt.Axes) -> None: def save_figure(fig: plt.Figure, output_dir: Path, stem: str) -> None: - fig.savefig( - output_dir / f"{stem}.png", dpi=220, bbox_inches="tight", facecolor="white" - ) + fig.savefig(output_dir / f"{stem}.png", dpi=220, bbox_inches="tight", facecolor="white") fig.savefig(output_dir / f"{stem}.pdf", bbox_inches="tight", facecolor="white") -def plot_decomposition( - group_rows: dict[str, list[dict[str, Any]]], output_dir: Path -) -> None: +def plot_decomposition(group_rows: dict[str, list[dict[str, Any]]], output_dir: Path) -> None: stages = [ ("rollout_time_s", "Rollout generation", "#E15759"), ("update_weights_time_s", "Weight update", "#F2B134"), @@ -236,10 +216,7 @@ def plot_decomposition( bottoms = np.zeros(2) for key, label, color in stages: values = np.asarray( - [ - statistics.fmean(float(row[key]) for row in group_rows[group]) - for group in groups - ] + [statistics.fmean(float(row[key]) for row in group_rows[group]) for group in groups] ) axis.bar(groups, values, bottom=bottoms, label=label, color=color, width=0.58) for index, value in enumerate(values): @@ -255,9 +232,7 @@ def plot_decomposition( ) bottoms += values for index, total in enumerate(bottoms): - axis.text( - index, total + 2, f"{total:.1f}s / step", ha="center", fontweight="bold" - ) + axis.text(index, total + 2, f"{total:.1f}s / step", ha="center", fontweight="bold") axis.set_title( "G11 vs G10 ยท Mean Step-Time Decomposition", fontsize=16, @@ -329,18 +304,14 @@ def plot_scaling(group_rows: dict[str, list[dict[str, Any]]], output_dir: Path) plt.close(fig) -def plot_time_series( - group_rows: dict[str, list[dict[str, Any]]], output_dir: Path -) -> None: +def plot_time_series(group_rows: dict[str, list[dict[str, Any]]], output_dir: Path) -> None: steps = np.arange(200) fig, axes = plt.subplots(2, 1, figsize=(13.5, 7.8), sharex=True) for group, raw_color, ma_color in ( ("G11", LIGHT_BLUE, BLUE), ("G10", LIGHT_RED, RED), ): - for axis, key in zip( - axes, ("rollout_time_s", "actor_train_time_s"), strict=True - ): + for axis, key in zip(axes, ("rollout_time_s", "actor_train_time_s"), strict=True): values = rows_array(group_rows[group], key) axis.plot(steps, values, color=raw_color, alpha=0.7, linewidth=0.9) axis.plot( @@ -369,9 +340,7 @@ def plot_time_series( plt.close(fig) -def plot_throughput( - group_rows: dict[str, list[dict[str, Any]]], output_dir: Path -) -> None: +def plot_throughput(group_rows: dict[str, list[dict[str, Any]]], output_dir: Path) -> None: fig, axes = plt.subplots(1, 2, figsize=(12.2, 5.3)) panels = [ ("rollout_tok_per_gpu_s", "Rollout throughput", "Tokens / GPU / s"), @@ -443,14 +412,10 @@ def summarize(group_rows: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: common_response_length = statistics.fmean( float(row["response_length_mean"]) for row in all_rows ) - common_total_length = statistics.fmean( - float(row["total_length_mean"]) for row in all_rows - ) + common_total_length = statistics.fmean(float(row["total_length_mean"]) for row in all_rows) summary: dict[str, Any] = {"groups": {}} for group, rows in group_rows.items(): - summary["groups"][group] = { - key: describe(rows_array(rows, key)) for key in keys - } + summary["groups"][group] = {key: describe(rows_array(rows, key)) for key in keys} summary["groups"][group]["rollout_length_regression"] = regression( rows_array(rows, "response_length_mean"), rows_array(rows, "rollout_time_s"), @@ -491,25 +456,17 @@ def summarize(group_rows: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: ] summary["gap"] = { "g11_minus_g10_step_time_s": step_gap, - "g11_minus_g10_total_hours": ( - g11["step_time_s"]["sum"] - g10["step_time_s"]["sum"] - ) - / 3600, + "g11_minus_g10_total_hours": (g11["step_time_s"]["sum"] - g10["step_time_s"]["sum"]) / 3600, "g10_step_time_reduction_fraction": 1 - g10["step_time_s"]["mean"] / g11["step_time_s"]["mean"], - "g10_end_to_end_speedup": g11["step_time_s"]["mean"] - / g10["step_time_s"]["mean"], + "g10_end_to_end_speedup": g11["step_time_s"]["mean"] / g10["step_time_s"]["mean"], "g10_rollout_throughput_speedup": g10["rollout_tok_per_gpu_s"]["mean"] / g11["rollout_tok_per_gpu_s"]["mean"], "g10_actor_throughput_speedup": g10["actor_train_tok_s"]["mean"] / g11["actor_train_tok_s"]["mean"], - "rollout_common_length_gap_s": g11["rollout_length_regression"][ - "prediction_at_common_x" - ] + "rollout_common_length_gap_s": g11["rollout_length_regression"]["prediction_at_common_x"] - g10["rollout_length_regression"]["prediction_at_common_x"], - "actor_common_length_gap_s": g11["actor_length_regression"][ - "prediction_at_common_x" - ] + "actor_common_length_gap_s": g11["actor_length_regression"]["prediction_at_common_x"] - g10["actor_length_regression"]["prediction_at_common_x"], "stage_contributions": {}, "bootstrap_mean_gap_ci95": {}, @@ -525,11 +482,27 @@ def summarize(group_rows: dict[str, list[dict[str, Any]]]) -> dict[str, Any]: rows_array(group_rows["G11"], key), rows_array(group_rows["G10"], key) ) summary["method_notes"] = { - "rollout_time": "Exact perf/rollout_time emitted by RolloutManager; tqdm integer seconds are retained only as an audit cross-check.", - "wait_residual": "train_wait - exact rollout_time - update_weights_time; captures wake/offload/orchestration outside the two named timers.", - "throughput": "Uses emitted perf/tokens_per_gpu_per_sec and perf/actor_train_tok_per_s, avoiding response-length confounding.", - "common_length": "Separate OLS fits for each arm evaluated at the pooled mean length; descriptive rather than causal because the two runs use different operator stacks, TE versions, revisions and generated sequences.", - "bootstrap": "Independent non-parametric bootstrap of per-step mean gaps with fixed seed 1234 and 20,000 draws.", + "rollout_time": ( + "Exact perf/rollout_time emitted by RolloutManager; " + "tqdm integer seconds are retained only as an audit cross-check." + ), + "wait_residual": ( + "train_wait - exact rollout_time - update_weights_time; " + "captures wake/offload/orchestration outside the two named timers." + ), + "throughput": ( + "Uses emitted perf/tokens_per_gpu_per_sec and perf/actor_train_tok_per_s, " + "avoiding response-length confounding." + ), + "common_length": ( + "Separate OLS fits for each arm evaluated at the pooled mean length; " + "descriptive rather than causal because the two runs use different operator stacks, " + "TE versions, revisions and generated sequences." + ), + "bootstrap": ( + "Independent non-parametric bootstrap of per-step mean gaps " + "with fixed seed 1234 and 20,000 draws." + ), } return summary @@ -545,9 +518,7 @@ def main() -> None: parser.add_argument("--g11-log", type=Path, help="sealed G11 run.log") parser.add_argument("--output-dir", type=Path, required=True) args = parser.parse_args() - if args.data_dir is not None and ( - args.g10_log is not None or args.g11_log is not None - ): + if args.data_dir is not None and (args.g10_log is not None or args.g11_log is not None): parser.error("use either --data-dir or both --g10-log/--g11-log") if args.data_dir is not None: g10_log = args.data_dir / "g10.run.log" @@ -564,8 +535,7 @@ def main() -> None: "G10": parse_log(g10_log), } group_rows = { - group: build_rows(group, records, progress) - for group, (records, progress) in parsed.items() + group: build_rows(group, records, progress) for group, (records, progress) in parsed.items() } all_rows = group_rows["G11"] + group_rows["G10"] write_csv(args.output_dir / "step-metrics.csv", all_rows) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/collect_results.py b/examples/vime_qwen3_8b_tp4_cp2_200/collect_results.py index 70366da5..43aab24b 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/collect_results.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/collect_results.py @@ -59,12 +59,8 @@ def _run_rows(run_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: if not validation.get("passed"): raise ValueError(f"sealed run has a failed validation: {run_dir}") records = _records(run_dir / "run.log") - validation_rows = { - int(row["step"]): row for row in validation["train_rollout_logprob"]["rows"] - } - steps = sorted( - set(records["rollout"]) | set(records["step"]) | set(records["perf"]) - ) + validation_rows = {int(row["step"]): row for row in validation["train_rollout_logprob"]["rows"]} + steps = sorted(set(records["rollout"]) | set(records["step"]) | set(records["perf"])) rows = [] for index in steps: rollout = records["rollout"].get(index, {}) @@ -79,9 +75,7 @@ def _run_rows(run_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: "seed": manifest["seed"], "rollout_seed": manifest["rollout_seed"], "step": index, - "framework_consistency": manifest["arm"][ - "framework_use_rollout_logprobs" - ], + "framework_consistency": manifest["arm"]["framework_use_rollout_logprobs"], "operator_case": manifest["arm"]["logp_case"], "reward": _value(rollout, "rollout/rewards"), "raw_reward": _value(rollout, "rollout/raw_reward"), @@ -95,9 +89,7 @@ def _run_rows(run_dir: Path) -> tuple[dict[str, Any], list[dict[str, Any]]]: "ppo_kl": _value(step, "train/ppo_kl"), "grad_norm": _value(step, "train/grad_norm"), "mean_abs_dlogp": _value(step, "train/train_rollout_logprob_abs_diff"), - "max_abs_dlogp": _value( - step, "train/train_current_rollout_logprob_max_abs_diff" - ), + "max_abs_dlogp": _value(step, "train/train_current_rollout_logprob_max_abs_diff"), "mismatch_count": exact.get("bitwise_mismatch_count"), "active_token_count": exact.get("active_token_count"), "rollout_time": _value(perf, "perf/rollout_time"), @@ -147,8 +139,7 @@ def _summaries(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: weighted_abs_numerator = sum( float(item["mean_abs_dlogp"]) * float(item["active_token_count"]) for item in items - if item.get("mean_abs_dlogp") is not None - and item.get("active_token_count") is not None + if item.get("mean_abs_dlogp") is not None and item.get("active_token_count") is not None ) token_total = sum(tokens) summaries.append( @@ -159,9 +150,7 @@ def _summaries(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: "round_count": len(items), "active_token_exposure": token_total, "bitwise_mismatch_count": sum(mismatches), - "bitwise_mismatch_rate": ( - sum(mismatches) / token_total if token_total else None - ), + "bitwise_mismatch_rate": (sum(mismatches) / token_total if token_total else None), "mean_abs_dlogp_token_weighted": ( weighted_abs_numerator / token_total if token_total else None ), @@ -170,9 +159,7 @@ def _summaries(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: mean(_finite(items, "reward")) if _finite(items, "reward") else None ), "raw_reward_mean": ( - mean(_finite(items, "raw_reward")) - if _finite(items, "raw_reward") - else None + mean(_finite(items, "raw_reward")) if _finite(items, "raw_reward") else None ), "truncated_ratio_mean": ( mean(_finite(items, "truncated_ratio")) @@ -180,9 +167,7 @@ def _summaries(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: else None ), "step_time_mean": ( - mean(_finite(items, "step_time")) - if _finite(items, "step_time") - else None + mean(_finite(items, "step_time")) if _finite(items, "step_time") else None ), "actor_tokens_per_second_mean": ( mean(_finite(items, "actor_tokens_per_second")) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/plot_results.py b/examples/vime_qwen3_8b_tp4_cp2_200/plot_results.py index f2fa4fbd..7afe10d2 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/plot_results.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/plot_results.py @@ -38,8 +38,7 @@ def _parse(value: str) -> Any: def _load(path: Path, phase: str | None) -> list[dict[str, Any]]: with path.open(encoding="utf-8", newline="") as handle: rows = [ - {key: _parse(value) for key, value in row.items()} - for row in csv.DictReader(handle) + {key: _parse(value) for key, value in row.items()} for row in csv.DictReader(handle) ] if phase is not None: rows = [row for row in rows if row["phase"] == phase] @@ -66,10 +65,7 @@ def _series( def _moving_average(values: list[float], window: int) -> list[float]: - return [ - mean(values[max(0, index - window + 1) : index + 1]) - for index in range(len(values)) - ] + return [mean(values[max(0, index - window + 1) : index + 1]) for index in range(len(values))] def _plot_series(axis, rows, metric, title, *, moving_average=1, symlog=False): @@ -86,12 +82,8 @@ def _plot_series(axis, rows, metric, title, *, moving_average=1, symlog=False): markersize=5, ) if any(spreads): - lower = [ - center - spread for center, spread in zip(centers, spreads, strict=True) - ] - upper = [ - center + spread for center, spread in zip(centers, spreads, strict=True) - ] + lower = [center - spread for center, spread in zip(centers, spreads, strict=True)] + upper = [center + spread for center, spread in zip(centers, spreads, strict=True)] axis.fill_between(steps, lower, upper, color=color, alpha=0.15) if symlog: axis.set_yscale("symlog", linthresh=1e-9) @@ -108,9 +100,7 @@ def _mismatch_rate_rows(rows: list[dict[str, Any]]) -> list[dict[str, Any]]: tokens = row.get("active_token_count") updated["mismatch_rate"] = ( float(mismatch) / float(tokens) - if isinstance(mismatch, (int, float)) - and isinstance(tokens, (int, float)) - and tokens + if isinstance(mismatch, (int, float)) and isinstance(tokens, (int, float)) and tokens else None ) result.append(updated) @@ -168,9 +158,7 @@ def _save_optimization(rows, output, dpi): import matplotlib.pyplot as plt figure, axes = plt.subplots(2, 2, figsize=(13, 8), constrained_layout=True) - _plot_series( - axes[0, 0], rows, "pg_loss", "GRPO policy-gradient loss", symlog=True - ) + _plot_series(axes[0, 0], rows, "pg_loss", "GRPO policy-gradient loss", symlog=True) _plot_series(axes[0, 1], rows, "pg_clipfrac", "Policy ratio clipped fraction") _plot_series(axes[1, 0], rows, "ppo_kl", "Training PPO KL", symlog=True) _plot_series(axes[1, 1], rows, "grad_norm", "Gradient norm", symlog=True) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/prepare_dapo_data.py b/examples/vime_qwen3_8b_tp4_cp2_200/prepare_dapo_data.py index 8027b182..6974856d 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/prepare_dapo_data.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/prepare_dapo_data.py @@ -70,11 +70,7 @@ def convert(source: Path, output: Path) -> dict[str, Any]: reward_model = row.get("reward_model") or {} label = reward_model.get("ground_truth") prompt = row.get("prompt") - if ( - not isinstance(prompt, list) - or not isinstance(label, str) - or not label - ): + if not isinstance(prompt, list) or not isinstance(label, str) or not label: raise ValueError(f"invalid prompt or ground truth for row {row_id}") record = { "prompt": prompt, @@ -84,9 +80,7 @@ def convert(source: Path, output: Path) -> dict[str, Any]: "source_index": row_id, "reward_style": reward_model.get("style"), } - destination.write( - json.dumps(record, ensure_ascii=False, separators=(",", ":")) - ) + destination.write(json.dumps(record, ensure_ascii=False, separators=(",", ":"))) destination.write("\n") rows_written += 1 partial.replace(output) diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/run.py b/examples/vime_qwen3_8b_tp4_cp2_200/run.py index 5bb0af6a..be32a5f3 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/run.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/run.py @@ -53,9 +53,7 @@ def validate_config(config: Mapping[str, Any]) -> None: or not isinstance(rollout, Mapping) or not isinstance(provider, Mapping) ): - raise ValueError( - "training, rollout, and linear_logp_provider sections are required" - ) + raise ValueError("training, rollout, and linear_logp_provider sections are required") expected = { "tensor_model_parallel_size": 4, "context_parallel_size": 2, @@ -66,15 +64,10 @@ def validate_config(config: Mapping[str, Any]) -> None: if training.get(name) != value: raise ValueError(f"training.{name} must be {value!r}") if rollout.get("top_p") != 1.0: - raise ValueError( - "rollout.top_p must remain 1.0 for the strict provider contract" - ) + raise ValueError("rollout.top_p must remain 1.0 for the strict provider contract") if provider.get("mode") != "strict": raise ValueError("linear_logp_provider.mode must be strict") - if ( - provider.get("path") - != "rl_engine.integrations.vime.linear_logp_provider.provider" - ): + if provider.get("path") != "rl_engine.integrations.vime.linear_logp_provider.provider": raise ValueError("example must use the RL-Kernel Vime provider") if provider.get("backend_id") != "rlkernel.linear_logp.bitwise.v1": raise ValueError("example must pin the deterministic vocab-parallel backend") @@ -87,13 +80,8 @@ def load_runtime_evidence(path: Path | None) -> dict[str, Any] | None: return None with path.open(encoding="utf-8") as handle: value = json.load(handle) - if ( - not isinstance(value, dict) - or value.get("schema_version") != RUNTIME_EVIDENCE_SCHEMA - ): - raise ValueError( - f"runtime evidence must use schema {RUNTIME_EVIDENCE_SCHEMA!r}" - ) + if not isinstance(value, dict) or value.get("schema_version") != RUNTIME_EVIDENCE_SCHEMA: + raise ValueError(f"runtime evidence must use schema {RUNTIME_EVIDENCE_SCHEMA!r}") return value @@ -112,19 +100,13 @@ def _operator_evidence_status(evidence: Mapping[str, Any] | None, operator: str) if not isinstance(comparison, Mapping) or comparison.get("passed") is not True: return "failed" required_identity = ("implementation_id", "backend_id", "contract_id") - if any( - not training.get(name) or not rollout.get(name) for name in required_identity - ): + if any(not training.get(name) or not rollout.get(name) for name in required_identity): return "failed" if training["implementation_id"] != rollout["implementation_id"]: return "failed" for metric in _OPERATOR_METRICS[operator]: value = comparison.get(metric) - if ( - not isinstance(value, (int, float)) - or isinstance(value, bool) - or value != 0.0 - ): + if not isinstance(value, (int, float)) or isinstance(value, bool) or value != 0.0: return "failed" return "passed" @@ -137,9 +119,7 @@ def validate_runtime_evidence(evidence: Mapping[str, Any] | None) -> None: for operator in _OPERATOR_METRICS: status = _operator_evidence_status(evidence, operator) if status == "failed": - raise ValueError( - f"runtime evidence for {operator} is incomplete or non-zero" - ) + raise ValueError(f"runtime evidence for {operator} is incomplete or non-zero") def _revision(path: Path) -> str | None: @@ -190,13 +170,9 @@ def build_report( ) -> dict[str, Any]: provider_active = PROVIDER_MARKER in log_text fallback_observed = any(marker in log_text for marker in FALLBACK_MARKERS) - strict_provider_passed = ( - status == "passed" and provider_active and not fallback_observed - ) + strict_provider_passed = status == "passed" and provider_active and not fallback_observed effective_status = ( - "passed" - if strict_provider_passed - else ("failed" if status == "passed" else status) + "passed" if strict_provider_passed else ("failed" if status == "passed" else status) ) attention_status = _operator_evidence_status(runtime_evidence, "attention") ffn_status = _operator_evidence_status(runtime_evidence, "ffn") @@ -230,9 +206,7 @@ def build_report( None if runtime_evidence_path is None else str(runtime_evidence_path) ), }, - "runtime_evidence": ( - None if runtime_evidence is None else dict(runtime_evidence) - ), + "runtime_evidence": (None if runtime_evidence is None else dict(runtime_evidence)), "revisions": { "vime": _revision(vime_root), "rl_kernel": _revision(rl_kernel_root), @@ -243,17 +217,13 @@ def build_report( def main(argv: list[str] | None = None) -> int: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=DEFAULT_CONFIG) - parser.add_argument( - "--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", ".")) - ) + parser.add_argument("--vime-root", type=Path, default=Path(os.environ.get("VIME_ROOT", "."))) parser.add_argument( "--rl-kernel-root", type=Path, default=Path(os.environ.get("RL_KERNEL_ROOT", ".")), ) - parser.add_argument( - "--output", type=Path, default=Path("qwen3_8b_tp4_cp2.validation.json") - ) + parser.add_argument("--output", type=Path, default=Path("qwen3_8b_tp4_cp2.validation.json")) parser.add_argument( "--runtime-evidence", type=Path, @@ -304,9 +274,7 @@ def main(argv: list[str] | None = None) -> int: runtime_evidence_path=args.runtime_evidence, ) args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text( - json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" - ) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") print(json.dumps(report, indent=2, sort_keys=True)) return 0 if report["status"] in {"passed", "not_run"} else 1 diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/validate_artifacts.py b/examples/vime_qwen3_8b_tp4_cp2_200/validate_artifacts.py index 670087fc..bc7a1705 100644 --- a/examples/vime_qwen3_8b_tp4_cp2_200/validate_artifacts.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/validate_artifacts.py @@ -52,14 +52,11 @@ def validate_readbacks(readbacks: list[dict[str, Any]]) -> dict[str, Any]: if value.get("fallbacks"): errors.append(f"{label} recorded fallback: {value['fallbacks']}") for module in _MODULES: - hook_count = sum( - module in value.get("installed_hooks", {}) for value in matching - ) + hook_count = sum(module in value.get("installed_hooks", {}) for value in matching) records = [ value["operators"][module] for value in matching - if isinstance(value.get("operators"), Mapping) - and module in value["operators"] + if isinstance(value.get("operators"), Mapping) and module in value["operators"] ] call_count = sum(int(record.get("call_count", 0)) for record in records) if hook_count == 0: @@ -74,9 +71,7 @@ def validate_readbacks(readbacks: list[dict[str, Any]]) -> dict[str, Any]: f"{label} logp used {backend!r}, expected {_STRICT_LOGP_BACKEND!r}" ) elif not backend.startswith(_BACKEND_PREFIXES): - errors.append( - f"{label} {module} used unexpected backend {backend!r}" - ) + errors.append(f"{label} {module} used unexpected backend {backend!r}") if _contains_triton(record): errors.append(f"{label} {module} used Triton") if _runtime_platform(record.get("provenance")) != "cuda": @@ -119,9 +114,7 @@ def compare_train_rollout_logps(paths: list[Path]) -> dict[str, Any]: if not isinstance(training_values, (list, tuple)) or not isinstance( rollout_values, (list, tuple) ): - errors.append( - f"{path} rollout_data lacks list log_probs/rollout_log_probs" - ) + errors.append(f"{path} rollout_data lacks list log_probs/rollout_log_probs") continue if len(training_values) != len(rollout_values): errors.append( @@ -130,9 +123,7 @@ def compare_train_rollout_logps(paths: list[Path]) -> dict[str, Any]: ) samples = [ {"log_probs": training, "rollout_log_probs": rollout} - for training, rollout in zip( - training_values, rollout_values, strict=False - ) + for training, rollout in zip(training_values, rollout_values, strict=False) ] for sample_index, sample in enumerate(samples): if not isinstance(sample, Mapping): @@ -140,9 +131,7 @@ def compare_train_rollout_logps(paths: list[Path]) -> dict[str, Any]: continue training = sample.get("log_probs") rollout = sample.get("rollout_log_probs") - if not isinstance(training, torch.Tensor) or not isinstance( - rollout, torch.Tensor - ): + if not isinstance(training, torch.Tensor) or not isinstance(rollout, torch.Tensor): errors.append( f"{path} sample {sample_index} lacks tensor log_probs/rollout_log_probs" ) @@ -214,9 +203,7 @@ def main(argv: list[str] | None = None) -> int: "error": str(exc), } args.output.parent.mkdir(parents=True, exist_ok=True) - args.output.write_text( - json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" - ) + args.output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") print(json.dumps(report, indent=2, sort_keys=True)) return 0 if report["passed"] else 1 diff --git a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py index 70e56ca9..4ff0dba8 100755 --- a/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py +++ b/examples/vime_qwen3_8b_tp4_cp2_200/validate_run.py @@ -40,9 +40,7 @@ "ffn": "ffn_case", "logp": "logp_case", } -RL_KERNEL_LINEAR_LOGP_PROVIDER = ( - "rl_engine.integrations.vime.linear_logp_provider.provider" -) +RL_KERNEL_LINEAR_LOGP_PROVIDER = "rl_engine.integrations.vime.linear_logp_provider.provider" VIME_NATIVE_LINEAR_LOGP_MARKER = ( "linear_logp native active: " "backend_id=vime.utils.ppo_utils.calculate_log_probs_and_entropy " @@ -82,14 +80,11 @@ def _parse_runtime_records(log_text: str) -> dict[str, dict[int, dict[str, Any]] def _validate_cudagraph(log_text: str, manifest: Mapping[str, Any]) -> dict[str, Any]: execution = manifest.get("vllm_execution", {}) - capture_sizes = ( - execution.get("capture_sizes", []) if isinstance(execution, Mapping) else [] - ) + capture_sizes = execution.get("capture_sizes", []) if isinstance(execution, Mapping) else [] compact_sizes = "[" + ",".join(str(value) for value in capture_sizes) + "]" checks = { "launcher_marker": any( - f"{marker}: {compact_sizes}" in log_text - for marker in CUDA_GRAPH_LAUNCHER_MARKERS + f"{marker}: {compact_sizes}" in log_text for marker in CUDA_GRAPH_LAUNCHER_MARKERS ), "engine_mode": bool(re.search(r"cudagraph_mode.*FULL_DECODE_ONLY", log_text)), "not_eager": "enforce_eager=False" in log_text, @@ -161,19 +156,12 @@ def _validate_readbacks( records = [ value["operators"][module] for value in matching - if isinstance(value.get("operators"), Mapping) - and module in value["operators"] + if isinstance(value.get("operators"), Mapping) and module in value["operators"] ] - installed_count = sum( - module in value.get("installed_hooks", {}) for value in matching - ) + installed_count = sum(module in value.get("installed_hooks", {}) for value in matching) call_count = sum(int(record.get("call_count", 0)) for record in records) - implementations = sorted( - {str(record.get("implementation", "")) for record in records} - ) - backend_ids = sorted( - {str(record.get("backend_id", "")) for record in records} - ) + implementations = sorted({str(record.get("implementation", "")) for record in records}) + backend_ids = sorted({str(record.get("backend_id", "")) for record in records}) case_ids = sorted({str(record.get("case_id", "")) for record in records}) native_megatron_logp = ( framework == "megatron" @@ -188,9 +176,7 @@ def _validate_readbacks( f"{label} production logp unexpectedly installed an RL-Kernel hook" ) if records: - errors.append( - f"{label} production logp unexpectedly entered provider readback" - ) + errors.append(f"{label} production logp unexpectedly entered provider readback") if not marker_present: errors.append( f"{label} production logp did not report Vime's native backend marker" @@ -203,9 +189,7 @@ def _validate_readbacks( "implementations": implementations, "backend_ids": backend_ids, "native_marker_present": marker_present, - "native_backend_id": ( - "vime.utils.ppo_utils.calculate_log_probs_and_entropy" - ), + "native_backend_id": ("vime.utils.ppo_utils.calculate_log_probs_and_entropy"), } continue if installed_count == 0: @@ -225,20 +209,17 @@ def _validate_readbacks( errors.append(f"{label} {module} did not report CUDA execution") if record.get("provenance", {}).get("fallback") is True: errors.append(f"{label} {module} provenance recorded fallback") - if expected == "rl_kernel" and not str( - record.get("backend_id", "") - ).startswith("rlkernel."): + if expected == "rl_kernel" and not str(record.get("backend_id", "")).startswith( + "rlkernel." + ): errors.append(f"{label} {module} did not use an RL-Kernel backend") provenance = record.get("provenance", {}) reported_backend_ids = _reported_backend_ids(record) - strict_execution = ( - isinstance(provenance, Mapping) - and ( - provenance.get("deterministic_linear_logp") is True - or ( - isinstance(provenance.get("execution"), Mapping) - and provenance["execution"].get("strict_backend") is True - ) + strict_execution = isinstance(provenance, Mapping) and ( + provenance.get("deterministic_linear_logp") is True + or ( + isinstance(provenance.get("execution"), Mapping) + and provenance["execution"].get("strict_backend") is True ) ) if expected == "production" and ( @@ -320,8 +301,7 @@ def _validate_runtime_logprobs( if len(rows) != expected_rounds: errors.append(f"observed {len(rows)} train steps, expected {expected_rounds}") bitwise_zero = bool(rows) and all( - row["bitwise_mismatch_count"] == 0.0 and row["max_abs_dlogp"] == 0.0 - for row in rows + row["bitwise_mismatch_count"] == 0.0 and row["max_abs_dlogp"] == 0.0 for row in rows ) if require_zero and not bitwise_zero: errors.append("R/R arm did not achieve bitwise-zero runtime metrics") @@ -334,9 +314,7 @@ def _validate_runtime_logprobs( ), "bitwise_zero": bitwise_zero, "rows": rows, - "total_active_token_exposure": sum( - row["active_token_count"] or 0.0 for row in rows - ), + "total_active_token_exposure": sum(row["active_token_count"] or 0.0 for row in rows), } @@ -345,9 +323,7 @@ def _inspect_offline_dumps(directory: Path) -> dict[str, Any]: comparable = 0 for path in paths: payload = torch.load(path, map_location="cpu", weights_only=False) - rollout_data = ( - payload.get("rollout_data", {}) if isinstance(payload, Mapping) else {} - ) + rollout_data = payload.get("rollout_data", {}) if isinstance(payload, Mapping) else {} if isinstance(rollout_data, Mapping) and "log_probs" in rollout_data: comparable += 1 return { @@ -357,7 +333,10 @@ def _inspect_offline_dumps(directory: Path) -> dict[str, Any]: "reason": ( None if paths and comparable == len(paths) - else "current VIME dump lacks captured training log_probs; runtime exact metrics are used" + else ( + "current VIME dump lacks captured training log_probs; " + "runtime exact metrics are used" + ) ), } @@ -371,9 +350,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: records = _parse_runtime_records(log_text) require_zero = all(str(arm[CASE_FIELDS[module]]) == "R/R" for module in MODULES) cudagraph = _validate_cudagraph(log_text, manifest) - readbacks = _validate_readbacks( - _load_readbacks(run_dir / "readbacks"), arm, log_text - ) + readbacks = _validate_readbacks(_load_readbacks(run_dir / "readbacks"), arm, log_text) logprobs = _validate_runtime_logprobs( records["step"], int(manifest["num_rollout"]), @@ -382,10 +359,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: ) global_errors = [] algorithm = manifest.get("algorithm", {}) - if ( - not isinstance(algorithm, Mapping) - or algorithm.get("advantage_estimator") != "grpo" - ): + if not isinstance(algorithm, Mapping) or algorithm.get("advantage_estimator") != "grpo": global_errors.append("manifest does not explicitly select GRPO") train_command = manifest.get("train_command", []) expected_algorithm_pair = ["--advantage-estimator", "grpo"] @@ -433,9 +407,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "production Megatron logp must not configure a linear_logp provider" ) if "--linear-logp-provider-mode" in train_command: - global_errors.append( - "production Megatron logp must not configure provider mode" - ) + global_errors.append("production Megatron logp must not configure provider mode") elif not has_provider or not has_strict_mode: global_errors.append( "RL-Kernel Megatron logp must configure the strict RL-Kernel provider" @@ -446,9 +418,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "recompute_num_layers": 1, } if manifest.get("training_memory") != expected_recompute: - global_errors.append( - "manifest does not contain the required recompute configuration" - ) + global_errors.append("manifest does not contain the required recompute configuration") if re.search(r"fallback=true", log_text, re.IGNORECASE): global_errors.append("run log contains fallback=true") if "Traceback (most recent call last)" in log_text: @@ -458,10 +428,7 @@ def validate_run(run_dir: Path) -> dict[str, Any]: "run_id": manifest.get("run_id"), "group": arm.get("group"), "passed": bool( - cudagraph["passed"] - and readbacks["passed"] - and logprobs["passed"] - and not global_errors + cudagraph["passed"] and readbacks["passed"] and logprobs["passed"] and not global_errors ), "errors": global_errors, "cudagraph": cudagraph, @@ -488,9 +455,7 @@ def main(argv: list[str] | None = None) -> int: "passed": False, "errors": [f"{type(exc).__name__}: {exc}"], } - output.write_text( - json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8" - ) + output.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") print(json.dumps(report, indent=2, sort_keys=True)) if args.seal and report["passed"]: (run_dir / "COMPLETE").touch(exist_ok=False) diff --git a/examples/vime_rocm_attention_ablation/run.py b/examples/vime_rocm_attention_ablation/run.py index 9a757541..d9ecb4a1 100644 --- a/examples/vime_rocm_attention_ablation/run.py +++ b/examples/vime_rocm_attention_ablation/run.py @@ -160,9 +160,7 @@ def validate(self, *, require_paths: bool) -> None: raise ValueError("Ray ports must be between 1024 and 65535") if abs(self.ray_port - self.ray_dashboard_port) < len(CASE_ORDER): raise ValueError("Ray GCS and dashboard port ranges overlap") - dashboard_ports = range( - self.ray_dashboard_port, self.ray_dashboard_port + len(CASE_ORDER) - ) + dashboard_ports = range(self.ray_dashboard_port, self.ray_dashboard_port + len(CASE_ORDER)) if any(10001 <= port <= 19999 for port in dashboard_ports): raise ValueError("Ray dashboard ports overlap the default client/worker range") if not require_paths: @@ -186,10 +184,7 @@ def validate(self, *, require_paths: bool) -> None: if not (self.rl_kernel_root / "rl_engine").is_dir(): raise FileNotFoundError("rl_kernel_root does not contain rl_engine") metrics_hook = ( - self.rl_kernel_root - / "examples" - / "vime_rocm_attention_ablation" - / "tis_metrics.py" + self.rl_kernel_root / "examples" / "vime_rocm_attention_ablation" / "tis_metrics.py" ) if not metrics_hook.is_file(): raise FileNotFoundError(f"Attention mismatch metrics hook is missing: {metrics_hook}") @@ -253,9 +248,7 @@ def frozen_parameters(self) -> dict[str, Any]: }, "seed": self.seed, "rollout_seed": self.rollout_seed, - "mismatch_metrics_hook": ( - "vime_rocm_attention_ablation.tis_metrics.metrics_only_tis" - ), + "mismatch_metrics_hook": ("vime_rocm_attention_ablation.tis_metrics.metrics_only_tis"), "ffn_case": "R/R", "logp_case": "R/R", "real_vocab_size": self.real_vocab_size, @@ -280,8 +273,7 @@ def _validate_checkpoint_marker(checkpoint: Path) -> None: marker = checkpoint / "latest_checkpointed_iteration.txt" if not marker.is_file(): raise FileNotFoundError( - "reference_checkpoint is not a Megatron checkpoint: missing " - f"{marker}" + "reference_checkpoint is not a Megatron checkpoint: missing " f"{marker}" ) value = marker.read_text(encoding="utf-8").strip() if value == "release": @@ -327,8 +319,7 @@ def _validate_rl_kernel_plugin_installation() -> None: if not matching: values = sorted({entry_point.value for entry_point in candidates}) raise RuntimeError( - "the visible rl_kernel vLLM plugin entry point has an unexpected target: " - f"{values!r}" + "the visible rl_kernel vLLM plugin entry point has an unexpected target: " f"{values!r}" ) installed_name = distribution.metadata.get("Name", "RL-Kernel") @@ -525,18 +516,14 @@ def build_arm_environment( "RL_KERNEL_VLLM_PADDED_VOCAB_SIZE": str(config.padded_vocab_size), "RL_KERNEL_VLLM_INTEGRATION": "1", "RL_KERNEL_READBACK_DIR": str((arm_dir / "readbacks").resolve()), - "RL_KERNEL_MISMATCH_SIDECAR_DIR": str( - (arm_dir / "mismatch_sidecars").resolve() - ), + "RL_KERNEL_MISMATCH_SIDECAR_DIR": str((arm_dir / "mismatch_sidecars").resolve()), "RLK_ABLATION_CASE_ID": case_id, "RLK_ABLATION_ARM_DIR": str(arm_dir.resolve()), "RLK_ABLATION_VIME_ROOT": str(config.vime_root.resolve()), "RLK_ABLATION_RL_KERNEL_ROOT": str(config.rl_kernel_root.resolve()), "RLK_ABLATION_MEGATRON_ROOT": str(config.megatron_root.resolve()), "RLK_ABLATION_MODEL_ROOT": str(config.model_root.resolve()), - "RLK_ABLATION_REFERENCE_CHECKPOINT": str( - config.reference_checkpoint.resolve() - ), + "RLK_ABLATION_REFERENCE_CHECKPOINT": str(config.reference_checkpoint.resolve()), "RLK_ABLATION_PROMPT_DATA": str(config.prompt_data.resolve()), "RLK_ABLATION_NUM_GPUS": str(config.num_gpus), "RLK_ABLATION_TP_SIZE": str(config.tensor_parallel_size), @@ -555,9 +542,7 @@ def build_arm_environment( "RLK_ABLATION_SEED": str(config.seed), "RLK_ABLATION_ROLLOUT_SEED": str(config.rollout_seed), "RLK_ABLATION_RAY_PORT": str(config.ray_port + arm_index), - "RLK_ABLATION_RAY_DASHBOARD_PORT": str( - config.ray_dashboard_port + arm_index - ), + "RLK_ABLATION_RAY_DASHBOARD_PORT": str(config.ray_dashboard_port + arm_index), "RLK_ABLATION_RAY_DASHBOARD_AGENT_PORT": str( config.ray_dashboard_port + arm_index + 10_000 ), @@ -763,9 +748,7 @@ def config_from_args(args: argparse.Namespace) -> MatrixConfig: "rl_kernel_root": _path_argument(args.rl_kernel_root, "RL_KERNEL_ROOT"), "megatron_root": _path_argument(args.megatron_root, "MEGATRON_ROOT"), "model_root": _path_argument(args.model_root, "MODEL_ROOT"), - "reference_checkpoint": _path_argument( - args.reference_checkpoint, "TORCH_DIST_ROOT" - ), + "reference_checkpoint": _path_argument(args.reference_checkpoint, "TORCH_DIST_ROOT"), "prompt_data": _path_argument(args.prompt_data, "PROMPT_DATA"), } missing = [name for name, value in values.items() if value is None] diff --git a/examples/vime_rocm_attention_ablation/run_full_pp_200.py b/examples/vime_rocm_attention_ablation/run_full_pp_200.py index 14ff1253..c3ed2831 100644 --- a/examples/vime_rocm_attention_ablation/run_full_pp_200.py +++ b/examples/vime_rocm_attention_ablation/run_full_pp_200.py @@ -82,8 +82,7 @@ def main() -> int: "passed": report["passed"], "errors": report["errors"], "metrics": report["metrics"], - "frozen_sources_match": frozen_before["fingerprint"] - == frozen_after["fingerprint"], + "frozen_sources_match": frozen_before["fingerprint"] == frozen_after["fingerprint"], } write_report(RUN_DIR / "single-arm-summary.json", summary) print(json.dumps(summary, indent=2, sort_keys=True), flush=True) diff --git a/examples/vime_rocm_attention_ablation/tis_metrics.py b/examples/vime_rocm_attention_ablation/tis_metrics.py index c71e4e30..ebef88f6 100644 --- a/examples/vime_rocm_attention_ablation/tis_metrics.py +++ b/examples/vime_rocm_attention_ablation/tis_metrics.py @@ -76,8 +76,7 @@ def _write_sidecar( _cpu_vector(value, label="train_log_probs") for value in train_log_probs ], "rollout_log_probs": [ - _cpu_vector(value, label="rollout_log_probs") - for value in rollout_log_probs + _cpu_vector(value, label="rollout_log_probs") for value in rollout_log_probs ], "loss_masks": [_cpu_vector(value, label="loss_masks") for value in loss_masks], "total_lengths": [int(value) for value in total_lengths], diff --git a/examples/vime_rocm_attention_ablation/validate_artifacts.py b/examples/vime_rocm_attention_ablation/validate_artifacts.py index a9e491b8..df657faa 100644 --- a/examples/vime_rocm_attention_ablation/validate_artifacts.py +++ b/examples/vime_rocm_attention_ablation/validate_artifacts.py @@ -102,9 +102,7 @@ def _contains_string(value: Any, needle: str) -> bool: def _runtime_platform(provenance: Any) -> str | None: values = _values_for_keys(provenance, {"runtime_platform", "platform"}) normalized = { - str(value).strip().lower() - for value in values - if isinstance(value, str) and value.strip() + str(value).strip().lower() for value in values if isinstance(value, str) and value.strip() } if normalized & {"rocm", "hip"}: return "rocm" @@ -233,9 +231,7 @@ def _validate_rlkernel_record( errors.append(f"{label} selected RL-Kernel but reported backend {backend_id!r}") if _runtime_platform(provenance) != "rocm": errors.append(f"{label} RL-Kernel route did not prove ROCm execution") - if _truthy_flag(provenance, _FALLBACK_KEYS) or _truthy_flag( - provenance, _REFERENCE_KEYS - ): + if _truthy_flag(provenance, _FALLBACK_KEYS) or _truthy_flag(provenance, _REFERENCE_KEYS): errors.append(f"{label} RL-Kernel route reported a fallback/reference path") fallback_values = _values_for_keys(provenance, {"fallback"}) reference_values = _values_for_keys(provenance, _REFERENCE_KEYS) @@ -248,9 +244,7 @@ def _validate_rlkernel_record( {"actual_backend", "backend_id"}, STRICT_ROCM_BACKEND_ID, ): - errors.append( - f"{label} did not prove strict ROCm backend {STRICT_ROCM_BACKEND_ID!r}" - ) + errors.append(f"{label} did not prove strict ROCm backend {STRICT_ROCM_BACKEND_ID!r}") if not _has_exact_value( provenance, {"strict_core_id", "core_id"}, @@ -262,9 +256,7 @@ def _validate_rlkernel_record( {"strict_schedule", "schedule_id"}, STRICT_ROCM_SCHEDULE_ID, ): - errors.append( - f"{label} did not prove strict ROCm schedule {STRICT_ROCM_SCHEDULE_ID!r}" - ) + errors.append(f"{label} did not prove strict ROCm schedule {STRICT_ROCM_SCHEDULE_ID!r}") production_ready = _values_for_keys(provenance, {"production_ready"}) if not production_ready or not any(value is True for value in production_ready): errors.append(f"{label} strict ROCm provenance is not production-ready") @@ -336,16 +328,12 @@ def _validate_rlkernel_record( tp_world_size = max(int(value) for value in tp_values) except (TypeError, ValueError): tp_world_size = 0 - collective_values = _values_for_keys( - provenance, {"deterministic_all_reduce_backend"} - ) + collective_values = _values_for_keys(provenance, {"deterministic_all_reduce_backend"}) expected_collective = ( "none" if tp_world_size == 1 else ROCM_DETERMINISTIC_COLLECTIVE_BACKEND_ID ) if tp_world_size <= 0 or expected_collective not in collective_values: - errors.append( - f"{label} did not prove the deterministic ROCm O-projection collective" - ) + errors.append(f"{label} did not prove the deterministic ROCm O-projection collective") elif framework == "megatron": cp_values = _values_for_keys(provenance, {"cp_world_size"}) try: @@ -364,20 +352,14 @@ def _validate_rlkernel_record( expected_tp_collective = ( "none" if tp_world_size == 1 else ROCM_DETERMINISTIC_COLLECTIVE_BACKEND_ID ) - qkv_collectives = _values_for_keys( - provenance, {"tp_qkv_dgrad_collective"} - ) - output_collectives = _values_for_keys( - provenance, {"tp_output_projection_collective"} - ) + qkv_collectives = _values_for_keys(provenance, {"tp_qkv_dgrad_collective"}) + output_collectives = _values_for_keys(provenance, {"tp_output_projection_collective"}) if ( tp_world_size <= 0 or expected_tp_collective not in qkv_collectives or expected_tp_collective not in output_collectives ): - errors.append( - f"{label} did not prove deterministic ROCm TP projection collectives" - ) + errors.append(f"{label} did not prove deterministic ROCm TP projection collectives") def _validate_production_record( @@ -403,16 +385,12 @@ def _validate_strict_dense_record( errors: list[str], ) -> None: provenance = record.get("provenance") - expected_backend = ( - STRICT_FFN_BACKEND_ID if module == "ffn" else STRICT_LINEAR_LOGP_BACKEND_ID - ) + expected_backend = STRICT_FFN_BACKEND_ID if module == "ffn" else STRICT_LINEAR_LOGP_BACKEND_ID if record.get("case_id") != "R/R" or record.get("implementation") != "rl_kernel": errors.append(f"{label} did not execute the fixed R/R route") if record.get("backend_id") != expected_backend: errors.append(f"{label} reported backend {record.get('backend_id')!r}") - expected_mode = ( - "compiled_hip_graph" if framework == "vllm" and module == "ffn" else "eager" - ) + expected_mode = "compiled_hip_graph" if framework == "vllm" and module == "ffn" else "eager" if record.get("execution_mode", "eager") != expected_mode: errors.append(f"{label} did not execute in {expected_mode} mode") if int(record.get("call_count", 0)) <= 0: @@ -438,9 +416,7 @@ def _validate_strict_dense_record( if framework == "megatron" else "rocm_vocab_parallel_logp_from_local_logits_tp" ) - if not _has_exact_value( - provenance, {"logprob_kernel_backend"}, ROCM_LOGP_KERNEL_BACKEND_ID - ): + if not _has_exact_value(provenance, {"logprob_kernel_backend"}, ROCM_LOGP_KERNEL_BACKEND_ID): errors.append(f"{label} did not prove the ROCm WS2 logp kernel") if not _has_exact_value(provenance, {"strict_entrypoint"}, expected_entrypoint): errors.append(f"{label} did not prove strict entrypoint {expected_entrypoint!r}") @@ -613,9 +589,7 @@ def slice_response_mask_for_cp( mask = _tensor(loss_mask, label="loss_masks") if response_length < 0 or total_length < response_length: - raise ValueError( - f"invalid total/response lengths: {total_length}/{response_length}" - ) + raise ValueError(f"invalid total/response lengths: {total_length}/{response_length}") if mask.numel() != response_length: raise ValueError( "full loss mask length does not match response_length: " @@ -629,9 +603,7 @@ def slice_response_mask_for_cp( return mask prompt_length = total_length - response_length - chunk_size = (total_length + 2 * context_parallel_size - 1) // ( - 2 * context_parallel_size - ) + chunk_size = (total_length + 2 * context_parallel_size - 1) // (2 * context_parallel_size) chunks = ( ( context_parallel_rank * chunk_size, @@ -789,9 +761,7 @@ def _sidecar_samples( } if any(not isinstance(value, (list, tuple)) for value in required_lists.values()): missing = [ - name - for name, value in required_lists.items() - if not isinstance(value, (list, tuple)) + name for name, value in required_lists.items() if not isinstance(value, (list, tuple)) ] raise ValueError(f"mismatch sidecar lacks list fields: {', '.join(missing)}") assert isinstance(training_values, (list, tuple)) @@ -829,9 +799,7 @@ def _sidecar_samples( training = _tensor(training_values[index], label="train_log_probs") rollout = _tensor(rollout_values[index], label="rollout_log_probs") mask = _tensor(loss_masks[index], label="loss_masks").to(torch.bool) - total_length = _positive_int( - total_lengths[index], label="mismatch sidecar total_lengths" - ) + total_length = _positive_int(total_lengths[index], label="mismatch sidecar total_lengths") response_length = _positive_int( response_lengths[index], label="mismatch sidecar response_lengths" ) @@ -987,9 +955,7 @@ def compare_train_rollout_logps( "train_rollout_logprob_abs_diff": mean_abs_diff, "mismatch_kl": mismatch_kl, "mismatch_k3_kl": mismatch_k3_kl, - "sample_count": len( - {sample["logical_key"] for sample in unique.values()} - ), + "sample_count": len({sample["logical_key"] for sample in unique.values()}), "element_count": element_count, "artifacts": [str(path) for path in paths], } diff --git a/rl_engine/distributed/collectives.py b/rl_engine/distributed/collectives.py index 0b8ecb9a..686dfb39 100644 --- a/rl_engine/distributed/collectives.py +++ b/rl_engine/distributed/collectives.py @@ -32,9 +32,7 @@ def _deterministic_all_reduce_(input: torch.Tensor, collective_handle: int) -> N from rl_engine import _C if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_all_reduce_input( - collective_handle, input, input - ) + _C.deterministic_collective_rocm_ipc_all_reduce_input(collective_handle, input, input) else: _C.deterministic_collective_all_reduce_fused(collective_handle, input, input) @@ -65,9 +63,7 @@ def _deterministic_staging_reserve_(staging: torch.Tensor, collective_handle: in from rl_engine import _C if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_prepare_staged( - collective_handle, staging - ) + _C.deterministic_collective_rocm_ipc_prepare_staged(collective_handle, staging) else: _C.deterministic_collective_prepare_staged(collective_handle, staging) @@ -91,9 +87,7 @@ def _deterministic_staged_all_reduce( output = torch.empty_like(staging) if torch.version.hip is not None: - _C.deterministic_collective_rocm_ipc_all_reduce_staged( - collective_handle, staging, output - ) + _C.deterministic_collective_rocm_ipc_all_reduce_staged(collective_handle, staging, output) else: _C.deterministic_collective_all_reduce_staged(collective_handle, staging, output) return output diff --git a/rl_engine/distributed/rocm_collectives.py b/rl_engine/distributed/rocm_collectives.py index 13d2ce3d..2d0a983b 100644 --- a/rl_engine/distributed/rocm_collectives.py +++ b/rl_engine/distributed/rocm_collectives.py @@ -21,6 +21,7 @@ _ROCM_IPC_CONTROL_BYTES = 256 _REDUCTION_DTYPES = (torch.float32, torch.float16, torch.bfloat16) + class TorchDistributedDeterministicCollective: """Correctness-first collectives using AllGather as transport only. @@ -199,8 +200,7 @@ def all_gather_many( if not values: raise ValueError("all_gather_many requires at least one input") return tuple( - self.all_gather(value, validate_signature=validate_signature) - for value in values + self.all_gather(value, validate_signature=validate_signature) for value in values ) def reduce_scatter( @@ -654,9 +654,7 @@ def __init__( ) self._ipc_handle = 0 self._ipc_staging: torch.Tensor | None = None - self._direct_staging_views: dict[ - tuple[tuple[int, ...], torch.dtype], torch.Tensor - ] = {} + self._direct_staging_views: dict[tuple[tuple[int, ...], torch.dtype], torch.Tensor] = {} self._initialize_ipc_transport() @property @@ -730,9 +728,7 @@ def prepare_direct_staging_views( numel = 1 for dim in shape: if dim < 0: - raise ValueError( - f"direct-staging dimensions must be non-negative, got {shape}" - ) + raise ValueError(f"direct-staging dimensions must be non-negative, got {shape}") numel *= dim size_bytes = numel * element_size if size_bytes > min( diff --git a/rl_engine/integrations/framework_operators.py b/rl_engine/integrations/framework_operators.py index 08b361c6..694e31cd 100644 --- a/rl_engine/integrations/framework_operators.py +++ b/rl_engine/integrations/framework_operators.py @@ -405,9 +405,7 @@ def _rocm_paged_kv_max_tokens() -> int | None: try: limit = int(value) except ValueError as exc: - raise RuntimeError( - "RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS must be an integer" - ) from exc + raise RuntimeError("RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS must be an integer") from exc if limit <= 0: raise RuntimeError("RL_KERNEL_ROCM_PAGED_KV_MAX_TOKENS must be positive") return limit @@ -1043,9 +1041,7 @@ def warmup_rocm_decode(self, impl: Any, *, dtype: torch.dtype) -> None: "context_parallel_size": 1, }, ) - runtime = bound.bind_accelerator_runtime( - torch.empty((1,), device=device, dtype=dtype) - ) + runtime = bound.bind_accelerator_runtime(torch.empty((1,), device=device, dtype=dtype)) page_size = 16 core = getattr(runtime, "_core", None) if getattr(core, "attention_backend", "ck") == "triton": @@ -1127,9 +1123,7 @@ def provenance(self) -> Mapping[str, Any]: }, } - def _record_phase_provenance( - self, phase: str, provenance: dict[str, Any] - ) -> None: + def _record_phase_provenance(self, phase: str, provenance: dict[str, Any]) -> None: self._last_provenance = provenance self._phase_provenance[phase] = provenance @@ -1231,9 +1225,10 @@ def _rocm_dense_prefill( ) if starts[0] != 0 or starts[-1] != num_actual: return None - if any(end <= start or end - start != length for start, end, length in zip( - starts[:-1], starts[1:], lengths, strict=True - )): + if any( + end <= start or end - start != length + for start, end, length in zip(starts[:-1], starts[1:], lengths, strict=True) + ): return None output_heads = output.view(output.size(0), impl.num_heads, impl.head_size) @@ -1281,19 +1276,22 @@ def _rocm_dense_prefill( output_group.copy_(result_output) if num_actual < output.size(0): output[num_actual:].zero_() - self._record_phase_provenance("prefill", { - "framework_layout": "vllm_dense_qkv_prefill", - "materialization": "direct_dense_qkv_to_aiter_ck", - "tp_world_size": tp_world, - "runtime_platform": "rocm", - "triton_used": True, - "prefill_request_count": num_prefills, - "prefill_token_count": num_actual, - "core_launch_count": num_prefills, - "deterministic_projection": _strict_attention_projection_provenance("rocm"), - "deterministic_all_reduce_backend": "unbound" if tp_world > 1 else "none", - "direct_output_buffer": True, - }) + self._record_phase_provenance( + "prefill", + { + "framework_layout": "vllm_dense_qkv_prefill", + "materialization": "direct_dense_qkv_to_aiter_ck", + "tp_world_size": tp_world, + "runtime_platform": "rocm", + "triton_used": True, + "prefill_request_count": num_prefills, + "prefill_token_count": num_actual, + "core_launch_count": num_prefills, + "deterministic_projection": _strict_attention_projection_provenance("rocm"), + "deterministic_all_reduce_backend": "unbound" if tp_world > 1 else "none", + "direct_output_buffer": True, + }, + ) return output def _rocm_direct_paged_metadata( @@ -1366,14 +1364,10 @@ def _rocm_direct_paged_metadata( return None query_starts_source = query_start_loc seq_lens_source = self._metadata_tensor(attn_metadata, "seq_lens") - max_seq_len = int( - getattr(attn_metadata, "max_seq_len", block_table.size(1) * block_size) - ) + max_seq_len = int(getattr(attn_metadata, "max_seq_len", block_table.size(1) * block_size)) configured_kv_limit = _rocm_paged_kv_max_tokens() kernel_max_seqlen_k = ( - max_seq_len - if configured_kv_limit is None - else min(max_seq_len, configured_kv_limit) + max_seq_len if configured_kv_limit is None else min(max_seq_len, configured_kv_limit) ) page_count = min( block_table.size(1), @@ -1401,9 +1395,7 @@ def _rocm_direct_paged_metadata( self._rocm_paged_metadata_owners.add(owner_id) return self._rocm_paged_metadata_value, True - query_start_loc = query_starts_source.to( - device=block_table.device, dtype=torch.int32 - ) + query_start_loc = query_starts_source.to(device=block_table.device, dtype=torch.int32) if not query_start_loc.is_contiguous(): query_start_loc = query_start_loc.contiguous() seq_lens = seq_lens_source.to(device=block_table.device, dtype=torch.int32) @@ -1424,28 +1416,22 @@ def _rocm_direct_paged_metadata( ) else: tokens = torch.arange(num_actual, dtype=torch.int32, device=block_table.device) - seq_of_token = torch.searchsorted( - query_start_loc[1:], tokens, right=True - ).to(torch.int32) + seq_of_token = torch.searchsorted(query_start_loc[1:], tokens, right=True).to( + torch.int32 + ) elif mode == "decode": query_starts = query_start_loc[: sequence_count + 1] query_ends = query_starts[1:] - query_indices = torch.arange( - num_actual, dtype=torch.int32, device=block_table.device + query_indices = torch.arange(num_actual, dtype=torch.int32, device=block_table.device) + request_indices = torch.searchsorted(query_ends, query_indices, right=True).to( + dtype=torch.long ) - request_indices = torch.searchsorted( - query_ends, query_indices, right=True - ).to(dtype=torch.long) request_indices = request_indices.clamp_max(sequence_count - 1) request_query_ends = query_ends.index_select(0, request_indices) request_seq_lens = seq_lens.index_select(0, request_indices) active_queries = query_indices < query_starts[-1] - seqused_k = request_seq_lens - ( - request_query_ends - query_indices - ) + 1 - seqused_k = torch.where( - active_queries, seqused_k, torch.ones_like(seqused_k) - ) + seqused_k = request_seq_lens - (request_query_ends - query_indices) + 1 + seqused_k = torch.where(active_queries, seqused_k, torch.ones_like(seqused_k)) pages = block_table.index_select(0, request_indices)[:, :page_count] query_start_loc = torch.arange( num_actual + 1, dtype=torch.int32, device=block_table.device @@ -1481,9 +1467,7 @@ def _rocm_direct_paged_metadata( # Reuse the row's first live page so masked loads see initialized KV. safe_page = torch.where(active_rows, pages[:, 0], torch.zeros_like(seqused_k)) columns = torch.arange(page_count, dtype=torch.int32, device=pages.device) - live_columns = active_rows[:, None] & ( - columns[None, :] * block_size < seqused_k[:, None] - ) + live_columns = active_rows[:, None] & (columns[None, :] * block_size < seqused_k[:, None]) pages = torch.where(live_columns, pages, safe_page[:, None]) tile_pages = max(1, 128 // block_size) guard_columns = (-page_count) % tile_pages @@ -1499,11 +1483,14 @@ def _rocm_direct_paged_metadata( ) kv_indptr = self._rocm_kv_indptr_cache.get(indptr_key) if kv_indptr is None: - kv_indptr = torch.arange( - sequence_count + 1, - dtype=torch.int32, - device=block_table.device, - ) * page_count + kv_indptr = ( + torch.arange( + sequence_count + 1, + dtype=torch.int32, + device=block_table.device, + ) + * page_count + ) self._rocm_kv_indptr_cache[indptr_key] = kv_indptr value = { "mode": mode, @@ -1588,29 +1575,30 @@ def _rocm_direct_paged( if tp_world > 1: projection_collective_backend = "unbound" if self._projection_collective_backend is not None: - projection_collective_backend = ( - self._projection_collective_backend() or "unbound" - ) - self._record_phase_provenance(metadata["mode"], { - "framework_layout": "vllm_paged_kv", - "materialization": "direct_vllm_paged_kv_to_aiter_batch_prefill_ck", - "dense_kv_materialized": False, - "tp_world_size": tp_world, - "runtime_platform": "rocm", - "triton_used": True, - "attention_phase": metadata["mode"], - "sequence_count": metadata["sequence_count"], - "query_token_count": num_actual, - "max_seqlen_k": metadata["max_seqlen_k"], - "configured_kv_limit": metadata["configured_kv_limit"], - "launch_group_count": 1, - "metadata_source": "vllm_gpu_sequence_level", - "metadata_reused_across_layers": reused, - "deterministic_projection": _strict_attention_projection_provenance("rocm"), - "deterministic_all_reduce_backend": projection_collective_backend, - "direct_output_buffer": True, - "operator": operator_provenance, - }) + projection_collective_backend = self._projection_collective_backend() or "unbound" + self._record_phase_provenance( + metadata["mode"], + { + "framework_layout": "vllm_paged_kv", + "materialization": "direct_vllm_paged_kv_to_aiter_batch_prefill_ck", + "dense_kv_materialized": False, + "tp_world_size": tp_world, + "runtime_platform": "rocm", + "triton_used": True, + "attention_phase": metadata["mode"], + "sequence_count": metadata["sequence_count"], + "query_token_count": num_actual, + "max_seqlen_k": metadata["max_seqlen_k"], + "configured_kv_limit": metadata["configured_kv_limit"], + "launch_group_count": 1, + "metadata_source": "vllm_gpu_sequence_level", + "metadata_reused_across_layers": reused, + "deterministic_projection": _strict_attention_projection_provenance("rocm"), + "deterministic_all_reduce_backend": projection_collective_backend, + "direct_output_buffer": True, + "operator": operator_provenance, + }, + ) return output def _materialization_groups( @@ -1697,9 +1685,7 @@ def _materialization_groups( .to(dtype=torch.int32) .contiguous() ) - cu_seqlens_q = torch.arange( - num_actual + 1, dtype=torch.int32, device=query.device - ) + cu_seqlens_q = torch.arange(num_actual + 1, dtype=torch.int32, device=query.device) kv_indptr = cu_seqlens_q * page_count groups = [ { @@ -1800,9 +1786,7 @@ def __call__( candidate_factory = getattr(runtime, "new_page_bounds_epoch", None) if callable(candidate_factory): page_bounds_epoch_factory = candidate_factory - block_table = self._metadata_tensor( - attn_metadata, "block_table", "block_table_tensor" - ) + block_table = self._metadata_tensor(attn_metadata, "block_table", "block_table_tensor") key_cache, value_cache = _vllm_kv_cache_views( kv_cache, head_size=int(impl.head_size), @@ -1918,32 +1902,31 @@ def __call__( if runtime_platform == "rocm" and tp_world > 1: projection_collective_backend = "unbound" if self._projection_collective_backend is not None: - projection_collective_backend = ( - self._projection_collective_backend() or "unbound" - ) - phase = ( - "decode" - if int(getattr(attn_metadata, "num_decodes", 0)) > 0 - else "prefill" + projection_collective_backend = self._projection_collective_backend() or "unbound" + phase = "decode" if int(getattr(attn_metadata, "num_decodes", 0)) > 0 else "prefill" + self._record_phase_provenance( + phase, + { + "framework_layout": "vllm_paged_kv", + "materialization": ( + "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" + if runtime_platform == "rocm" + else "direct_paged_fa4" + ), + "dense_kv_materialized": False, + "tp_world_size": tp_world, + "tp_group_bound": tp_group is not None, + "runtime_platform": runtime_platform, + "triton_used": runtime_platform == "rocm", + "deterministic_projection": _strict_attention_projection_provenance( + runtime_platform + ), + "deterministic_all_reduce_backend": projection_collective_backend, + "direct_output_buffer": direct_output_buffer, + **metadata_summary, + "operator": last_operator_provenance, + }, ) - self._record_phase_provenance(phase, { - "framework_layout": "vllm_paged_kv", - "materialization": ( - "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" - if runtime_platform == "rocm" - else "direct_paged_fa4" - ), - "dense_kv_materialized": False, - "tp_world_size": tp_world, - "tp_group_bound": tp_group is not None, - "runtime_platform": runtime_platform, - "triton_used": runtime_platform == "rocm", - "deterministic_projection": _strict_attention_projection_provenance(runtime_platform), - "deterministic_all_reduce_backend": projection_collective_backend, - "direct_output_buffer": direct_output_buffer, - **metadata_summary, - "operator": last_operator_provenance, - }) return output diff --git a/rl_engine/integrations/linear_logp.py b/rl_engine/integrations/linear_logp.py index 667f989f..935a358d 100644 --- a/rl_engine/integrations/linear_logp.py +++ b/rl_engine/integrations/linear_logp.py @@ -450,8 +450,7 @@ def _rocm_contract( tp_rank=rank, tp_world_size=world, vocab_shard_bounds=tuple( - (index * local_vocab, (index + 1) * local_vocab) - for index in range(world) + (index * local_vocab, (index + 1) * local_vocab) for index in range(world) ), real_vocab_size=real_vocab_size, padded_vocab_size=global_vocab_size, diff --git a/rl_engine/integrations/megatron_runtime.py b/rl_engine/integrations/megatron_runtime.py index 596b0645..7ea94d64 100644 --- a/rl_engine/integrations/megatron_runtime.py +++ b/rl_engine/integrations/megatron_runtime.py @@ -93,7 +93,9 @@ def _save_megatron_layer_diagnostic( try: enabled_layers = {int(value.strip()) for value in requested_layers.split(",")} except ValueError as exc: - raise RuntimeError("RL_KERNEL_ALIGNMENT_LAYERS must be comma-separated integers") from exc + raise RuntimeError( + "RL_KERNEL_ALIGNMENT_LAYERS must be comma-separated integers" + ) from exc if layer not in enabled_layers: return if input_value.ndim == 3: @@ -136,8 +138,7 @@ def _save_megatron_layer_diagnostic( payload, output_dir / ( - f"megatron-pid{os.getpid()}-rank{rank:05d}-layer{layer:02d}-" - f"call{call_index:08d}.pt" + f"megatron-pid{os.getpid()}-rank{rank:05d}-layer{layer:02d}-" f"call{call_index:08d}.pt" ), ) @@ -252,8 +253,7 @@ def _strict_rocm_rope_positions( if cu_seqlens is None: raise RuntimeError("strict ROCm THD RoPE requires cu_seqlens") values = tuple( - int(value) - for value in cu_seqlens.detach().to(device="cpu", dtype=torch.int64).tolist() + int(value) for value in cu_seqlens.detach().to(device="cpu", dtype=torch.int64).tolist() ) if len(values) < 2 or values[0] != 0: raise RuntimeError("strict ROCm THD RoPE received invalid cu_seqlens") @@ -359,9 +359,7 @@ def strict_apply_rotary_pos_emb( def _install_torch_dist_object_compatibility() -> None: """Normalize the PyTorch DCP object shape expected by this Megatron revision.""" - strategy = importlib.import_module( - "megatron.core.dist_checkpointing.strategies.torch" - ) + strategy = importlib.import_module("megatron.core.dist_checkpointing.strategies.torch") original = strategy._replace_sharded_keys_with_state_dict_keys if getattr(original, "__rl_kernel_dcp_object_compatibility__", False): return @@ -793,6 +791,7 @@ def attention_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: callback_backend(reduce_from_tp), ) if hasattr(qkv, "layer_norm_weight"): + def te_qkv_forward(module: Any, input_value: torch.Tensor) -> Any: normalized = _fused_rms_norm_input(module, input_value, "linear_qkv") normalized = strict_tp_copy(module, core_attention, normalized) diff --git a/rl_engine/integrations/rocm_ablation.py b/rl_engine/integrations/rocm_ablation.py index 4a1fb040..1110b4eb 100644 --- a/rl_engine/integrations/rocm_ablation.py +++ b/rl_engine/integrations/rocm_ablation.py @@ -356,20 +356,12 @@ def _validate_route( backend_ids = tuple(sorted({str(record.get("backend_id", "")) for record in records})) provenance = [record.get("provenance", {}) for record in records] runtime_platforms = tuple( - sorted( - set().union( - *(_nested_strings(item, "runtime_platform") for item in provenance) - ) - ) + sorted(set().union(*(_nested_strings(item, "runtime_platform") for item in provenance))) ) actual_backends = tuple( - sorted( - set().union(*(_nested_strings(item, "actual_backend") for item in provenance)) - ) - ) - strict_core_ids = set().union( - *(_nested_strings(item, "strict_core_id") for item in provenance) + sorted(set().union(*(_nested_strings(item, "actual_backend") for item in provenance))) ) + strict_core_ids = set().union(*(_nested_strings(item, "strict_core_id") for item in provenance)) strict_schedules = set().union( *(_nested_strings(item, "strict_schedule") for item in provenance) ) @@ -379,9 +371,7 @@ def _validate_route( if expected is Implementation.PRODUCTION: native_backend = f"{framework}.production.attention" if any(backend != native_backend for backend in backend_ids): - errors.append( - f"{label} did not execute framework-native Attention: {backend_ids}" - ) + errors.append(f"{label} did not execute framework-native Attention: {backend_ids}") else: if any(not backend.startswith("rlkernel.attention.") for backend in backend_ids): errors.append(f"{label} did not execute the RL-Kernel Attention wrapper") @@ -394,8 +384,7 @@ def _validate_route( ) if STRICT_ROCM_ATTENTION_CORE not in strict_core_ids: errors.append( - f"{label} did not prove strict AITER/CK core " - f"{STRICT_ROCM_ATTENTION_CORE!r}" + f"{label} did not prove strict AITER/CK core " f"{STRICT_ROCM_ATTENTION_CORE!r}" ) if STRICT_ROCM_ATTENTION_SCHEDULE not in strict_schedules: errors.append( @@ -453,8 +442,7 @@ def run_rocm_attention_ablation( if occupied: joined = ", ".join(str(path) for path in occupied) raise FileExistsError( - "refusing to mix ROCm ablation evidence with existing case directories: " - + joined + "refusing to mix ROCm ablation evidence with existing case directories: " + joined ) validate_rocm_host() output_dir.mkdir(parents=True, exist_ok=True) diff --git a/rl_engine/integrations/runtime.py b/rl_engine/integrations/runtime.py index 083783e4..6b7d8f3b 100644 --- a/rl_engine/integrations/runtime.py +++ b/rl_engine/integrations/runtime.py @@ -178,9 +178,7 @@ def record_execution( ) if execution_provenance is None: raw_provenance = getattr(selected, "provenance", {}) - provenance = ( - dict(raw_provenance) if isinstance(raw_provenance, Mapping) else {} - ) + provenance = dict(raw_provenance) if isinstance(raw_provenance, Mapping) else {} else: provenance = dict(execution_provenance) actual_backend = _actual_backend(provenance) or backend_id @@ -225,9 +223,7 @@ def record_execution( flush=True, ) if implementation is Implementation.RL_KERNEL and fallback: - self.record_fallback( - normalized, f"operator provenance selected {actual_backend}" - ) + self.record_fallback(normalized, f"operator provenance selected {actual_backend}") raise RuntimeError( f"{self.framework} {normalized} strict RL-Kernel route reported fallback" ) @@ -242,8 +238,7 @@ def readback(self) -> dict[str, Any]: "fallbacks": list(self._fallbacks), "profile_calls": dict(self._profile_counts), "operators": { - module: readback.to_dict() - for module, readback in self._readbacks.items() + module: readback.to_dict() for module, readback in self._readbacks.items() }, } diff --git a/rl_engine/integrations/vime/linear_logp_provider.py b/rl_engine/integrations/vime/linear_logp_provider.py index 476ba5db..e9cb942e 100644 --- a/rl_engine/integrations/vime/linear_logp_provider.py +++ b/rl_engine/integrations/vime/linear_logp_provider.py @@ -290,7 +290,8 @@ def _provider_impl(request: Any, *, linear_logp: Any = None) -> LinearLogpResult and isinstance(request_logits, torch.Tensor) and request_logits.ndim == 2 and request_logits.dtype in (torch.bfloat16, torch.float16, torch.float32) - and request_logits.shape == ( + and request_logits.shape + == ( hidden.size(0), projection.weight.size(0), ) diff --git a/rl_engine/integrations/vllm_runtime.py b/rl_engine/integrations/vllm_runtime.py index 822ef41a..3b44ea41 100644 --- a/rl_engine/integrations/vllm_runtime.py +++ b/rl_engine/integrations/vllm_runtime.py @@ -834,6 +834,7 @@ def prepare_tables(instance: Any, device: torch.device) -> tuple[torch.Tensor, t device=device, theta=theta, ) + def register_table(name: str, table: torch.Tensor) -> None: if isinstance(instance, torch.nn.Module): buffers = instance._buffers @@ -915,9 +916,7 @@ def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: raise RuntimeError("vLLM splitting operators were not finalized before model init") if torch.version.hip is None: # Preserve the CUDA full-graph path introduced by PR 377. - splitting_ops[:] = [ - op for op in splitting_ops if op != DETERMINISTIC_ALL_REDUCE_OP - ] + splitting_ops[:] = [op for op in splitting_ops if op != DETERMINISTIC_ALL_REDUCE_OP] return from vllm import envs as vllm_envs @@ -933,16 +932,12 @@ def _configure_strict_ffn_compilation(vllm_config: Any | None = None) -> None: # vLLM's AOT key cannot see implementations behind torch custom ops. # Keep its normal config/code/compiler hashing under an RL-Kernel ABI # namespace so an older custom-op artifact cannot be replayed silently. - os.environ["VLLM_CACHE_ROOT"] = os.path.join( - cache_root, cache_namespace - ) + os.environ["VLLM_CACHE_ROOT"] = os.path.join(cache_root, cache_namespace) compilation.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE # ROCm IPC generations are allocated and consumed on device. Replayed # reductions therefore advance their generation instead of reusing the # capture-time payload, so these ops can remain in the full HIP graph. - splitting_ops[:] = [ - op for op in splitting_ops if op not in _ROCM_STATEFUL_GRAPH_SPLITTING_OPS - ] + splitting_ops[:] = [op for op in splitting_ops if op not in _ROCM_STATEFUL_GRAPH_SPLITTING_OPS] def _patch_rocm_weight_cache_refresh() -> None: @@ -991,7 +986,8 @@ def wrapped_init(instance: Any, *args: Any, **kwargs: Any) -> None: _handle, tp_world_size = operator.bind_packed_inference(instance) if not compiled_evidence_armed: execution_mode = ( - "compiled_hip_graph" if getattr(torch.version, "hip", None) is not None + "compiled_hip_graph" + if getattr(torch.version, "hip", None) is not None else "compiled_cuda_graph" ) register_packed_inference_observer( @@ -1097,9 +1093,7 @@ def deterministic_linear_apply( output_2d = rocm_linear_all_reduce( x_2d, layer.weight, - collective_handle=int( - getattr(layer, _STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT) - ), + collective_handle=int(getattr(layer, _STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT)), ) return output_2d.reshape(*x.shape[:-1], layer.weight.shape[0]) direct_output = None @@ -1223,9 +1217,7 @@ def bind_o_proj_collective(module: Any) -> None: raise RuntimeError("strict ROCm o_proj staging allocation failed") if register_rocm_linear_staging is None: raise RuntimeError("strict ROCm o_proj staging registry is unavailable") - compiled_slot = register_rocm_linear_staging( - int(collective._handle), staging - ) + compiled_slot = register_rocm_linear_staging(int(collective._handle), staging) setattr(module, _STRICT_O_PROJ_COMPILED_COLLECTIVE_SLOT, compiled_slot) setattr(module, _STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER, True) @@ -1253,9 +1245,7 @@ def strict_row_parallel_forward(instance: Any, input_: torch.Tensor) -> Any: output_parallel = instance.quant_method.apply(instance, input_parallel, bias_) if instance.reduce_results and instance.tp_size > 1: - if bool( - getattr(instance, _STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER, False) - ): + if bool(getattr(instance, _STRICT_O_PROJ_FUSED_ALL_REDUCE_MARKER, False)): output = output_parallel elif bool(getattr(instance, _STRICT_DIRECT_STAGING_MARKER, False)): output = deterministic_all_reduce_staged( @@ -1298,11 +1288,7 @@ def rotary_init_wrapped(instance: Any, *args: Any, **kwargs: Any) -> None: instance._forward_method = instance.forward_cuda cache = getattr(instance, "cos_sin_cache", None) prepare = getattr(instance, "_rl_kernel_prepare_strict_rocm_tables", None) - if ( - isinstance(cache, torch.Tensor) - and cache.is_cuda - and callable(prepare) - ): + if isinstance(cache, torch.Tensor) and cache.is_cuda and callable(prepare): prepare(cache.device) setattr(rotary_cls, _STRICT_ROTARY_INIT_MARKER, rotary_init) @@ -1444,9 +1430,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: if dtype in (torch.float16, torch.bfloat16): operator.warmup_rocm_decode(self, dtype=dtype) - def _split_kv_cache( - self, kv_cache: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: + def _split_kv_cache(self, kv_cache: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: if torch.version.hip is not None and operator is not None: if ( kv_cache.ndim != 4 @@ -1454,8 +1438,7 @@ def _split_kv_cache( or kv_cache.size(-1) != 2 * int(self.head_size) ): raise RuntimeError( - "RL-Kernel ROCm KV cache must use " - "[blocks, block, heads, 2 * head_size]" + "RL-Kernel ROCm KV cache must use " "[blocks, block, heads, 2 * head_size]" ) return kv_cache.split(int(self.head_size), dim=-1) return super()._split_kv_cache(kv_cache) diff --git a/rl_engine/kernels/ops/pytorch/ffn/ffn.py b/rl_engine/kernels/ops/pytorch/ffn/ffn.py index ed1c761b..13adaa6d 100644 --- a/rl_engine/kernels/ops/pytorch/ffn/ffn.py +++ b/rl_engine/kernels/ops/pytorch/ffn/ffn.py @@ -172,9 +172,7 @@ def _qwen3_ffn_packed_tp_inference_rocm( down_weight, direct_input, ) - _C.deterministic_collective_rocm_ipc_all_reduce_staged( - runtime_handle, direct_input, output - ) + _C.deterministic_collective_rocm_ipc_all_reduce_staged(runtime_handle, direct_input, output) return output.reshape(*input_shape[:-1], down_weight.shape[0]) else: # Profiling and uncaptured prefill can exceed the decode capture bound. @@ -199,9 +197,7 @@ def _qwen3_ffn_packed_tp_inference_rocm_fake( collective_handle: int, ) -> Tensor: del fused_gate_up_weight, collective_handle - return rmsnorm_output.new_empty( - (*rmsnorm_output.shape[:-1], down_weight.shape[0]) - ) + return rmsnorm_output.new_empty((*rmsnorm_output.shape[:-1], down_weight.shape[0])) def qwen3_ffn_packed_inference( @@ -249,8 +245,8 @@ def qwen3_ffn_packed_inference( rmsnorm_output.numel() // input_shape[-1], down_weight.shape[0], ) - direct_staging = None if collective is None else getattr( - collective, "direct_staging_view", None + direct_staging = ( + None if collective is None else getattr(collective, "direct_staging_view", None) ) direct_output = ( None @@ -835,20 +831,14 @@ def prepare_packed_inference( ) if staging is None: raise RuntimeError("packed ROCm rollout FFN staging allocation failed") - collective_handle = _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE.get( - runtime_handle, 0 - ) + collective_handle = _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE.get(runtime_handle, 0) if collective_handle == 0: # Keep the AOT graph identity stable across worker processes; # resolve its process-local C++ handle inside the custom op. collective_handle = len(_PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE) + 1 - _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE[runtime_handle] = ( - collective_handle - ) + _PACKED_INFERENCE_SLOT_BY_RUNTIME_HANDLE[runtime_handle] = collective_handle binding = _PACKED_INFERENCE_STAGING_BY_HANDLE.get(collective_handle) - stable_output = ( - torch.empty_like(staging) if binding is None else binding[2] - ) + stable_output = torch.empty_like(staging) if binding is None else binding[2] _PACKED_INFERENCE_STAGING_BY_HANDLE[collective_handle] = ( runtime_handle, staging, diff --git a/rl_engine/kernels/ops/rocm/attention/paged_gather.py b/rl_engine/kernels/ops/rocm/attention/paged_gather.py index e43eb9a5..ef7f354f 100644 --- a/rl_engine/kernels/ops/rocm/attention/paged_gather.py +++ b/rl_engine/kernels/ops/rocm/attention/paged_gather.py @@ -104,9 +104,7 @@ def fused_paged_kv_gather_bhsd( raise ValueError("paged gather output buffers have the wrong BHSD shape") if not k_out.is_contiguous() or not v_out.is_contiguous(): raise ValueError("paged gather output buffers must be contiguous") - if not ( - k_cache.device == v_cache.device == page_rows.device == k_out.device == v_out.device - ): + if not (k_cache.device == v_cache.device == page_rows.device == k_out.device == v_out.device): raise ValueError("paged gather tensors must share one device") if k_out.dtype != k_cache.dtype or v_out.dtype != v_cache.dtype: raise ValueError("paged gather output dtype must match the cache") diff --git a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py index 14cad16a..cddb3f7a 100644 --- a/rl_engine/kernels/ops/rocm/attention/strict_runtime.py +++ b/rl_engine/kernels/ops/rocm/attention/strict_runtime.py @@ -179,11 +179,7 @@ def _position_plan( key = (id(query_position_ids), id(key_position_ids), cp_world_size, causal) cached = self._position_plan_cache.get(key) - if ( - cached is not None - and cached[0] is query_position_ids - and cached[1] is key_position_ids - ): + if cached is not None and cached[0] is query_position_ids and cached[1] is key_position_ids: return cached[2:] if cp_world_size == 1: @@ -204,9 +200,7 @@ def _position_plan( key_sort = query_sort else: key_sort = torch.argsort(global_key_positions, dim=1) - query_positions_sorted = torch.gather( - global_query_positions, 1, query_sort - ) + query_positions_sorted = torch.gather(global_query_positions, 1, query_sort) key_positions_sorted = torch.gather(global_key_positions, 1, key_sort) _validate_global_positions( query_positions_sorted, @@ -399,9 +393,7 @@ def forward_paged_varlen_with_lse( self._require_rocm(q) direct = getattr(self._core, "forward_paged_varlen_with_lse", None) - if not callable(direct) or not bool( - getattr(self._core, "supports_paged_schedule", False) - ): + if not callable(direct) or not bool(getattr(self._core, "supports_paged_schedule", False)): raise RuntimeError("strict ROCm direct paged-varlen CK is unavailable") if q.ndim != 3: raise ValueError("packed paged Q must use [tokens, heads, head_dim]") @@ -549,9 +541,7 @@ def forward_paged_with_lse( bounds_ok = torch.all((page_table >= 0) & (page_table < k_cache.size(0))) torch._assert_async(bounds_ok, "page_table entries are outside the KV cache") if cu_seqlens_q is None: - cu_seqlens_q = torch.arange( - q.size(0) + 1, dtype=torch.int32, device=q.device - ) + cu_seqlens_q = torch.arange(q.size(0) + 1, dtype=torch.int32, device=q.device) if kv_indptr is None: kv_indptr = torch.arange( q.size(0) + 1, dtype=torch.int32, device=q.device @@ -796,9 +786,7 @@ def forward_paged_with_lse( if use_fused_gather else "logical_kv_gather_then_dense_core" ), - "paged_kernel": ( - "triton_fused_kv_gather_bhsd" if use_fused_gather else "none" - ), + "paged_kernel": ("triton_fused_kv_gather_bhsd" if use_fused_gather else "none"), "lse_returned": bool(return_lse), "launch_granularity": "one_batch_row_one_kv_group", "tp_degree_invariant": True, @@ -960,6 +948,7 @@ def _gather_paged_rows_by_page_count( rows = pages.size(0) page_size = k_cache.size(1) + def _gather(cache: torch.Tensor) -> torch.Tensor: selected = cache.index_select(0, flat_pages) flat = selected.reshape( diff --git a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py index 7f58862f..57831167 100644 --- a/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py +++ b/rl_engine/kernels/ops/rocm/loss/vocab_parallel_logp.py @@ -78,15 +78,11 @@ def _cached_active_mask( return cached, all_active -def _cached_shard_starts( - bounds: tuple[tuple[int, int], ...], device: torch.device -) -> torch.Tensor: +def _cached_shard_starts(bounds: tuple[tuple[int, int], ...], device: torch.device) -> torch.Tensor: key = (_device_key(device), bounds) cached = _SHARD_START_CACHE.get(key) if cached is None: - cached = torch.tensor( - [start for start, _ in bounds], dtype=torch.long, device=device - ) + cached = torch.tensor([start for start, _ in bounds], dtype=torch.long, device=device) _SHARD_START_CACHE[key] = cached if len(_SHARD_START_CACHE) > _METADATA_CACHE_LIMIT: _SHARD_START_CACHE.popitem(last=False) @@ -119,16 +115,11 @@ def _gather_target_logit_cached( if sharding.tp_world_size == 1: stacked = local_contrib.unsqueeze(0) else: - if ( - not torch.distributed.is_available() - or not torch.distributed.is_initialized() - ): + if not torch.distributed.is_available() or not torch.distributed.is_initialized(): raise LogprobContractError( "vocab-parallel logprob requires initialized torch.distributed" ) - gathered = [ - torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size) - ] + gathered = [torch.empty_like(local_contrib) for _ in range(sharding.tp_world_size)] torch.distributed.all_gather(gathered, local_contrib, group=tp_group) stacked = torch.stack(gathered, dim=0) @@ -176,9 +167,7 @@ class _RocmVocabParallelLogprobFunction(torch.autograd.Function): """ROCm tile statistics and backward with the shared WS2 merge contract.""" @staticmethod - def forward( - ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, all_active - ): + def forward(ctx, local_logits, target_1d, active_mask, contract, tp_group, tile, all_active): sharding = contract.sharding shard = local_logits.contiguous() local_tiles = sharding.local_vocab_size // tile @@ -190,9 +179,7 @@ def forward( sharding.real_vocab_size, local_tiles, ) - tile_counts = [ - (end - start) // tile for start, end in sharding.vocab_shard_bounds - ] + tile_counts = [(end - start) // tile for start, end in sharding.vocab_shard_bounds] m_all, s_all = _gather_tile_stats( local_m.contiguous(), local_s.contiguous(), @@ -205,9 +192,7 @@ def forward( if all_active else torch.where(active_mask, target_1d, torch.zeros_like(target_1d)) ) - target_logit = _gather_target_logit_cached( - shard, safe_target, contract, tp_group - ).float() + target_logit = _gather_target_logit_cached(shard, safe_target, contract, tp_group).float() lse = _merge_tile_partials(m_all, s_all) selected_logp = ( target_logit - lse @@ -245,15 +230,9 @@ def backward(ctx, grad_logp, grad_lse): ).contiguous() else: coef_logp = lse.new_zeros((rows,)) - target_local = torch.full( - (rows,), -1, dtype=torch.long, device=shard.device - ) + target_local = torch.full((rows,), -1, dtype=torch.long, device=shard.device) has_lse_grad = grad_lse is not None - coef_lse = ( - grad_lse.float().contiguous() - if has_lse_grad - else lse.new_zeros((rows,)) - ) + coef_lse = grad_lse.float().contiguous() if has_lse_grad else lse.new_zeros((rows,)) grad = _HipKernels.backward( shard, lse.contiguous(), @@ -281,18 +260,12 @@ def _apply_with_kernels( tile = _tile_size(contract, num_vocab_tiles) _validate_invocation(local_logits, target_ids, contract, tp_group) - target_1d = target_ids.reshape(-1).to( - device=local_logits.device, dtype=torch.long - ) + target_1d = target_ids.reshape(-1).to(device=local_logits.device, dtype=torch.long) active_mask, all_active = _cached_active_mask(contract, local_logits.device) if validate: - _validate_active_targets( - target_1d, active_mask, contract.sharding.real_vocab_size - ) + _validate_active_targets(target_1d, active_mask, contract.sharding.real_vocab_size) if contract.sharding.tp_world_size > 1: - _preflight_cross_rank_agreement( - contract, tp_group, num_vocab_tiles, True - ) + _preflight_cross_rank_agreement(contract, tp_group, num_vocab_tiles, True) selected_logp, lse = _RocmVocabParallelLogprobFunction.apply( local_logits, target_1d, active_mask, contract, tp_group, tile, all_active diff --git a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py index 14871417..48fb8af3 100644 --- a/rl_engine/kernels/ops/rocm/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/rocm/matmul/det_gemm.py @@ -262,9 +262,7 @@ def _det_gemm_linear_all_reduce_inference( direct_input, ) det_gemm_linear(a, weight, out=direct_input, inference_schedule=True) - _C.deterministic_collective_rocm_ipc_all_reduce_staged( - runtime_handle, direct_input, output - ) + _C.deterministic_collective_rocm_ipc_all_reduce_staged(runtime_handle, direct_input, output) return output else: # Profiling and uncaptured prefill can exceed the decode capture bound. diff --git a/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py b/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py index 360b4c5b..93571e51 100644 --- a/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py +++ b/rl_engine/kernels/ops/rocm/rotary_embedding/rope.py @@ -38,9 +38,7 @@ def _deterministic_rope_apply_rocm_fake( return torch.empty_like(x) -@torch.library.custom_op( - "rl_kernel::deterministic_rope_apply_token_major_rocm", mutates_args=() -) +@torch.library.custom_op("rl_kernel::deterministic_rope_apply_token_major_rocm", mutates_args=()) def _deterministic_rope_apply_token_major_rocm( x: Tensor, positions: Tensor, diff --git a/rl_engine/kernels/ops/triton/activation/swiglu.py b/rl_engine/kernels/ops/triton/activation/swiglu.py index ad11ecc8..0e8f08fd 100644 --- a/rl_engine/kernels/ops/triton/activation/swiglu.py +++ b/rl_engine/kernels/ops/triton/activation/swiglu.py @@ -63,9 +63,7 @@ def _swiglu_fwd_kernel(gate_ptr, up_ptr, y_ptr, n_elements, BLOCK: tl.constexpr) @triton.jit -def _swiglu_bwd_kernel( - dy_ptr, gate_ptr, silu_grad_ptr, d_up_ptr, n_elements, BLOCK: tl.constexpr -): +def _swiglu_bwd_kernel(dy_ptr, gate_ptr, silu_grad_ptr, d_up_ptr, n_elements, BLOCK: tl.constexpr): pid = tl.program_id(0) offs = pid * BLOCK + tl.arange(0, BLOCK) mask = offs < n_elements diff --git a/rl_engine/kernels/ops/triton/ffn/ffn.py b/rl_engine/kernels/ops/triton/ffn/ffn.py index f618031a..0f28465a 100644 --- a/rl_engine/kernels/ops/triton/ffn/ffn.py +++ b/rl_engine/kernels/ops/triton/ffn/ffn.py @@ -41,12 +41,8 @@ class Qwen3FFNForwardWeights: _source_versions: tuple[int | None, int | None, int | None] = field(repr=False) _packed_data_ptrs: tuple[int, int, int] = field(repr=False) _packed_versions: tuple[int, int, int] = field(repr=False) - _source_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field( - repr=False - ) - _source_strides: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field( - repr=False - ) + _source_shapes: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field(repr=False) + _source_strides: tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...]] = field(repr=False) def refresh_( self, @@ -125,9 +121,7 @@ def _validate_ffn_inputs( if tensor.dtype != torch.bfloat16: raise TypeError(f"{name} must have dtype bfloat16, got {tensor.dtype}.") if not tensor.is_cuda: - raise RuntimeError( - f"{name} must be on a CUDA/ROCm GPU device, got '{tensor.device}'." - ) + raise RuntimeError(f"{name} must be on a CUDA/ROCm GPU device, got '{tensor.device}'.") if tensor.device != rmsnorm_output.device: raise RuntimeError( f"all FFN inputs must be on {rmsnorm_output.device}, " @@ -160,9 +154,7 @@ def _validate_forward_weight_sources( if weight.dtype != torch.bfloat16: raise TypeError(f"{name} must have dtype bfloat16, got {weight.dtype}.") if not weight.is_cuda: - raise RuntimeError( - f"{name} must be on a CUDA/ROCm GPU device, got '{weight.device}'." - ) + raise RuntimeError(f"{name} must be on a CUDA/ROCm GPU device, got '{weight.device}'.") if tuple(up_weight.shape) != tuple(gate_weight.shape): raise ValueError( @@ -242,10 +234,7 @@ def refresh_qwen3_ffn_forward_weights( """ if not isinstance(forward_weights, Qwen3FFNForwardWeights): - raise TypeError( - "forward_weights must be created by " - "pack_qwen3_ffn_forward_weights." - ) + raise TypeError("forward_weights must be created by " "pack_qwen3_ffn_forward_weights.") sources = _validate_forward_weight_sources(gate_weight, up_weight, down_weight) if any( source is not original @@ -279,17 +268,13 @@ def refresh_qwen3_ffn_forward_weights( "repack and recapture CUDA Graphs." ) if not target.is_contiguous() or target.data_ptr() != expected_ptr: - raise RuntimeError( - f"packed {name} storage changed; repack and recapture CUDA Graphs." - ) + raise RuntimeError(f"packed {name} storage changed; repack and recapture CUDA Graphs.") with torch.inference_mode(False), torch.no_grad(): for target, source in zip(packed, sources, strict=True): target.copy_(source.t()) - forward_weights._source_versions = tuple( - _tracked_tensor_version(weight) for weight in sources - ) + forward_weights._source_versions = tuple(_tracked_tensor_version(weight) for weight in sources) forward_weights._packed_versions = tuple(int(weight._version) for weight in packed) return forward_weights @@ -301,10 +286,7 @@ def _validate_forward_weights( down_weight: Tensor, ) -> None: if not isinstance(forward_weights, Qwen3FFNForwardWeights): - raise TypeError( - "forward_weights must be created by " - "pack_qwen3_ffn_forward_weights." - ) + raise TypeError("forward_weights must be created by " "pack_qwen3_ffn_forward_weights.") sources = (gate_weight, up_weight, down_weight) names = ("gate_weight", "up_weight", "down_weight") @@ -336,9 +318,7 @@ def _validate_forward_weights( if not isinstance(weight, Tensor): raise TypeError(f"packed {name} must be a torch.Tensor.") if tuple(weight.shape) != shape: - raise ValueError( - f"packed {name} must have shape {shape}, got {tuple(weight.shape)}." - ) + raise ValueError(f"packed {name} must have shape {shape}, got {tuple(weight.shape)}.") if weight.dtype != torch.bfloat16: raise TypeError(f"packed {name} must have dtype bfloat16, got {weight.dtype}.") if weight.device != gate_weight.device: @@ -556,10 +536,8 @@ def backward(ctx: Any, grad_output: Tensor) -> tuple[Any, ...]: grad_rmsnorm_from_gate = _gemm(grad_gate, gate_weight) grad_rmsnorm_from_up = _gemm(grad_up, up_weight) if ctx.sequence_parallel: - grad_rmsnorm_from_gate, grad_rmsnorm_from_up = ( - tp_collective.reduce_scatter_many( - (grad_rmsnorm_from_gate, grad_rmsnorm_from_up) - ) + grad_rmsnorm_from_gate, grad_rmsnorm_from_up = tp_collective.reduce_scatter_many( + (grad_rmsnorm_from_gate, grad_rmsnorm_from_up) ) elif tp_collective is not None: grad_rmsnorm_from_gate = _all_reduce_inplace( @@ -615,9 +593,7 @@ def qwen3_ffn( down_weight, ) if not isinstance(sequence_parallel, bool): - raise TypeError( - f"sequence_parallel must be a bool, got {type(sequence_parallel)!r}." - ) + raise TypeError(f"sequence_parallel must be a bool, got {type(sequence_parallel)!r}.") return _TritonDeterministicFFNFunction.apply( rmsnorm_output, gate_weight, diff --git a/rl_engine/kernels/ops/triton/matmul/det_gemm.py b/rl_engine/kernels/ops/triton/matmul/det_gemm.py index 1fd370d6..260771cc 100644 --- a/rl_engine/kernels/ops/triton/matmul/det_gemm.py +++ b/rl_engine/kernels/ops/triton/matmul/det_gemm.py @@ -137,11 +137,7 @@ def _gfx942_qwen_tree_leaf_config( logical_shape, _DEFAULT_TREE_LEAF_CONFIG, ) - if ( - transpose_output - and preserve_a_strides - and (m_size, n_size) in _QWEN_WGRAD_OUTPUT_SHAPES - ): + if transpose_output and preserve_a_strides and (m_size, n_size) in _QWEN_WGRAD_OUTPUT_SHAPES: return _GFX942_QWEN_WGRAD_LEAF_CONFIGS.get( k_size, _DEFAULT_TREE_LEAF_CONFIG, @@ -204,13 +200,16 @@ class _DeviceTreePlan: rocm_fused_reduction_pairs: tuple[ tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ... ] - rocm_leaf_reduction: tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - ] | None + rocm_leaf_reduction: ( + tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + ] + | None + ) rocm_fused_reduction_pairs_after_leaf: tuple[ tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor], ... ] @@ -247,8 +246,7 @@ def visit(begin: int, end: int) -> tuple[int, int]: root, max_height = visit(0, k_size) reduction_levels = tuple( - tuple(reductions_by_height.get(height, ())) - for height in range(1, max_height + 1) + tuple(reductions_by_height.get(height, ())) for height in range(1, max_height + 1) ) return _TreePlan( leaf_starts=tuple(leaf_starts), @@ -281,11 +279,10 @@ def indices(values: tuple[int, ...]) -> torch.Tensor: leaf_reduction = None fused_pairs_after_leaf = [] if _device_arch(device_index) == "gfx942": + def build_fused_pairs(start_level: int): result = [] - for level_index in range( - start_level, len(host.reduction_levels) - 1, 2 - ): + for level_index in range(start_level, len(host.reduction_levels) - 1, 2): first = { output: (lower, upper) for lower, upper, output in host.reduction_levels[level_index] @@ -410,9 +407,7 @@ def _det_gemm_tree_leaf_kernel( ).to(tl.float32) acc += a[:, None] * b[None, :] output_offsets = ( - leaf_node * (M * N) - + offs_m[:, None].to(tl.int64) * N - + offs_n[None, :].to(tl.int64) + leaf_node * (M * N) + offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64) ) output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) tl.store( @@ -498,14 +493,9 @@ def _det_gemm_tree_leaf_reduce_rocm_kernel( upper_acc += a[:, None] * b[None, :] # Match both leaf BF16 stores followed by the original first-level FP32 # add and BF16 store. Only the intermediate global-memory trip is gone. - result = lower.to(tl.float32) + upper_acc.to( - workspace_ptr.dtype.element_ty - ).to(tl.float32) + result = lower.to(tl.float32) + upper_acc.to(workspace_ptr.dtype.element_ty).to(tl.float32) output_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) - offsets = ( - offs_m[:, None].to(tl.int64) * N - + offs_n[None, :].to(tl.int64) - ) + offsets = offs_m[:, None].to(tl.int64) * N + offs_n[None, :].to(tl.int64) if WRITE_OUTPUT: tl.store( output_ptr + offsets, @@ -615,18 +605,18 @@ def _det_gemm_tree_reduce_two_levels_rocm_kernel( node1 = tl.load(node1_ptr + operation).to(tl.int64) node2 = tl.load(node2_ptr + operation).to(tl.int64) node3 = tl.load(node3_ptr + operation).to(tl.int64) - value0 = tl.load( - workspace_ptr + node0 * elements + offsets, mask=mask, other=0.0 - ).to(tl.float32) - value1 = tl.load( - workspace_ptr + node1 * elements + offsets, mask=mask, other=0.0 - ).to(tl.float32) - value2 = tl.load( - workspace_ptr + node2 * elements + offsets, mask=mask, other=0.0 - ).to(tl.float32) - value3 = tl.load( - workspace_ptr + node3 * elements + offsets, mask=mask, other=0.0 - ).to(tl.float32) + value0 = tl.load(workspace_ptr + node0 * elements + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + value1 = tl.load(workspace_ptr + node1 * elements + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + value2 = tl.load(workspace_ptr + node2 * elements + offsets, mask=mask, other=0.0).to( + tl.float32 + ) + value3 = tl.load(workspace_ptr + node3 * elements + offsets, mask=mask, other=0.0).to( + tl.float32 + ) # Match the two original kernel boundaries exactly: each first-level # FP32 add is rounded to BF16 before the second-level FP32 add. lower = (value0 + value1).to(workspace_ptr.dtype.element_ty).to(tl.float32) @@ -674,10 +664,7 @@ def _copy_tree_root_transposed_kernel( elements = M * N mask = (offsets_m[:, None] < M) & (offsets_n[None, :] < N) values = tl.load( - workspace_ptr - + root * elements - + offsets_m[:, None] * N - + offsets_n[None, :], + workspace_ptr + root * elements + offsets_m[:, None] * N + offsets_n[None, :], mask=mask, ) # Store the already-rounded BF16 root directly in [N, M] layout. @@ -836,9 +823,7 @@ def _triton_tree_gemm( if out.dtype != torch.bfloat16: raise TypeError(f"Triton tree GEMM output must be BF16, got {out.dtype}") if out.device != a.device: - raise RuntimeError( - f"Triton tree GEMM output must be on {a.device}, got {out.device}" - ) + raise RuntimeError(f"Triton tree GEMM output must be on {a.device}, got {out.device}") if not out.is_contiguous(): raise ValueError("Triton tree GEMM output buffer must be contiguous") if out.requires_grad: @@ -851,9 +836,7 @@ def _triton_tree_gemm( ) device_index = a.device.index if a.device.index is not None else torch.cuda.current_device() is_gfx942 = _device_arch(device_index) == "gfx942" - use_inference_schedule = ( - torch.is_inference_mode_enabled() or inference_schedule - ) + use_inference_schedule = torch.is_inference_mode_enabled() or inference_schedule leaf_config = _tree_leaf_config( a.device, m_size, @@ -869,10 +852,7 @@ def _triton_tree_gemm( and not transpose_output and not preserve_a_strides ): - leaf_config = ( - _gfx942_qwen_tp4_decode_leaf_config(m_size, k_size, n_size) - or leaf_config - ) + leaf_config = _gfx942_qwen_tp4_decode_leaf_config(m_size, k_size, n_size) or leaf_config tiles_m = triton.cdiv(m_size, leaf_config.block_m) tiles_n = triton.cdiv(n_size, leaf_config.block_n) leaf_grid = ( @@ -938,9 +918,7 @@ def _triton_tree_gemm( reduction_block = 256 if fuse_leaf_reduction: blocks = triton.cdiv(m_size * n_size, reduction_block) - for pair_index, nodes in enumerate( - plan.rocm_fused_reduction_pairs_after_leaf - ): + for pair_index, nodes in enumerate(plan.rocm_fused_reduction_pairs_after_leaf): second_level = pair_index * 2 + 2 operations = plan.host.reduction_levels[second_level] write_final_output = second_level == len(plan.host.reduction_levels) - 1 @@ -1003,9 +981,7 @@ def _triton_tree_gemm( direct_root_output and level_index == len(plan.host.reduction_levels) - 1 ) if write_final_output and len(operations) != 1: - raise RuntimeError( - "the final deterministic GEMM tree level must contain one root" - ) + raise RuntimeError("the final deterministic GEMM tree level must contain one root") if write_final_output: grid = (triton.cdiv(m_size * n_size, reduction_block),) # direct_root_output is true only for gfx942; CUDA retains its diff --git a/tests/distributed/test_qwen_ffn_topology.py b/tests/distributed/test_qwen_ffn_topology.py index 5cb72a09..8fe750fe 100644 --- a/tests/distributed/test_qwen_ffn_topology.py +++ b/tests/distributed/test_qwen_ffn_topology.py @@ -27,6 +27,7 @@ pack_qwen3_ffn_forward_weights, qwen3_ffn, ) + _IS_ROCM = getattr(torch.version, "hip", None) is not None _EXTERNAL_WORLD_SIZE = int(os.environ.get("WORLD_SIZE", "1")) @@ -251,12 +252,12 @@ def _run_topology( expected_output = inference_reference[token_start:token_end] assert torch.equal(inference, training.detach()), f"{name}: train/infer mismatch" assert torch.equal(inference, expected_output), f"{name}: inference mismatch vs TP=1" - assert torch.equal(training.detach(), training_reference.detach()[token_start:token_end]), ( - f"{name}: training forward mismatch vs TP=1" - ) - assert torch.equal(inputs[0].grad, reference_inputs[0].grad[token_start:token_end]), ( - f"{name}: hidden grad mismatch vs TP=1" - ) + assert torch.equal( + training.detach(), training_reference.detach()[token_start:token_end] + ), f"{name}: training forward mismatch vs TP=1" + assert torch.equal( + inputs[0].grad, reference_inputs[0].grad[token_start:token_end] + ), f"{name}: hidden grad mismatch vs TP=1" expected_weight_grads = ( (inputs[1].grad, reference_inputs[1].grad[feature_start:feature_end], "gate"), diff --git a/tests/test_det_gemm.py b/tests/test_det_gemm.py index 0c4408c7..508ab2d6 100644 --- a/tests/test_det_gemm.py +++ b/tests/test_det_gemm.py @@ -44,9 +44,11 @@ HAS_SUPPORTED_GPU = torch.cuda.is_available() and ( IS_ROCM or torch.cuda.get_device_capability()[0] >= 8 ) -IS_GFX942 = IS_ROCM and torch.cuda.is_available() and str( - getattr(torch.cuda.get_device_properties(0), "gcnArchName", "") -).startswith("gfx942") +IS_GFX942 = ( + IS_ROCM + and torch.cuda.is_available() + and str(getattr(torch.cuda.get_device_properties(0), "gcnArchName", "")).startswith("gfx942") +) pytestmark = pytest.mark.skipif( not HAS_SUPPORTED_GPU, @@ -483,9 +485,9 @@ def test_target_shapes_invariance(name, gemm, shape): row = _rand(1, K) big = _rand(64, K) big[0] = row[0] - assert torch.equal(gemm(row, b)[0], gemm(big, b)[0]), ( - f"{name}: batch-invariance broken at shape {shape}" - ) + assert torch.equal( + gemm(row, b)[0], gemm(big, b)[0] + ), f"{name}: batch-invariance broken at shape {shape}" @pytest.mark.skipif(not _HAS_TRITON, reason="Triton is unavailable") @@ -654,10 +656,14 @@ def test_triton_wgrad_reads_positive_stride_transpose_view_raw_bytes(): assert not activation_t.is_contiguous() assert all(stride > 0 for stride in activation_t.stride()) - legacy = _triton_tree_gemm( - activation_t.contiguous(), - grad_output, - ).t().contiguous() + legacy = ( + _triton_tree_gemm( + activation_t.contiguous(), + grad_output, + ) + .t() + .contiguous() + ) output_buffer = torch.empty( (output_size, input_size), dtype=torch.bfloat16, diff --git a/tests/test_framework_runtime_adapters.py b/tests/test_framework_runtime_adapters.py index 2fa2900a..a34537ca 100644 --- a/tests/test_framework_runtime_adapters.py +++ b/tests/test_framework_runtime_adapters.py @@ -29,18 +29,18 @@ _packed_local_sequence_layout, _vllm_kv_cache_views, ) -from rl_engine.integrations.megatron_runtime import ( - _deterministic_reduce_from_tensor_model_parallel_region, - _install_torch_dist_object_compatibility, - _patch_strict_attention_projections, - install_megatron_integration, -) +from rl_engine.integrations.megatron_runtime import ( + _deterministic_reduce_from_tensor_model_parallel_region, + _install_torch_dist_object_compatibility, + _patch_strict_attention_projections, + install_megatron_integration, +) from rl_engine.integrations.runtime import FrameworkOperatorIntegration from rl_engine.integrations.state import clear_active_integration -from rl_engine.integrations.vllm_runtime import ( - _patch_qwen3_strict_model, - _patch_strict_rocm_rotary_embedding, - _register_attention_backend, +from rl_engine.integrations.vllm_runtime import ( + _patch_qwen3_strict_model, + _patch_strict_rocm_rotary_embedding, + _register_attention_backend, configure_vllm_environment, ) from rl_engine.kernels.attention_contract import ( @@ -53,44 +53,39 @@ ReductionSpec, ShardingSpec, ) -from rl_engine.kernels.ops.cuda.attention.strict_runtime import ( - StrictCUDAAttentionRuntime, -) - - -def test_torch_dist_object_compatibility_deserializes_scalar_bytes_io(monkeypatch): - strategy_name = "megatron.core.dist_checkpointing.strategies.torch" - strategy = ModuleType(strategy_name) - calls = [] - - def replace(state_dict, flat_mapping, rename_mapping): - calls.append((state_dict, flat_mapping, rename_mapping)) - return state_dict - - strategy._replace_sharded_keys_with_state_dict_keys = replace - monkeypatch.setitem(sys.modules, strategy_name, strategy) - - _install_torch_dist_object_compatibility() - installed = strategy._replace_sharded_keys_with_state_dict_keys - _install_torch_dist_object_compatibility() - payload = __import__("io").BytesIO() - torch.save([{"recipe": "checkpoint object"}], payload) - - assert strategy._replace_sharded_keys_with_state_dict_keys is installed - assert installed({"state": payload}, "flat", "rename") == { - "state": [{"recipe": "checkpoint object"}] - } - assert calls == [ - ({"state": [{"recipe": "checkpoint object"}]}, "flat", "rename") - ] - - -def test_framework_adapters_do_not_construct_registered_kernels_directly(): +from rl_engine.kernels.ops.cuda.attention.strict_runtime import ( + StrictCUDAAttentionRuntime, +) + + +def test_torch_dist_object_compatibility_deserializes_scalar_bytes_io(monkeypatch): + strategy_name = "megatron.core.dist_checkpointing.strategies.torch" + strategy = ModuleType(strategy_name) + calls = [] + + def replace(state_dict, flat_mapping, rename_mapping): + calls.append((state_dict, flat_mapping, rename_mapping)) + return state_dict + + strategy._replace_sharded_keys_with_state_dict_keys = replace + monkeypatch.setitem(sys.modules, strategy_name, strategy) + + _install_torch_dist_object_compatibility() + installed = strategy._replace_sharded_keys_with_state_dict_keys + _install_torch_dist_object_compatibility() + payload = __import__("io").BytesIO() + torch.save([{"recipe": "checkpoint object"}], payload) + + assert strategy._replace_sharded_keys_with_state_dict_keys is installed + assert installed({"state": payload}, "flat", "rename") == { + "state": [{"recipe": "checkpoint object"}] + } + assert calls == [({"state": [{"recipe": "checkpoint object"}]}, "flat", "rename")] + + +def test_framework_adapters_do_not_construct_registered_kernels_directly(): source_path = ( - Path(__file__).parents[1] - / "rl_engine" - / "integrations" - / "framework_operators.py" + Path(__file__).parents[1] / "rl_engine" / "integrations" / "framework_operators.py" ) tree = ast.parse(source_path.read_text(encoding="utf-8")) forbidden = { @@ -164,8 +159,8 @@ def test_vllm_rlkernel_attention_overrides_selected_flash_attn_backend( configure_vllm_environment(plan) - expected = "ROCM_AITER_FA" if torch.version.hip is not None else "FLASH_ATTN" - assert os.environ["VLLM_ATTENTION_BACKEND"] == expected + expected = "ROCM_AITER_FA" if torch.version.hip is not None else "FLASH_ATTN" + assert os.environ["VLLM_ATTENTION_BACKEND"] == expected def test_vllm_rocm_attention_selects_aiter_metadata_backend(monkeypatch): @@ -331,12 +326,12 @@ def get(self, tensor, *, topology): get_tensor_model_parallel_rank=lambda: 0, get_context_parallel_group=lambda: "cp-group", ) - monkeypatch.setattr(framework_operators, "_megatron_parallel_state", lambda: parallel_state) - monkeypatch.setattr( - framework_operators, - "_require_attention_accelerator", - lambda tensor: "cuda", - ) + monkeypatch.setattr(framework_operators, "_megatron_parallel_state", lambda: parallel_state) + monkeypatch.setattr( + framework_operators, + "_require_attention_accelerator", + lambda tensor: "cuda", + ) packed = SimpleNamespace( qkv_format="thd", cu_seqlens_q=torch.tensor([0, 8, 16], dtype=torch.int32), @@ -364,9 +359,9 @@ def get(self, tensor, *, topology): [0, 1, 6, 7], [0, 1, 6, 7], ] - - -def test_megatron_attention_binds_rocm_core_and_schedule(monkeypatch): + + +def test_megatron_attention_binds_rocm_core_and_schedule(monkeypatch): calls = [] class Operator: @@ -420,7 +415,7 @@ def get(self, tensor, *, topology): assert adapter.provenance["execution"]["runtime_platform"] == "rocm" -def test_vllm_attention_routes_paged_cache_to_rocm_runtime(monkeypatch): +def test_vllm_attention_routes_paged_cache_to_rocm_runtime(monkeypatch): runtime_calls = [] class Runtime: @@ -466,112 +461,112 @@ def get(self, tensor, *, topology): output = adapter(impl, object(), query, query, query, kv_cache, metadata) - assert output.shape == (1, 16) - assert len(runtime_calls) == 1 - assert runtime_calls[0][0].shape == (1, 2, 1, 8) - assert torch.equal(runtime_calls[0][3]["cu_seqlens_q"], torch.tensor([0, 1])) - assert torch.equal(runtime_calls[0][3]["kv_indptr"], torch.tensor([0, 1])) - assert adapter.provenance["execution"]["runtime_platform"] == "rocm" - assert ( - adapter.provenance["execution"]["materialization"] - == "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" - ) - assert adapter.provenance["execution"]["dense_kv_materialized"] is False - - -def test_vllm_rocm_metadata_stays_on_device_and_is_reused_across_layers(monkeypatch): - original_tolist = torch.Tensor.tolist - tolist_calls = 0 - - def counting_tolist(tensor): - nonlocal tolist_calls - tolist_calls += 1 - return original_tolist(tensor) - - monkeypatch.setattr(torch.Tensor, "tolist", counting_tolist) - query = torch.zeros(2, 2, 8, dtype=torch.bfloat16) - block_table = torch.tensor([[0], [1]], dtype=torch.int32) - metadata = SimpleNamespace( - query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), - seq_lens=torch.tensor([3, 5], dtype=torch.int32), - max_seq_len=5, - ) - adapter = VllmAttentionOperator() - first_layer = object() - second_layer = object() - - first, _ = adapter._materialization_groups( - metadata, - query=query, - block_table=block_table, - block_size=8, - num_actual=2, - cache_owner=first_layer, - ) - second, summary = adapter._materialization_groups( - metadata, - query=query, - block_table=block_table, - block_size=8, - num_actual=2, - cache_owner=second_layer, - ) - - assert first is second - assert torch.equal(first[0]["seqused_k"], torch.tensor([3, 5], dtype=torch.int32)) - assert torch.equal(first[0]["cu_seqlens_q"], torch.tensor([0, 1, 2], dtype=torch.int32)) - assert torch.equal(first[0]["kv_indptr"], torch.tensor([0, 1, 2], dtype=torch.int32)) - assert tolist_calls == 0 - assert summary["metadata_reused_across_layers"] is True - - -def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): - cache = torch.arange(2 * 3 * 4 * 10).reshape(2, 3, 4, 10) + assert output.shape == (1, 16) + assert len(runtime_calls) == 1 + assert runtime_calls[0][0].shape == (1, 2, 1, 8) + assert torch.equal(runtime_calls[0][3]["cu_seqlens_q"], torch.tensor([0, 1])) + assert torch.equal(runtime_calls[0][3]["kv_indptr"], torch.tensor([0, 1])) + assert adapter.provenance["execution"]["runtime_platform"] == "rocm" + assert ( + adapter.provenance["execution"]["materialization"] + == "direct_vllm_paged_kv_to_aiter_batch_prefill_ck" + ) + assert adapter.provenance["execution"]["dense_kv_materialized"] is False + + +def test_vllm_rocm_metadata_stays_on_device_and_is_reused_across_layers(monkeypatch): + original_tolist = torch.Tensor.tolist + tolist_calls = 0 + + def counting_tolist(tensor): + nonlocal tolist_calls + tolist_calls += 1 + return original_tolist(tensor) + + monkeypatch.setattr(torch.Tensor, "tolist", counting_tolist) + query = torch.zeros(2, 2, 8, dtype=torch.bfloat16) + block_table = torch.tensor([[0], [1]], dtype=torch.int32) + metadata = SimpleNamespace( + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + seq_lens=torch.tensor([3, 5], dtype=torch.int32), + max_seq_len=5, + ) + adapter = VllmAttentionOperator() + first_layer = object() + second_layer = object() + + first, _ = adapter._materialization_groups( + metadata, + query=query, + block_table=block_table, + block_size=8, + num_actual=2, + cache_owner=first_layer, + ) + second, summary = adapter._materialization_groups( + metadata, + query=query, + block_table=block_table, + block_size=8, + num_actual=2, + cache_owner=second_layer, + ) + + assert first is second + assert torch.equal(first[0]["seqused_k"], torch.tensor([3, 5], dtype=torch.int32)) + assert torch.equal(first[0]["cu_seqlens_q"], torch.tensor([0, 1, 2], dtype=torch.int32)) + assert torch.equal(first[0]["kv_indptr"], torch.tensor([0, 1, 2], dtype=torch.int32)) + assert tolist_calls == 0 + assert summary["metadata_reused_across_layers"] is True + + +def test_vllm_current_flash_attention_kv_cache_layout_is_materialized(): + cache = torch.arange(2 * 3 * 4 * 10).reshape(2, 3, 4, 10) key, value = _vllm_kv_cache_views(cache, head_size=5) assert key.shape == (2, 4, 3, 5) assert value.shape == (2, 4, 3, 5) assert torch.equal(key, cache.transpose(1, 2)[..., :5]) - assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) - - -def test_vllm_rocm_native_kv_cache_with_two_heads_is_not_treated_as_pair_axis(): - cache = torch.arange(3 * 2 * 4 * 10).reshape(3, 2, 4, 10) - - key, value = _vllm_kv_cache_views( - cache, - head_size=5, - num_kv_heads=2, - platform="rocm", - ) - - assert key.shape == (3, 4, 2, 5) - assert value.shape == (3, 4, 2, 5) - assert torch.equal(key, cache.transpose(1, 2)[..., :5]) - assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) - - -def test_vllm_rocm_rlkernel_token_major_kv_cache_is_zero_copy(): - cache = torch.arange(3 * 4 * 2 * 10).reshape(3, 4, 2, 10) - - key, value = _vllm_kv_cache_views( - cache, - head_size=5, - num_kv_heads=2, - platform="rocm", - ) - - assert key.shape == (3, 4, 2, 5) - assert value.shape == (3, 4, 2, 5) - assert key.data_ptr() == cache.data_ptr() - assert value.untyped_storage().data_ptr() == cache.untyped_storage().data_ptr() - assert torch.equal(key, cache[..., :5]) - assert torch.equal(value, cache[..., 5:]) - assert key.stride(1) >= key.size(2) * key.stride(2) - - -def test_vllm_rocm_kv_cache_pair_axis_is_materialized(): + assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) + + +def test_vllm_rocm_native_kv_cache_with_two_heads_is_not_treated_as_pair_axis(): + cache = torch.arange(3 * 2 * 4 * 10).reshape(3, 2, 4, 10) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=2, + platform="rocm", + ) + + assert key.shape == (3, 4, 2, 5) + assert value.shape == (3, 4, 2, 5) + assert torch.equal(key, cache.transpose(1, 2)[..., :5]) + assert torch.equal(value, cache.transpose(1, 2)[..., 5:]) + + +def test_vllm_rocm_rlkernel_token_major_kv_cache_is_zero_copy(): + cache = torch.arange(3 * 4 * 2 * 10).reshape(3, 4, 2, 10) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=2, + platform="rocm", + ) + + assert key.shape == (3, 4, 2, 5) + assert value.shape == (3, 4, 2, 5) + assert key.data_ptr() == cache.data_ptr() + assert value.untyped_storage().data_ptr() == cache.untyped_storage().data_ptr() + assert torch.equal(key, cache[..., :5]) + assert torch.equal(value, cache[..., 5:]) + assert key.stride(1) >= key.size(2) * key.stride(2) + + +def test_vllm_rocm_kv_cache_pair_axis_is_materialized(): cache = torch.arange(3 * 2 * 4 * 1 * 5).reshape(3, 2, 4, 1, 5) key, value = _vllm_kv_cache_views( @@ -598,18 +593,18 @@ def test_vllm_rocm_lhbnc_kv_cache_is_normalized_to_block_major(): assert torch.equal(value, cache[1].permute(1, 2, 0, 3)) -def test_vllm_rocm_flattened_kv_cache_is_unpacked(): - cache = torch.arange(3 * 2 * 4 * 15).reshape(3, 2, 4, 15) - - key, value = _vllm_kv_cache_views( - cache, - head_size=5, - num_kv_heads=3, - platform="rocm", - ) - - assert key.shape == (3, 4, 3, 5) - assert value.shape == (3, 4, 3, 5) +def test_vllm_rocm_flattened_kv_cache_is_unpacked(): + cache = torch.arange(3 * 2 * 4 * 15).reshape(3, 2, 4, 15) + + key, value = _vllm_kv_cache_views( + cache, + head_size=5, + num_kv_heads=3, + platform="rocm", + ) + + assert key.shape == (3, 4, 3, 5) + assert value.shape == (3, 4, 3, 5) assert torch.equal(key.flatten(2), cache[:, 0]) assert torch.equal(value.flatten(2), cache[:, 1]) @@ -723,7 +718,7 @@ def all_reduce(self, value): assert torch.equal(value.grad, torch.ones_like(value)) -def test_vllm_qwen3_strict_model_installs_without_debug_environment(monkeypatch): +def test_vllm_qwen3_strict_model_installs_without_debug_environment(monkeypatch): monkeypatch.delenv("RL_KERNEL_MODEL_DEBUG_DIR", raising=False) class RMSNorm: @@ -775,52 +770,52 @@ def __init__(self): method.apply(LinearLayer(), value), torch.full((1, 3), -1.0), ) - assert torch.equal( - norm.forward_cuda(value), - torch.nn.functional.rms_norm(value, (2,), norm.weight, 1e-6), - ) - - -def test_vllm_rocm_rotary_reuses_one_table_for_query_and_key(monkeypatch): - calls = [] - - class FakeOperator: - def __call__(self, value, positions): - calls.append(("single", tuple(value.shape), positions.clone())) - return value + 1 - - def forward_pair(self, query, key, positions): - calls.append(("pair", tuple(query.shape), tuple(key.shape), positions.clone())) - return query + 2, key + 3 - - class Rotary: - head_size = 4 - rotary_dim = 4 - - def forward_cuda(self, positions, query, key=None): - return query, key - - monkeypatch.setattr(torch.version, "hip", "test") - monkeypatch.setattr( - "rl_engine.kernels.ops.rocm.rotary_embedding.rope.RocmDeterministicRoPEOp", - FakeOperator, - ) - _patch_strict_rocm_rotary_embedding(Rotary) - rotary = Rotary() - positions = torch.tensor([7, 2]) - query = torch.arange(24, dtype=torch.float32).reshape(2, 12) - key = torch.arange(8, dtype=torch.float32).reshape(2, 4) - - query_out, key_out = rotary.forward_cuda(positions, query, key) - assert torch.equal(query_out, query + 2) - assert torch.equal(key_out, key + 3) - assert calls[0][0] == "pair" - assert calls[0][1:3] == ((3, 2, 4), (1, 2, 4)) - - query_only, absent_key = rotary.forward_cuda(positions, query) - assert torch.equal(query_only, query + 1) - assert absent_key is None - assert calls[1][0] == "single" + assert torch.equal( + norm.forward_cuda(value), + torch.nn.functional.rms_norm(value, (2,), norm.weight, 1e-6), + ) + + +def test_vllm_rocm_rotary_reuses_one_table_for_query_and_key(monkeypatch): + calls = [] + + class FakeOperator: + def __call__(self, value, positions): + calls.append(("single", tuple(value.shape), positions.clone())) + return value + 1 + + def forward_pair(self, query, key, positions): + calls.append(("pair", tuple(query.shape), tuple(key.shape), positions.clone())) + return query + 2, key + 3 + + class Rotary: + head_size = 4 + rotary_dim = 4 + + def forward_cuda(self, positions, query, key=None): + return query, key + + monkeypatch.setattr(torch.version, "hip", "test") + monkeypatch.setattr( + "rl_engine.kernels.ops.rocm.rotary_embedding.rope.RocmDeterministicRoPEOp", + FakeOperator, + ) + _patch_strict_rocm_rotary_embedding(Rotary) + rotary = Rotary() + positions = torch.tensor([7, 2]) + query = torch.arange(24, dtype=torch.float32).reshape(2, 12) + key = torch.arange(8, dtype=torch.float32).reshape(2, 4) + + query_out, key_out = rotary.forward_cuda(positions, query, key) + assert torch.equal(query_out, query + 2) + assert torch.equal(key_out, key + 3) + assert calls[0][0] == "pair" + assert calls[0][1:3] == ((3, 2, 4), (1, 2, 4)) + + query_only, absent_key = rotary.forward_cuda(positions, query) + assert torch.equal(query_only, query + 1) + assert absent_key is None + assert calls[1][0] == "single" def test_vllm_logp_replaces_every_duplicate_sampled_token_column(): @@ -983,13 +978,13 @@ def test_strict_readback_accepts_cuda_without_triton(): integration.assert_strict_ready() -def test_strict_readback_accepts_rocm_without_triton(): - plan = IntegrationPlan.from_case_ids(attention="R/R") - integration = FrameworkOperatorIntegration( - framework="megatron", - target="training", - plan=plan, - rl_kernel_operators={ +def test_strict_readback_accepts_rocm_without_triton(): + plan = IntegrationPlan.from_case_ids(attention="R/R") + integration = FrameworkOperatorIntegration( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators={ "attention": _ReadbackOperator( { "runtime_platform": "rocm", @@ -1002,17 +997,17 @@ def test_strict_readback_accepts_rocm_without_triton(): integration.record_installed_hook("attention", "test.attention") integration.execute("attention", lambda value: value, "x") - integration.assert_strict_ready() - - -def test_production_readback_infers_platform_from_real_execution_tensors(): - plan = IntegrationPlan.from_case_ids(attention="P/P") - integration = FrameworkOperatorIntegration( - framework="megatron", - target="training", - plan=plan, - rl_kernel_operators={}, - ) + integration.assert_strict_ready() + + +def test_production_readback_infers_platform_from_real_execution_tensors(): + plan = IntegrationPlan.from_case_ids(attention="P/P") + integration = FrameworkOperatorIntegration( + framework="megatron", + target="training", + plan=plan, + rl_kernel_operators={}, + ) value = torch.zeros(2) integration.execute("attention", lambda tensor: tensor + 1, value) @@ -1041,5 +1036,5 @@ def native(actual_request): integration.execute("logp", native, request) provenance = integration.readback()["operators"]["logp"]["provenance"] - assert provenance["actual_backend"] == "production.logp.test" - assert provenance["runtime_platform"] == "cpu" + assert provenance["actual_backend"] == "production.logp.test" + assert provenance["runtime_platform"] == "cpu" diff --git a/tests/test_qwen_ffn.py b/tests/test_qwen_ffn.py index f0dade06..01feaaa6 100644 --- a/tests/test_qwen_ffn.py +++ b/tests/test_qwen_ffn.py @@ -178,9 +178,7 @@ def test_qwen3_ffn_forward_backward_matches_fp32_reference(): assert actual.shape == hidden.shape assert actual.dtype is torch.bfloat16 - torch.testing.assert_close( - actual.cpu().float(), expected.detach(), atol=5e-2, rtol=2e-2 - ) + torch.testing.assert_close(actual.cpu().float(), expected.detach(), atol=5e-2, rtol=2e-2) for actual_input, reference_input in zip(actual_inputs, reference_inputs, strict=True): assert actual_input.grad.dtype is torch.bfloat16 torch.testing.assert_close( diff --git a/tests/test_rocm_e2e_ablation.py b/tests/test_rocm_e2e_ablation.py index 295936af..deec6378 100644 --- a/tests/test_rocm_e2e_ablation.py +++ b/tests/test_rocm_e2e_ablation.py @@ -44,9 +44,7 @@ def _readback(case, *, side: str, calls: int = 1, strict: bool | None = None): implementation = case.implementation_for(target) is_strict = implementation is Implementation.RL_KERNEL if strict is None else strict backend = ( - "rlkernel.attention.deterministic.v1" - if is_strict - else f"{framework}.production.attention" + "rlkernel.attention.deterministic.v1" if is_strict else f"{framework}.production.attention" ) provenance = ( { @@ -198,9 +196,7 @@ def test_strict_side_rejects_triton_and_missing_fixed_schedule(): def test_production_side_requires_framework_native_backend_identity(): case = rocm_attention_ablation_matrix(["P/P"])[0] bad = _readback(case, side="training") - bad["operators"]["attention"]["backend_id"] = ( - "rlkernel.attention.deterministic.v1" - ) + bad["operators"]["attention"]["backend_id"] = "rlkernel.attention.deterministic.v1" _routes, errors = validate_case_readbacks( case, diff --git a/tests/test_rocm_packed_ffn.py b/tests/test_rocm_packed_ffn.py index 70dcf47e..060d8d4b 100644 --- a/tests/test_rocm_packed_ffn.py +++ b/tests/test_rocm_packed_ffn.py @@ -101,12 +101,8 @@ def test_rocm_weight_gradient_uses_parameter_layout_and_is_bitwise_stable(): def test_rocm_det_linear_preserves_autograd_and_bitwise_gradients(): - inputs = torch.randn( - (16, 8), device="cuda", dtype=torch.bfloat16, requires_grad=True - ) - weight = torch.randn( - (12, 8), device="cuda", dtype=torch.bfloat16, requires_grad=True - ) + inputs = torch.randn((16, 8), device="cuda", dtype=torch.bfloat16, requires_grad=True) + weight = torch.randn((12, 8), device="cuda", dtype=torch.bfloat16, requires_grad=True) grad_output = torch.randn((16, 12), device="cuda", dtype=torch.bfloat16) op = DetGemmOp() diff --git a/tests/test_vime_linear_logp_provider.py b/tests/test_vime_linear_logp_provider.py index d8071a6f..2dad2baf 100644 --- a/tests/test_vime_linear_logp_provider.py +++ b/tests/test_vime_linear_logp_provider.py @@ -106,9 +106,7 @@ def test_provider_entropy_preserves_vime_semantics_and_autograd(): result = provider(request) reference_logits = request.logits.detach().clone().requires_grad_(True) log_probs = torch.log_softmax(reference_logits[:, :7], dim=-1) - reference_logp = log_probs[ - torch.arange(reference_logits.size(0)), request.target_ids - ] + reference_logp = log_probs[torch.arange(reference_logits.size(0)), request.target_ids] reference_entropy = -(log_probs.exp() * log_probs).sum(dim=-1) torch.testing.assert_close(result.logp.squeeze(-1), reference_logp) @@ -141,9 +139,7 @@ def from_local_logits(self, local_logits, target_ids, **_kwargs): import rl_engine.integrations.vime.linear_logp_provider as provider_module - monkeypatch.setattr( - provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp() - ) + monkeypatch.setattr(provider_module, "_default_strict_linear_logp", lambda: FakeLinearLogp()) request = _structural_request() result = provider(request) @@ -168,9 +164,7 @@ def provider(actual_request, *, linear_logp): ) wrapper = SimpleNamespace(backend_id="fake-linear-logp", provenance={}) - monkeypatch.setattr( - framework_operators, "_require_nvidia_cuda", lambda *_args: None - ) + monkeypatch.setattr(framework_operators, "_require_nvidia_cuda", lambda *_args: None) result = MegatronLogpOperator(provider, linear_logp=wrapper)(request) assert observed["context"] is request.context