From e963310ba736f8638e0459c49dcc4acd74060838 Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 24 Jun 2026 19:44:51 +0800 Subject: [PATCH 01/61] =?UTF-8?q?[test]=20=E4=BF=AE=E6=AD=A3=E9=93=BE?= =?UTF-8?q?=E5=BC=8F=E5=A4=8D=E7=94=A8=E6=B5=8B=E8=AF=95=EF=BC=9A=E5=88=A0?= =?UTF-8?q?=E5=AD=A4=E5=84=BF=E6=B5=8B=E8=AF=95=EF=BC=8C=E7=94=A8=E8=87=AA?= =?UTF-8?q?=E7=84=B6=E9=93=BE=E5=9C=BA=E6=99=AF=E6=9B=BF=E6=8D=A2=E9=94=99?= =?UTF-8?q?=E8=AF=AF=E9=A2=84=E8=AE=BE?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 删除 test_logprob_extended.py：core.logprob 已删、函数无调用，测试成孤儿导致收集中断 - 移除 test_prefix_last_in_chain_reuse_resolves_to_ancestor_with_packed_slot 及其 helper： 手工拼的 row2->row1(缺 slot 的 reuser) 链不可能由 trie 检测产生，且断言(解析到祖先) 与 _resolve_provider_for_position 的 keep_start-1 扩展正确行为相反 - 新增 test_chain_reuse_resolves_position_to_provider_with_matching_label： 用 [1,2,3,4,5]/[1,2,3,7,8]/[1,2,3,7,9] 自然链验证 seq2 位置2(token 3) 解析到 直接 provider seq1(label 7) 而非 seq0(label 4)，prefix-last 落在 seq1 真实 packed slot --- .../tests/unit_test/test_logprob_extended.py | 245 ------------------ .../tests/unit_test/test_runtime_context.py | 136 +++++----- 2 files changed, 61 insertions(+), 320 deletions(-) delete mode 100644 prefix-sharing/tests/unit_test/test_logprob_extended.py diff --git a/prefix-sharing/tests/unit_test/test_logprob_extended.py b/prefix-sharing/tests/unit_test/test_logprob_extended.py deleted file mode 100644 index 7de974b9..00000000 --- a/prefix-sharing/tests/unit_test/test_logprob_extended.py +++ /dev/null @@ -1,245 +0,0 @@ -"""Tests for logprob module — tensor helpers and core math. - -Extends the existing test_logprob.py which covers the list-based -`restore_prefix_last_logprobs` and `build_provider_prefix_last_values`. -This file adds coverage for the torch tensor helpers and validation paths. -""" - -from __future__ import annotations - -import pytest - -torch = pytest.importorskip("torch") - -from prefix_sharing.core.config import PrefixSharingConfig -from prefix_sharing.core.logprob import ( - compute_token_logprobs_from_logits, - gather_provider_prefix_last_logits, - restore_prefix_last_logprobs, - restore_prefix_last_logprobs_tensor, -) -from prefix_sharing.core.planner import PrefixSharingPlanner - - -def _make_plan(batch_sizes, prefix_lens): - """Build a PrefixSharingPlan with controlled provider/reuser layout.""" - config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=1) - planner = PrefixSharingPlanner(config) - sequences = [] - next_token = 100 - provider_seqs = {} - for i, (size, p) in enumerate(zip(batch_sizes, prefix_lens)): - if p == 0: - seq = list(range(next_token, next_token + size)) - next_token += size - provider_seqs[i] = seq - sequences.append(seq) - else: - provider_idx = max(j for j in range(i) if prefix_lens[j] == 0) - provider_seq = provider_seqs[provider_idx] - suffix = list(range(next_token, next_token + size - p)) - next_token += size - p - sequences.append(provider_seq[:p] + suffix) - return planner.plan(sequences) - - -# ------------------------------------------------------------------ -# restore_prefix_last_logprobs (list API) validation -# ------------------------------------------------------------------ - - -def test_restore_prefix_last_logprobs_wrong_suffix_length(): - plan = _make_plan([5, 4], [0, 3]) - # Pass only 1 row instead of 2 - with pytest.raises(ValueError, match="suffix_logprobs length"): - restore_prefix_last_logprobs([[0.1]], [0.0, 0.2], plan) - - -def test_restore_prefix_last_logprobs_wrong_provider_length(): - plan = _make_plan([5, 4], [0, 3]) - # Pass only 1 value instead of 2 - with pytest.raises(ValueError, match="provider_prefix_last_logprobs length"): - restore_prefix_last_logprobs([[0.1, 0.2], [0.3]], [0.0], plan) - - -def test_restore_prefix_last_logprobs_output_slot_out_of_range(): - """output_slot < 0 or > len(row) raises ValueError.""" - # We need a plan where restore spec's output_slot is out of range. - # This is hard to trigger with real plans since they compute slots - # correctly. Test the validation logic directly: - plan = _make_plan([5, 4], [0, 3]) - suffix_logprobs = [[0.1, 0.2, 0.3, 0.4, 0.5], [0.6, 0.7]] - provider_prefix_last_logprobs = [0.0, 0.8] - # The actual plan should work correctly - result = restore_prefix_last_logprobs(suffix_logprobs, provider_prefix_last_logprobs, plan) - # Verify reuser row gets the restored value prepended - assert len(result[1]) == 3 # 1 restored + 2 suffix - - -# ------------------------------------------------------------------ -# compute_token_logprobs_from_logits -# ------------------------------------------------------------------ - - -def test_compute_token_logprobs_matches_manual(): - """Verify compute_token_logprobs_from_logits matches log_softmax + gather.""" - vocab = 10 - seq_len = 5 - logits = torch.randn(seq_len, vocab) - labels = torch.randint(0, vocab, (seq_len,)) - - result = compute_token_logprobs_from_logits(logits, labels) - - # Manual: log_softmax then gather - log_probs = torch.log_softmax(logits, dim=-1) - expected = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) - - assert torch.allclose(result, expected, atol=1e-5) - - -def test_compute_token_logprobs_batched(): - """Batched logits [batch, seq, vocab] with labels [batch, seq].""" - batch, seq, vocab = 3, 5, 10 - logits = torch.randn(batch, seq, vocab) - labels = torch.randint(0, vocab, (batch, seq)) - - result = compute_token_logprobs_from_logits(logits, labels) - - log_probs = torch.log_softmax(logits, dim=-1) - expected = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) - - assert torch.allclose(result, expected, atol=1e-5) - - -# ------------------------------------------------------------------ -# restore_prefix_last_logprobs_tensor -# ------------------------------------------------------------------ - - -def test_restore_prefix_last_logprobs_tensor_reuser_gets_prepended(): - """Reuser row should have first_suffix_logprob prepended.""" - plan = _make_plan([5, 4], [0, 3]) - - # Create suffix_logprobs tensor: [batch, max_suffix_len] - # Provider row: 5 tokens (kept=5) - # Reuser row: 1 token (kept=4-3=1) - max_suffix = max(plan.kept_lengths_q) - suffix_logprobs = torch.zeros(plan.batch_size, max_suffix) - - # Fill provider row with distinct values - for j in range(plan.kept_lengths_q[0]): - suffix_logprobs[0, j] = 0.1 * (j + 1) - - # Fill reuser row with distinct values - for j in range(plan.kept_lengths_q[1]): - suffix_logprobs[1, j] = 0.2 * (j + 1) - - # First suffix logprobs: provider=0, reuser=some value - first_suffix = torch.zeros(plan.batch_size) - first_suffix[1] = 0.99 # restored value for reuser - - result = restore_prefix_last_logprobs_tensor(suffix_logprobs, first_suffix, plan) - - # Provider row should be unchanged (no restore needed) - # Reuser row should have 0.99 prepended, then suffix values - # Result shape: [batch, max_restored_len] - assert result.shape[0] == plan.batch_size - assert result.shape[1] >= plan.kept_lengths_q[0] - - # Reuser row's first position should be the restored value (float32 precision) - assert torch.allclose(result[1, 0], torch.tensor(0.99), atol=1e-5) - - -def test_restore_prefix_last_logprobs_tensor_batch_size_mismatch(): - plan = _make_plan([5, 4], [0, 3]) - - # Wrong batch dimension - suffix_logprobs = torch.zeros(1, 5) # batch=1, but plan.batch_size=2 - first_suffix = torch.zeros(2) - - with pytest.raises(ValueError, match="suffix_logprobs batch"): - restore_prefix_last_logprobs_tensor(suffix_logprobs, first_suffix, plan) - - -def test_restore_prefix_last_logprobs_tensor_first_suffix_mismatch(): - plan = _make_plan([5, 4], [0, 3]) - - suffix_logprobs = torch.zeros(2, 5) - first_suffix = torch.zeros(1) # batch=1, but plan.batch_size=2 - - with pytest.raises(ValueError, match="first_suffix_logprobs batch"): - restore_prefix_last_logprobs_tensor(suffix_logprobs, first_suffix, plan) - - -def test_restore_prefix_last_logprobs_tensor_provider_only(): - """When all rows are providers (no reuser), no restore happens.""" - plan = _make_plan([5, 3], [0, 0]) # both providers - - max_suffix = max(plan.kept_lengths_q) - suffix_logprobs = torch.randn(plan.batch_size, max_suffix) - first_suffix = torch.zeros(plan.batch_size) - - result = restore_prefix_last_logprobs_tensor(suffix_logprobs, first_suffix, plan) - - # Provider rows unchanged (no prepend) - # Logical length = kept_lengths_q for each row - for i in range(plan.batch_size): - logical_len = plan.kept_lengths_q[i] - assert torch.allclose( - result[i, :logical_len], - suffix_logprobs[i, :logical_len], - atol=1e-5, - ) - - -# ------------------------------------------------------------------ -# gather_provider_prefix_last_logits -# ------------------------------------------------------------------ - - -def test_gather_provider_prefix_last_logits_correctness(): - """Gather logits at provider prefix-last position for each reuser.""" - plan = _make_plan([5, 4], [0, 3]) - - batch, seq, vocab = plan.batch_size, 6, 10 - logits_by_batch = torch.randn(batch, seq, vocab) - - result = gather_provider_prefix_last_logits(logits_by_batch, plan) - - assert result.shape == (batch, vocab) - # Provider row should be zeros (not a reuser) - assert torch.all(result[0] == 0) - # Reuser row should contain logits from provider at prefix_last_pos - # provider_prefix_last_pos for reuser row 1 should be plan.prefix_lens[1]-1 = 2 - spec = plan.prefix_last_restore[0] - expected = logits_by_batch[spec.provider_idx_in_batch, spec.provider_prefix_last_pos] - assert torch.allclose(result[1], expected, atol=1e-5) - - -def test_gather_provider_prefix_last_logits_multiple_reusers(): - """Multiple reusers sharing the same provider.""" - plan = _make_plan([8, 6, 5], [0, 3, 3]) - - batch, seq, vocab = plan.batch_size, 8, 10 - logits_by_batch = torch.randn(batch, seq, vocab) - - result = gather_provider_prefix_last_logits(logits_by_batch, plan) - - # All reuser rows should get logits from provider (row 0) at their respective positions - for spec in plan.prefix_last_restore: - expected = logits_by_batch[spec.provider_idx_in_batch, spec.provider_prefix_last_pos] - assert torch.allclose(result[spec.reuse_idx_in_batch], expected, atol=1e-5) - - # Provider row should be zeros - assert torch.all(result[0] == 0) - - -def test_gather_provider_prefix_last_logits_no_reusers(): - """When all rows are providers, output should be all zeros.""" - plan = _make_plan([5, 3], [0, 0]) - - batch, seq, vocab = plan.batch_size, 5, 10 - logits_by_batch = torch.randn(batch, seq, vocab) - - result = gather_provider_prefix_last_logits(logits_by_batch, plan) - assert torch.all(result == 0) \ No newline at end of file diff --git a/prefix-sharing/tests/unit_test/test_runtime_context.py b/prefix-sharing/tests/unit_test/test_runtime_context.py index 3a985ee9..1de2ae92 100644 --- a/prefix-sharing/tests/unit_test/test_runtime_context.py +++ b/prefix-sharing/tests/unit_test/test_runtime_context.py @@ -72,87 +72,73 @@ def test_prefix_sharing_runtime_context_uses_padded_layout_for_restore_indices() assert ctx.stats.kept_padded_tokens == 8 -def _chain_reuse_runtime_state(): - """Chain-reuse: row0=provider, row1=reuse(0,prefix=3), row2=reuse(1,prefix=3). - - Row2's prefix-last predict position (prefix_len-1 = 2) falls on row1's - keep_start-1 (row1 keep_start=3). Under v080 physical trimming row1's - packed region has no slot for position 2, so the prefix-last logits must - be fetched from the chain ancestor row0 whose packed region strictly - contains position 2 (keep_start=0 <= 2 < keep_end=8). - - The chain is forced via a hand-built PrefixDetectionResult (the trie - detector naturally routes row2 to row0 because it records the *first* - inserter as provider at each depth; we need the transitive edge - row2->row1 specifically to exercise the keep_start-1 miss path). +def test_chain_reuse_resolves_position_to_provider_with_matching_label(): + """Chain-reuse must route a deeper reuser to the provider whose + *continuation* matches, not the root. + + Trie detection produces a natural chain here: + seq0: [1, 2, 3, 4, 5] provider (root) + seq1: [1, 2, 3, 7, 8] reuse(seq0), shared prefix [1,2,3] (len 3) + seq2: [1, 2, 3, 7, 9] reuse(seq1), shared prefix [1,2,3,7] (len 4) + + The trie routes seq2 to seq1 (not seq0): at position 3 both seq1 and seq2 + hold token 7, whereas seq0 holds 4. So seq2's restore for position 2 + (token "3", which predicts token 7) must reference seq1 — whose label at + that position is 7 — rather than seq0, whose label there is 4. + + The keep_start-1 extension in ``_resolve_provider_for_position`` keeps the + boundary position on the direct provider seq1 (seq1 keeps [3, 5), so + keep_start-1 = 2 covers position 2). Positions shared identically with + the root (0, 1) walk up to seq0. The prefix-last (position 3) lands on a + real packed slot of seq1 because seq1's kept range [3, 5) contains it. """ - from prefix_sharing.core.prefix_detector import ( - PrefixDetectionResult, - PrefixReuseSpec, - ) - planner = PrefixSharingPlanner(PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)) - sequences = [ - [100, 101, 102, 103, 104, 105, 106, 107], # row0 provider, len 8 - [100, 101, 102, 108, 109, 110], # row1 reuse(0) prefix=3, len 6 - [100, 101, 102, 111, 112, 113], # row2 reuse(1) prefix=3, len 6 - ] - detection = PrefixDetectionResult( - batch_size=3, - reuse_specs=( - PrefixReuseSpec(reuse_idx_in_batch=1, provider_idx_in_batch=0, prefix_len=3), - PrefixReuseSpec(reuse_idx_in_batch=2, provider_idx_in_batch=1, prefix_len=3), - ), - groups=(), - group_ids=[0, 0, 1], - provider_index=[0, 0, 1], - prefix_lens=[0, 3, 3], - is_provider=[True, False, False], - ) - prefix_sharing_plan = planner.plan_from_detection( - sequences, detection, forward_id=10, micro_batch_id=20, + plan = planner.plan( + [[1, 2, 3, 4, 5], [1, 2, 3, 7, 8], [1, 2, 3, 7, 9]], + forward_id=10, + micro_batch_id=20, ) - # Sanity: the chain is row2 -> row1 -> row0 with prefix_len 3 each. - assert prefix_sharing_plan.provider_index[1] == 0 - assert prefix_sharing_plan.provider_index[2] == 1 - assert prefix_sharing_plan.prefix_lens[1] == 3 - assert prefix_sharing_plan.prefix_lens[2] == 3 - return PrefixSharingRuntimeState( - prefix_sharing_plan=prefix_sharing_plan, + # Chain: seq2 -> seq1 -> seq0; seq2 shares the longer prefix [1,2,3,7]. + assert plan.provider_index == [0, 0, 1] + assert plan.prefix_lens == [0, 3, 4] + + runtime_state = PrefixSharingRuntimeState( + prefix_sharing_plan=plan, attention_backend=None, - packed_batch_layout=PackedBatchLayout.from_valid_lengths(prefix_sharing_plan.kept_lengths_q), + packed_batch_layout=PackedBatchLayout.from_valid_lengths(plan.kept_lengths_q), parallel_info=MegatronParallelInfo(), ) - - -def test_prefix_last_in_chain_reuse_resolves_to_ancestor_with_packed_slot(): - """Regression: prefix-last in chain-reuse must fetch logits from the - chain ancestor whose packed region strictly contains the predict position, - not the intermediate reuser whose keep_start-1 has no packed slot. - - Before the fix _build_prefix_last_restore_indices left provider_1d_pos=-1 - (sentinel) for such specs, causing vocab_logprobs save to skip them and - the downstream restore to KeyError on the saved-logits lookup. - """ - runtime_state = _chain_reuse_runtime_state() with prefix_sharing_runtime_context(runtime_state) as ctx: - # Collect prefix-last (non-interior) indices for row2 (reuse_idx=2). - row2_plast = [ - idx for idx in ctx.prefix_last_restore_indices - if idx.reuse_idx_in_batch == 2 and not idx.is_shared_prefix_interior - ] - assert len(row2_plast) == 1, "row2 should have exactly one prefix-last restore" - spec = row2_plast[0] - - # target_2d_pos = prefix_len-1 = 2 (prefix-last predict position). - assert spec.target_2d_pos == 2 - # The fix: provider_1d_pos must NOT be the -1 sentinel. It must point - # into row0's packed region (row0 has 8 tokens, position 2 → packed - # index 2 since row0 starts at offset 0 in the packed tensor). - assert spec.provider_1d_pos != -1, ( - "prefix-last in chain-reuse must resolve provider_1d_pos to a " - "chain ancestor with a real packed slot, not the -1 sentinel" + seq2 = { + idx.target_2d_pos: idx + for idx in ctx.prefix_last_restore_indices + if idx.reuse_idx_in_batch == 2 + } + + # Positions 0, 1 (tokens 1, 2): labels 2, 3 are identical across all + # rows, so they resolve up the chain to the root seq0. + assert seq2[0].provider_idx_in_batch == 0 + assert seq2[0].label_value == 2 + assert seq2[1].provider_idx_in_batch == 0 + assert seq2[1].label_value == 3 + + # Position 2 (token "3"): seq2/seq1 predict 7 here, seq0 predicts 4. + # Must resolve to the direct provider seq1 (matching label 7), and + # stay there via the keep_start-1 extension — NOT walk up to seq0. + assert seq2[2].provider_idx_in_batch == 1, ( + "seq2 position 2 (token 3 -> label 7) must resolve to seq1 " + "(matching continuation 7), not seq0 (continuation 4)" ) - assert spec.provider_1d_pos == 2, ( - f"expected packed index 2 (row0 offset 0 + pos 2), got {spec.provider_1d_pos}" + assert seq2[2].label_value == 7 + + # Prefix-last (position 3, predicts token 9): seq1 keeps [3, 5) which + # strictly contains position 3, so it resolves to a real packed slot + # on seq1 (not the -1 sentinel). + plast = seq2[3] + assert plast.is_shared_prefix_interior is False + assert plast.provider_idx_in_batch == 1 + assert plast.label_value == 9 + assert plast.provider_1d_pos != -1, ( + "prefix-last must resolve to a real packed slot on the direct " + "provider seq1, not the -1 sentinel" ) From 16877df1a4eca2e91a75b5c64f0bc21d2e0a494c Mon Sep 17 00:00:00 2001 From: Boundless Date: Thu, 25 Jun 2026 15:29:29 +0800 Subject: [PATCH 02/61] =?UTF-8?q?[refactor]=20logprob/entropy=20restore=20?= =?UTF-8?q?=E6=94=B9=E4=B8=BA=20build=5Fkv=20=E5=BC=8F=E5=8C=BA=E9=97=B4?= =?UTF-8?q?=E6=8B=BC=E6=8E=A5=EF=BC=8C=E6=B8=85=E7=90=86=E5=86=97=E4=BD=99?= =?UTF-8?q?=20interior-spec=20=E6=95=B0=E6=8D=AE=E7=BB=93=E6=9E=84?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - verl_mcore/context/vocab_logprobs/observability: restore 由逐元素 2D 注入改为整段切片 + prefix-last 重算 - planner: 仅生成 prefix-last spec，删除 interior spec 与 is_shared_prefix_interior 标记 - 删除死字段 reuse_1d_pos、valid_indices 链、provider_predict_pos、reuse_first_suffix_label_pos - 同步单测与 overview 文档；更新精度诊断工具 --- docs/overview.md | 42 +-- docs/overview.puml | 8 +- .../prefix_sharing/core/observability.py | 5 +- prefix-sharing/prefix_sharing/core/planner.py | 93 ++----- .../prefix_sharing/integrations/context.py | 96 ++----- .../prefix_sharing/integrations/verl_mcore.py | 251 +++++++----------- .../vocab_logprobs.py | 24 +- .../prefix_sharing/tools/cmp_diag_verl080.py | 104 +++++++- .../prefix_sharing/tools/diagnostic_dump.py | 118 +++++++- .../tools/diagnostic_dump_verl080.py | 10 +- .../test_verl_megatron_runtime_helpers.py | 51 ++-- .../tests/unit_test/test_observability.py | 4 +- .../tests/unit_test/test_planner.py | 77 ++---- .../unit_test/test_restore_unfold_verl080.py | 29 +- .../tests/unit_test/test_runtime_context.py | 84 +++--- 15 files changed, 507 insertions(+), 489 deletions(-) diff --git a/docs/overview.md b/docs/overview.md index 7dcf5e32..22032124 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -537,24 +537,36 @@ Integrations 层实现具体的框架适配逻辑，可被独立测试和迭代 | `PrefixSharingRuntimeState` | 跨层传递的运行时状态，包含 plan、backend、packed layout 和 Megatron 并行信息 | | `PackedPrefixLastRestoreIndex` | runtime context 中派生出的 THD 1D 读写索引 | -**restore_suffix_first_log_probs_from_prefix 逻辑**: +**restore_reuser_prefix_columns_2d 逻辑**（build_kv 式区间拼接，详见 `verl_mcore.py`）: ```python -def restore_suffix_first_log_probs_from_prefix(logits, labels, log_probs, vocab_fn): +def restore_reuser_prefix_columns_2d(output, vocab_parallel_log_probs_fn, ...): ctx = current_prefix_sharing_context() - if ctx is None or not ctx.prefix_last_restore_indices: - return log_probs - if not ctx.parallel_info.is_pipeline_last_stage: - return log_probs - - restored = log_probs.clone() - for index in ctx.prefix_last_restore_indices: - # 从 provider 的 prefix-last 位置取 logits - provider_logits = logits[0:1, index.provider_1d_pos:index.provider_1d_pos+1, :] - # 用 reuser 的第一个 suffix label 计算 logprob - reuse_label = labels[0:1, index.reuse_1d_pos:index.reuse_1d_pos+1] - restored[0, index.reuse_1d_pos] = vocab_fn(provider_logits, reuse_label) - return restored + plan = ctx.prefix_sharing_plan + if ctx is None or not plan.has_sharing: + return output + + log_probs = output["log_probs"] # [B, L] 2D + prefix_last_by_row = {idx.reuse_idx_in_batch: idx + for idx in ctx.prefix_last_restore_indices} + # 按 batch 顺序遍历：provider 必在 reuser 之前（检测器不变量）， + # 故读到 provider 行时它已恢复完。 + for i in range(1, B): + if provider_index[i] == i or prefix_lens[i] <= 0: + continue # provider / 非 reuser + p = provider_index[i] + P = prefix_lens[i] + # interior [0, P-2]：整段从 provider 已恢复行 bulk 复制 + log_probs[i, 0:P-1] = log_probs[p, 0:P-1] + # prefix-last (P-1)：用 reuser 自己的首个 suffix label 重算 + idx = prefix_last_by_row.get(i) + if idx is not None: + saved = ctx.prefix_last_logits_saved[(i, idx.target_2d_pos)] + log_probs[i, P-1] = vocab_parallel_log_probs_fn( + saved, idx.label_value).reshape(()) + else: # suffix_len==0：列被下游 mask + log_probs[i, P-1] = log_probs[p, P-1] + return output ``` #### 4.2.2 context.py - 运行时上下文管理 diff --git a/docs/overview.puml b/docs/overview.puml index 1ba72d10..171d8729 100644 --- a/docs/overview.puml +++ b/docs/overview.puml @@ -94,7 +94,8 @@ package "prefix-sharing/integrations/ (框架适配层)" as IntegrationLayer #FF +reuse_idx_in_batch: int +provider_idx_in_batch: int +provider_1d_pos: int - +reuse_1d_pos: int + +target_2d_pos: int + +label_value: int } class "megatron_runtime.py" as MegatronRuntime #FFCC80 { @@ -189,10 +190,9 @@ package "prefix-sharing/core/ (核心语义层)" as CoreLayer #F3E5F5 { 描述如何从provider恢复reuser的prefix-last位置logprob +reuse_idx_in_batch: int +provider_idx_in_batch: int - +provider_predict_pos: int - +reuse_first_suffix_label_pos: int +group_id: int - +is_shared_prefix_interior: bool + +target_2d_pos: int + +label_value: int } class "planner.py\nPrefixSharingPlan" as Plan #CE93D8 { diff --git a/prefix-sharing/prefix_sharing/core/observability.py b/prefix-sharing/prefix_sharing/core/observability.py index 5caef03a..7c3e9edd 100644 --- a/prefix-sharing/prefix_sharing/core/observability.py +++ b/prefix-sharing/prefix_sharing/core/observability.py @@ -115,7 +115,10 @@ def from_plan( expected_reused_prefix_tokens_per_layer=sum( prefix_sharing_plan.prefix_lens[index] for index in reuser_indices ), - expected_restore_count=len(prefix_sharing_plan.prefix_last_restore), + # Counts reuser rows restored (one bulk slice per reuser), matching + # record_restore's per-reuser count. interior specs are no longer + # restored per-position, so don't count them here. + expected_restore_count=len(reuser_indices), ) def layer(self, layer_id: int) -> PrefixSharingLayerStats: diff --git a/prefix-sharing/prefix_sharing/core/planner.py b/prefix-sharing/prefix_sharing/core/planner.py index c16a1c22..7a872a72 100644 --- a/prefix-sharing/prefix_sharing/core/planner.py +++ b/prefix-sharing/prefix_sharing/core/planner.py @@ -61,49 +61,35 @@ @dataclass(frozen=True) class PrefixLastRestoreSpec: - """Plan for one reuse row's logprob restore. - - Two kinds of restore exist: - - * **Prefix-last restore** (``is_shared_prefix_interior=False``, default): the - reuser's first suffix token is predicted from the shared prefix-last - position. Since different reusers may have different first suffix labels, - the provider's full logits at ``prefix_len - 1`` must be stored and the - logprob computed per reuser using that reuser's label. - - * **Shared-prefix interior restore** (``is_shared_prefix_interior=True``): - a token that lives strictly inside the shared prefix. Its - logprob is ``log_softmax(logits[pos-1])[label=pos]``, where both the - logits position and the label belong to the shared prefix. The logprob - is therefore **identical for all reusers** and can be computed once from - the provider's logits, stored as a non-detached scalar tensor, and - reused for every reuser without re-reading logits. + """Plan for one reuse row's prefix-last logprob restore. + + A reuser's first suffix token is predicted from the shared prefix-last + position (``logits[prefix_len - 1]`` predicting ``input_ids[prefix_len]``). + Since different reusers may have different first suffix labels, the + provider's full logits at ``prefix_len - 1`` must be stored and the logprob + recomputed per reuser using that reuser's label. + + Interior prefix positions (``[0, prefix_len - 2]``) are **not** represented + here: their logprob is identical across the shared prefix (same logits + + same labels), so the 2D restore bulk-slices the whole interval off the + direct provider's already-restored row — no per-position spec is needed. """ reuse_idx_in_batch: int provider_idx_in_batch: int - provider_predict_pos: int - """Logits position in provider used for this restore (logits[p] predicts - token at p+1). Equals ``target_2d_pos`` numerically; kept as a separate - field to document the provider-side lookup semantic and to drive provider - chain resolution in ``context.py``.""" - reuse_first_suffix_label_pos: int group_id: int - is_shared_prefix_interior: bool = False target_2d_pos: int = -1 - """Absolute 2D position in output tensor where restored logprob belongs. - - This is the label position in the original (untrimmed) sequence: - log_probs[i] = log_softmax(logits[i])[label[i]] - For shared-prefix interior: ``target_2d_pos = prefix_label_pos - 1`` (label index). - For prefix-last: ``target_2d_pos = prefix_len - 1`` (first-suffix label index). + """The prefix-last position (``prefix_len - 1``): the last token of the + shared prefix. By the log_probs layout invariant + (``log_probs[i] = log_softmax(logits[i])[label[i+1]]``) this single index + is **both** the provider logits row that predicts the reuser's first + suffix token (read side) and the output column the recomputed logprob is + written to (write side) — so one field covers both roles. """ label_value: int = -1 - """The actual token ID used as label for logprob computation. - - For shared-prefix interior: token at ``prefix_label_pos`` (shared prefix, same for provider/reuser). - For prefix-last: token at ``prefix_len`` (reuser's first suffix token). - This value is needed because trimmed packed labels don't contain these positions. + """The reuser's first suffix token ID (``input_ids[prefix_len]``), used as + the label for logprob recompute. Needed because the trimmed packed labels + don't contain this position. """ @@ -282,47 +268,20 @@ def plan_from_detection( kept_len = suffix_len q_offset = prefix_len - # --- Shared-prefix interior token restore --- - # Response tokens inside the shared prefix (positions - # 1 .. prefix_len-1) are trimmed from the Q path but - # still need logprob entries for PPO loss. Their labels are - # inside the shared prefix so the logprob is identical for - # provider and reuser: compute once from provider logits, - # store as a non-detached scalar tensor. - # Note: prompt positions (0..prompt_len-1) are also - # restored, but downstream label_mask is False for them - # so they don't affect loss — restoring the whole prefix - # column uniformly keeps planning simple. - for prefix_label_pos in range(1, prefix_len): - # Shared-prefix interior token logprob: computed from - # logits[prefix_label_pos-1] predicting token at - # prefix_label_pos. Writes to 2D position - # prefix_label_pos - 1 (the label position). - restore_specs.append( - PrefixLastRestoreSpec( - reuse_idx_in_batch=index, - provider_idx_in_batch=provider_index[index], - provider_predict_pos=prefix_label_pos - 1, - reuse_first_suffix_label_pos=prefix_label_pos, - group_id=group_ids[index], - is_shared_prefix_interior=True, - target_2d_pos=prefix_label_pos - 1, - label_value=input_ids[index][prefix_label_pos], - ) - ) - # --- Prefix-last restore (first suffix token) --- + # Interior prefix positions [0, prefix_len-2] are restored by + # the 2D restore, which bulk-slices them off the direct + # provider's already-restored row (identical across the shared + # prefix), so the planner only emits the prefix-last token. # The logits at position prefix_len-1 predict the first suffix # token whose label is input_ids[prefix_len] (differs per - # reuser). The restored logprob is written to 2D position + # reuser); the restored logprob is written to 2D position # prefix_len - 1 (the label slot for first-suffix prediction). if prefix_len > 0 and suffix_len > 0: restore_specs.append( PrefixLastRestoreSpec( reuse_idx_in_batch=index, provider_idx_in_batch=provider_index[index], - provider_predict_pos=prefix_len - 1, - reuse_first_suffix_label_pos=prefix_len, group_id=group_ids[index], target_2d_pos=prefix_len - 1, label_value=input_ids[index][prefix_len], diff --git a/prefix-sharing/prefix_sharing/integrations/context.py b/prefix-sharing/prefix_sharing/integrations/context.py index 49329c8e..54d6b75f 100644 --- a/prefix-sharing/prefix_sharing/integrations/context.py +++ b/prefix-sharing/prefix_sharing/integrations/context.py @@ -25,8 +25,6 @@ class PackedPrefixLastRestoreIndex: reuse_idx_in_batch: int provider_idx_in_batch: int provider_1d_pos: int - reuse_1d_pos: int - is_shared_prefix_interior: bool = False target_2d_pos: int = -1 """Absolute 2D position in output where restored logprob is written.""" label_value: int = -1 @@ -49,12 +47,10 @@ class PrefixSharingRuntimeContext: cloned from packed logits after temperature scaling, before any in-place modification by entropy/logprob computation. - Only populated for non-interior (prefix-last) restore specs. - Interior specs use direct 2D copy instead. + Populated for prefix-last restore specs (one per reuser-with-suffix); + interior prefix columns are bulk-copied by the 2D restore and need no + saved logits. """ - valid_indices: list | None = None - """Per-row tensor positions of valid tokens in the original 2D tensors. - Used to map planner's valid-space target_2d_pos to tensor-space columns.""" stats: PrefixSharingStats | None = None def __init__(self, runtime_state: Any, store: PrefixAttentionStore) -> None: @@ -70,13 +66,9 @@ def __init__(self, runtime_state: Any, store: PrefixAttentionStore) -> None: ) # Provider packed logits saved for prefix-last logprob recompute in 2D # space. Populated lazily by the verl vocab-logprobs patch for each - # non-interior (prefix-last) restore spec; read by + # prefix-last restore spec (one per reuser-with-suffix); read by # restore_reuser_prefix_columns_2d. Always present (possibly empty). self.prefix_last_logits_saved: dict[tuple[int, int], Any] = {} - # valid_indices maps planner's valid-space target_2d_pos to tensor - # columns; used by the 2D restore path in verl_mcore. May be absent - # on minimal runtime states (e.g. unit-test fixtures). - self.valid_indices = getattr(runtime_state, "valid_indices", None) # stats drives the per-micro-batch audit log. Prefer an explicit # stats object carried by the runtime state; otherwise derive one # from the plan so the audit summary is always available. @@ -94,75 +86,35 @@ def current_prefix_sharing_context() -> PrefixSharingRuntimeContext | None: return _current_context.get() -def _resolve_provider_for_position( - plan: PrefixSharingPlan, - provider_idx: int, - target_pos: int, -) -> int: - """Walk up the provider chain to find the nearest ancestor whose - packed layout contains ``target_pos`` (absolute original position). - - In chain-reuse scenarios (row 0 → row 1 → row 2), intermediate - providers are reusers with truncated packed layouts. The root - provider may also be shorter than the reuser's prefix — in that - case intermediate providers contribute the extended range. - - The range is extended left by one position (keep_start - 1) for - reusers: that position holds the logprob for the token at - keep_start which is still inside the shared prefix. Its value was - computed via prefix-last recompute with the correct shared label - and is safe to copy from. - """ - while True: - keep_start, keep_end = plan.input_keep_ranges[provider_idx] - if keep_start - 1 <= target_pos < keep_end: - return provider_idx - if plan.is_reuser(provider_idx): - provider_idx = plan.provider_index[provider_idx] - else: - # Can't walk further; return current (should not happen - # with valid detection, but guard) - return provider_idx - - def _build_prefix_last_restore_indices( prefix_sharing_plan: PrefixSharingPlan, packed_batch_layout: PackedBatchLayout, ) -> list[PackedPrefixLastRestoreIndex]: + """Build prefix-last restore indices — one per reuser-with-suffix. + + The planner now emits only prefix-last specs (interior prefix columns are + bulk-sliced by the 2D restore off the direct provider's already-restored + row, so no per-position index is needed for them). The prefix-last logits + always live in the **direct provider's** packed region — chain reuse forces + ``prefix_len_reuser > prefix_len_provider`` (a reuser only becomes someone's + provider by extending the trie *beyond* its own prefix), so the reuser's + prefix-last position ``prefix_len - 1`` is at or beyond the direct + provider's ``keep_start`` and thus inside its computed suffix region. No + chain walk up to ancestors is needed. + """ indices = [] for spec in prefix_sharing_plan.prefix_last_restore: - reuse_idx = spec.reuse_idx_in_batch - # Resolve through chain reuse to the nearest provider whose - # packed layout contains provider_predict_pos. - target_pos = spec.provider_predict_pos - resolved_provider = _resolve_provider_for_position( - prefix_sharing_plan, spec.provider_idx_in_batch, target_pos - ) - - # Interior: provider and reuser share the same prefix tokens, - # so the logprob at target_pos is identical — just copy. - # 2D restore uses provider_idx_in_batch + valid_indices mapping, - # no packed-index lookup needed. 1d_pos_in_provider is only used - # for prefix-last entries to fetch saved provider logits. - provider_offset = ( - target_pos - prefix_sharing_plan.input_keep_ranges[resolved_provider][0] - ) - if provider_offset >= 0: - # target_pos is inside the resolved provider's packed region. - pos_1d_in_provider = packed_batch_layout.packed_index(resolved_provider, provider_offset) - else: - # Interior spec resolved to an intermediate reuser's keep_start-1 - pos_1d_in_provider = -1 - - reuse_1d = -1 # sentinel: no slot in reuser packed region - + provider = spec.provider_idx_in_batch # direct provider + target_pos = spec.target_2d_pos + # offset of prefix-last within the direct provider's packed region. + # Proven >= 0 (see docstring); guard cheaply to surface regressions. + provider_offset = target_pos - prefix_sharing_plan.input_keep_ranges[provider][0] + pos_1d_in_provider = packed_batch_layout.packed_index(provider, provider_offset) indices.append( PackedPrefixLastRestoreIndex( - reuse_idx_in_batch=reuse_idx, - provider_idx_in_batch=resolved_provider, + reuse_idx_in_batch=spec.reuse_idx_in_batch, + provider_idx_in_batch=provider, provider_1d_pos=pos_1d_in_provider, - reuse_1d_pos=reuse_1d, - is_shared_prefix_interior=spec.is_shared_prefix_interior, target_2d_pos=spec.target_2d_pos, label_value=spec.label_value, ) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index 4c09c8da..ef6490c4 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -43,10 +43,6 @@ class PrefixSharingRuntimeState: attention_backend: Any packed_batch_layout: PackedBatchLayout parallel_info: MegatronParallelInfo - valid_indices: list | None = None - """Per-row tensor positions of valid (non-padding) tokens in the - original 2D tensors. Used to map planner's valid-space target_2d_pos - to tensor-space columns (needed when sequences have left padding).""" kept_position_ids: Any | None = None @@ -229,7 +225,6 @@ def build_prefix_sharing_micro_batch_verl070( attention_backend=get_backend_instance(config, backend), packed_batch_layout=packed_batch_layout, parallel_info=parallel_info, - valid_indices=valid_indices, ) print( f"[PS][prepare] PATH 6 DONE: returning (trimmed_micro_batch, " @@ -240,57 +235,58 @@ def build_prefix_sharing_micro_batch_verl070( def restore_reuser_prefix_columns_2d( output: dict[str, Any], - label_2d: Any, vocab_parallel_log_probs_fn: Any, vocab_parallel_entropy_fn: Any = None, ) -> dict[str, Any]: - """Restore reuser prefix columns in 2D space after postprocess_packed_seqs. - - All prefix-column restoration happens purely in 2D [B, L] space, - consolidating the previous three-phase approach (packed compute → - cache → 2D inject) into a single post-forward step. - - For each :class:`PackedPrefixLastRestoreIndex` in the runtime context: - - - **Interior response** (shared-prefix token): logprob and entropy - are identical between provider and reuser because the label is the - same shared token and the logits are the same (same KV). Directly - copy from the provider's 2D row. - - - **Prefix-last token**: entropy is still the same (same logits), - so copy from provider's 2D row. Logprob depends on the label - which differs (reuser's first suffix token ≠ provider's), so - recompute from saved provider packed logits + reuser's 2D label. - - Must be called while ``prefix_sharing_runtime_context`` is still - active (i.e. before the context manager exits), and after - ``postprocess_packed_seqs`` has produced the 2D output dict. + """Restore reuser prefix columns in 2D space — build_kv-style slice + concat. + + Mirrors :meth:`TorchReferenceBackend.build_kv` + (``torch.cat([provider_kv[:prefix_len], own_suffix])``): instead of writing + each prefix token one scalar at a time, the whole prefix interval is sliced + off the **direct provider's already-restored 2D row** and only the single + prefix-last logprob is recomputed. + + Per reuser row ``i`` with direct provider ``p = provider_index[i]`` and + ``P = prefix_lens[i]`` (columns are identity-mapped in the unfolded 2D + tensor, so ``target_2d_pos`` == column): + + - **interior ``[0, P-2]``**: ``log_probs[i, 0:P-1] = log_probs[p, 0:P-1]`` + (bulk copy). Identical across the shared prefix (same logits + labels), + and ``p`` was restored earlier in the batch-order loop, so its row already + holds correct values — no per-position provider resolution needed. + - **prefix-last ``P-1``**: recompute ``log_probs[i, P-1]`` from the saved + provider logits + the reuser's own first-suffix label (differs from the + provider's). When the reuser has no suffix (``suffix_len == 0``) the + planner emits no prefix-last spec; that column is masked downstream, so + the provider's value is copied as a safe placeholder. + - **entropy ``[0, P-1]``**: ``entropy[i, 0:P] = entropy[p, 0:P]`` (whole + prefix copied, including prefix-last — entropy is label-independent). + + Rows are visited in ``range(B)`` order so a provider is always restored + before any reuser that reads it (the same online-detector invariant + ``build_kv`` relies on). Args: - output: Output dict from forward, with ``log_probs`` [B, L] and - optionally ``entropy`` [B, L] in 2D space. - label_2d: Original 2D label ``[B, L]`` or ``None``. Fallback label - source for prefix-last recompute when ``index.label_value < 0``; - ignored whenever ``index.label_value >= 0`` (the common path — - the per-index label is populated framework-agnostically by the - context for both v070 and v080). Pass ``None`` when every restore - index carries a valid ``label_value`` (verl080 path — avoids - materialising a dense ``[B, L_max]`` long tensor just to look up - a handful of entries). **Position 2 (required, no default)** to - match v070's positional-arg call site - ``(output, label, log_probs_fn, entropy_fn)``. - vocab_parallel_log_probs_fn: Function to compute logprob from - packed logits [1, 1, V//tp] and label [1, 1] → scalar. - Typically :func:`verl.utils.megatron.tensor_parallel.vocab_parallel_log_probs_from_logits`. - vocab_parallel_entropy_fn: Optional function to compute entropy - from packed logits [1, V//tp] → scalar. + output: Output dict with ``log_probs`` [B, L] and optionally + ``entropy`` [B, L] in 2D space (unfolded from the trimmed + NestedTensor by :func:`restore_via_2d_unfold_verl080`). + vocab_parallel_log_probs_fn: ``logits [1, V//tp]``, ``label [1]`` → + scalar; used only for the prefix-last recompute. + vocab_parallel_entropy_fn: Retained for call-site compatibility; + unused (entropy is copied, never recomputed). Returns: ``output`` with ``log_probs`` and ``entropy`` mutated in-place. """ ctx = current_prefix_sharing_context() - if ctx is None or not ctx.prefix_last_restore_indices: + if ctx is None: + return output + plan = ctx.prefix_sharing_plan + # Guard on reuser presence (not on prefix_last_restore_indices): a batch + # whose reusers all have suffix_len == 0 emits no prefix-last spec but still + # needs its interior prefix columns restored. + if not plan.has_sharing: return output import torch @@ -299,96 +295,54 @@ def restore_reuser_prefix_columns_2d( if log_probs is None: return output entropy = output.get("entropy") - # Map planner's valid-space target_2d_pos (0-based within valid content) - # to tensor-space 2D columns. postprocess_packed_seqs places tokens at - # their ORIGINAL attention_mask positions, so a valid-space position p - # maps to valid_indices[row][p] in the [B, L] output. - valid_indices = ctx.valid_indices - - def _map_2d_col(row: int, valid_pos: int) -> int: - if valid_indices is not None: - vi = valid_indices[row] - if vi is not None and 0 <= valid_pos < len(vi): - return int(vi[valid_pos].item()) - return valid_pos - - # ── 诊断：取 saved logits 之前，打印 saved dict 的 key 和期望的非 interior key ── - # 用于定位 vocab_parallel_log_probs_from_logits patch 是否真的在 verl080 - # logits_processor 闭包调用点上命中、保存了 provider prefix-last logits。 - # saved 空 + expected 非空 → 保存侧 patch 完全没生效（没进 patched_fn 循环） - # saved 有但缺某 key → 部分保存，看缺的是不是 (1,15) 这类 - # saved == expected → 保存正常，KeyError 另有原因 - _saved_keys = sorted(ctx.prefix_last_logits_saved.keys()) - _expected_keys = sorted( - (i.reuse_idx_in_batch, i.target_2d_pos) - for i in ctx.prefix_last_restore_indices - if not i.is_shared_prefix_interior - ) - _interior_n = len(ctx.prefix_last_restore_indices) - len(_expected_keys) - print( - f"[PS][diag] prefix_last_logits_saved keys ({len(_saved_keys)})=" - f"{_saved_keys}", - flush=True, - ) - print( - f"[PS][diag] expected non-interior keys ({len(_expected_keys)})=" - f"{_expected_keys} (interior={_interior_n}, not read from saved)", - flush=True, - ) - non_interior_count = 0 - for index in ctx.prefix_last_restore_indices: - reuser_row = index.reuse_idx_in_batch - provider_row = index.provider_idx_in_batch - valid_col = index.target_2d_pos - # Map valid-space col to per-row tensor-space columns: - # postprocess_packed_seqs places tokens at original - # attention_mask positions, which may differ between rows. - provider_col = _map_2d_col(provider_row, valid_col) - reuser_col = _map_2d_col(reuser_row, valid_col) - - if index.is_shared_prefix_interior: - # Shared-prefix interior: token is in shared prefix → logprob and entropy - # are identical to the provider's (same label, same logits). - log_probs[reuser_row, reuser_col] = log_probs[provider_row, provider_col] - if entropy is not None: - entropy[reuser_row, reuser_col] = entropy[provider_row, provider_col] - else: - # Prefix-last: entropy is the same (same logits), so copy - # from provider's 2D row. Logprob differs because reuser's - # first suffix token ≠ provider's. - non_interior_count += 1 - if entropy is not None: - entropy[reuser_row, reuser_col] = entropy[provider_row, provider_col] - - # Recompute logprob from saved provider packed logits - # with reuser's own label. - # vocab_parallel_log_probs_from_logits expects: - # logits: [N, V//tp] labels: [N] - saved_key = (reuser_row, valid_col) + provider_index = plan.provider_index + prefix_lens = plan.prefix_lens + + # prefix-last lookup keyed by reuser row. ``prefix_last_restore_indices`` + # now carries only prefix-last entries (one per reuser-with-suffix); + # interior is handled by the bulk slice below. + prefix_last_by_row = { + idx.reuse_idx_in_batch: idx for idx in ctx.prefix_last_restore_indices + } + + restored_reusers = 0 + # Row 0 is always a provider (nothing precedes it to reuse), so start at 1. + # A reuser's provider always has a smaller batch index (online-detector + # invariant), so it is already restored when we reach row i. + for i in range(1, len(prefix_lens)): + prefix_len = prefix_lens[i] + if provider_index[i] == i or prefix_len <= 0: + continue # provider / non-reuser: row already complete + provider = provider_index[i] + + # interior [0, prefix_len-2]: bulk-copy from the provider's restored row. + if prefix_len - 1 > 0: + log_probs[i, 0:prefix_len - 1] = log_probs[provider, 0:prefix_len - 1] + + # prefix-last (position prefix_len-1): recompute with the reuser's label. + idx = prefix_last_by_row.get(i) + if idx is not None: + saved_key = (i, idx.target_2d_pos) provider_logits = ctx.prefix_last_logits_saved[saved_key] # [1, V//tp] - # Prefer the per-index label_value (framework-agnostic; populated - # by the context for both v070 and v080). Fall back to dense - # label_2d for backward compat — avoids materialising a - # [B, L_max] long tensor just to read a handful of entries. - if index.label_value >= 0: - reuser_label = torch.tensor( - [index.label_value], dtype=torch.long, device=log_probs.device, - ) # [1] - elif label_2d is not None: - reuser_label = label_2d[reuser_row:reuser_row + 1, reuser_col:reuser_col + 1].view(1) # [1] - else: - raise RuntimeError( - "prefix-last restore needs a label but index.label_value<0 " - "and label_2d is None" - ) - log_probs[reuser_row, reuser_col] = vocab_parallel_log_probs_fn( - provider_logits, # [1, V//tp] - reuser_label, # [1] + reuser_label = torch.tensor( + [idx.label_value], dtype=torch.long, device=log_probs.device, + ) # [1] + log_probs[i, prefix_len - 1] = vocab_parallel_log_probs_fn( + provider_logits, reuser_label, ).reshape(()) + else: + # suffix_len == 0: no prefix-last spec; column is masked downstream. + log_probs[i, prefix_len - 1] = log_probs[provider, prefix_len - 1] + + # entropy [0, prefix_len-1]: whole prefix copied (label-independent). + if entropy is not None: + entropy[i, 0:prefix_len] = entropy[provider, 0:prefix_len] + + restored_reusers += 1 if ctx.stats is not None: - ctx.stats.record_restore(len(ctx.prefix_last_restore_indices)) + ctx.stats.record_restore(restored_reusers) return output @@ -410,15 +364,13 @@ def restore_via_2d_unfold_verl080( 1. 展开裁剪后 NestedTensor 各行为完整 2D ``[B, L_max]``（reuser prefix 区段 left-pad 0，尾部 right-pad 0 到 L_max） - 2. 构造 ``label_2d``：仅在 prefix-last 列填 ``index.label_value``（reuser - suffix_0 token id）；interior 分支不读 label，suffix 区段不动 - 3. 复用 :func:`restore_reuser_prefix_columns_2d`（现有、已测）：interior 直接 - 复制 provider 2D 值，prefix-last 用存的 logits + label_value 重算 - 4. 按各 ``original_lengths`` 切片压回 NestedTensor (jagged) + 2. 复用 :func:`restore_reuser_prefix_columns_2d`：interior 整段从直接 + provider 的已恢复 2D 行 bulk 切片复制，prefix-last 用存的 logits + + ``index.label_value`` 重算 + 3. 按各 ``original_lengths`` 切片压回 NestedTensor (jagged) - ``valid_indices`` 映射：新方案 left-pad 后 valid-content 0-based 偏移即 2D 列号， - state 不传 ``valid_indices``，``restore_reuser_prefix_columns_2d`` 内的 - ``_map_2d_col`` 自动回退 identity（``return valid_pos``）。 + 列映射为 identity：left-pad 后 valid-content 的 0-based 偏移即 2D 列号， + ``target_2d_pos`` 直接当列索引用，无需 ``valid_indices`` / 列映射表。 Must be called inside ``prefix_sharing_runtime_context`` (reads ``current_prefix_sharing_context``), after the vocab_logprobs patch has saved @@ -438,7 +390,13 @@ def restore_via_2d_unfold_verl080( import torch ctx = current_prefix_sharing_context() - if ctx is None or not ctx.prefix_last_restore_indices: + if ctx is None: + return output + plan = ctx.prefix_sharing_plan + # Guard on reuser presence, not on prefix_last_restore_indices: a batch + # whose reusers all have suffix_len == 0 has no prefix-last spec but still + # needs interior prefix columns restored. + if not plan.has_sharing: return output log_probs_nested = output.get("log_probs") @@ -447,7 +405,6 @@ def restore_via_2d_unfold_verl080( entropy_nested = output.get("entropy") has_entropy = entropy_nested is not None and _is_nested_tensor(entropy_nested) - plan = ctx.prefix_sharing_plan original_lengths = plan.original_lengths input_keep_ranges = plan.input_keep_ranges B = len(original_lengths) @@ -468,15 +425,14 @@ def restore_via_2d_unfold_verl080( ) # --- Step 2: 复用 restore_reuser_prefix_columns_2d --- - # label 由 index.label_value 提供（context 框架无关透传），无需构造 - # [B, L_max] dense label_2d —— 每行至多一个 prefix-last，全表只读几个元素。 - # identity 列映射：target_2d_pos 即 2D 列号（无 left padding）。 + # build_kv 式区间拼接：interior 整段从直接 provider 的已恢复 2D 行切片， + # prefix-last 用 index.label_value + saved logits 重算。identity 列映射 + # （target_2d_pos 即 2D 列号，无 left padding）。 output_2d: dict[str, Any] = {"log_probs": logp_2d} if entropy_2d is not None: output_2d["entropy"] = entropy_2d output_2d = restore_reuser_prefix_columns_2d( output_2d, - None, # label_2d — verl080 用 index.label_value 透传，不构造 dense label 表 vocab_parallel_log_probs_fn, vocab_parallel_entropy_fn, ) @@ -486,12 +442,11 @@ def restore_via_2d_unfold_verl080( if entropy_2d is not None: output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths) - _n_total = len(ctx.prefix_last_restore_indices) - _n_interior = sum(1 for i in ctx.prefix_last_restore_indices if i.is_shared_prefix_interior) + _n_prefix_last = len(ctx.prefix_last_restore_indices) print( f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, " - f"restored {_n_total} indices " - f"({_n_total - _n_interior} prefix-last, {_n_interior} interior)", + f"restored reusers={_n_prefix_last} (prefix-last entries; " + f"interior handled by bulk slice)", flush=True, ) return output diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py index 2192e072..c4dcda99 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py @@ -8,8 +8,9 @@ 仍在、context 激活时，把 prefix-last 重算所需的 provider logits 保存到 ``ctx.prefix_last_logits_saved``，供后续 restore 使用。 -保存条件：仅非 interior（prefix-last）的 index。interior 走 2D 复制路径， -不需要 logits。 +保存条件：``ctx.prefix_last_restore_indices`` 现仅含 prefix-last（每 reuser 一条）， +interior 由 restore 侧 build_kv 式 bulk 切片处理，不读 logits。逐条保存 provider +直接对应位置的 vocab 维 logits。 """ from __future__ import annotations @@ -58,8 +59,6 @@ def patched_fn(logits, labels): flush=True, ) for _idx in ctx.prefix_last_restore_indices: - if _idx.is_shared_prefix_interior: - continue print( f"[PS-diag][packed-align] reuser={_idx.reuse_idx_in_batch} " f"provider={_idx.provider_idx_in_batch} " @@ -70,19 +69,18 @@ def patched_fn(logits, labels): # ##### [PS-diag] 验证 packed 坐标对齐 end ##### for index in ctx.prefix_last_restore_indices: - # interior 走 2D 复制路径，不需要 logits；只保存 prefix-last。 - if index.is_shared_prefix_interior: - continue + # prefix_last_restore_indices 现在只含 prefix-last（interior 由 + # restore 侧 bulk 切片处理），逐条保存其 provider 的 vocab 维 logits。 pos = index.provider_1d_pos key = (index.reuse_idx_in_batch, index.target_2d_pos) if pos < 0: - # 不应再发生：_build_prefix_last_restore_indices 已对 prefix-last - # 二次 strict 解析到 packed 真含 target_pos 的祖先。若到这里说明 - # 解析逻辑有遗漏，直接 raise 暴露，避免下游 restore 静默 KeyError。 + # 不应发生：prefix-last 必落在直接 provider 的 packed 区段内 + # （见 _build_prefix_last_restore_indices 文档）。raise 暴露，避免 + # 下游 restore 静默 KeyError。 raise RuntimeError( - f"[vocab_logprobs] prefix-last spec got provider_1d_pos<0 " - f"after strict resolve; key={key} provider_1d_pos={pos}. " - f"_build_prefix_last_restore_indices 解析逻辑可能有遗漏。" + f"[vocab_logprobs] prefix-last spec got provider_1d_pos<0; " + f"key={key} provider_1d_pos={pos}. " + f"prefix-last 应在直接 provider 的 packed 区段内。" ) # clone 保留 autograd 图（restore 重算 logp 要走反向传播，禁止 detach）。 saved = logits_2d[pos:pos + 1, :].clone() # [1, V//tp] diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 5dd223bf..259ffb86 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -150,6 +150,46 @@ def _load_tensor(dir_path: str, filename: str) -> torch.Tensor | None: return torch.load(fp, weights_only=True).float() if os.path.exists(fp) else None +def _load_manifest(dir_path: str) -> dict | None: + """Load ``parallel_info.json`` written by the dump layer (topology + scopes). + + Returns None when absent (single-card or pre-manifest dumps) → callers fall + back to tp_size==1 behavior (plain filenames, single-card compatible). + """ + fp = os.path.join(dir_path, "parallel_info.json") + if not os.path.exists(fp): + return None + try: + with open(fp, encoding="utf-8") as f: + return json.load(f) + except Exception: + return None + + +def _load_logits(dir_path: str, manifest: dict | None = None) -> torch.Tensor | None: + """Load packed logits, gathering tp vocab shards to full vocab when tp>1. + + tp_size==1 (or no manifest) → single ``logits.pt`` (single-card compatible). + tp_size>1 → concat ``logits_tp{0..tp-1}.pt`` on the vocab (last) dim, + reconstructing ``[N, V]`` so ON-vs-OFF compares on the same + full-vocab coordinate system as single-card. A missing shard + aborts the reconstruction (returns None) rather than silently + comparing partial vocab. + """ + if manifest is None: + manifest = _load_manifest(dir_path) + tp_size = (manifest or {}).get("tp_size", 1) + if tp_size <= 1: + return _load_tensor(dir_path, "logits.pt") + shards = [] + for t in range(tp_size): + s = _load_tensor(dir_path, f"logits_tp{t}.pt") + if s is None: + return None + shards.append(s) + return torch.cat(shards, dim=-1) + + def _load_packed_meta(dir_path: str, cu_fname: str = "cu_seqlens_q.pt") -> dict | None: """加载 cu_seqlens + prefix_lens（suffix 对齐所需）。""" @@ -336,8 +376,8 @@ def cmp_first_token(dir_on: str, dir_off: str) -> list[CheckResult]: name="first_token_attn", metrics=_first_token_metrics(a0[0], b0[0]))) - lo = _load_tensor(dir_on, "logits.pt") - lf = _load_tensor(dir_off, "logits.pt") + lo = _load_logits(dir_on) + lf = _load_logits(dir_off) if lo is not None and lf is not None: lo_first, lf_first = _logits_first_token(lo, lf) results.append(CheckResult( @@ -348,8 +388,8 @@ def cmp_first_token(dir_on: str, dir_off: str) -> list[CheckResult]: def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: """全 packed logits suffix 对齐 + per-token cosine。""" - lo = _load_tensor(dir_on, "logits.pt") - lf = _load_tensor(dir_off, "logits.pt") + lo = _load_logits(dir_on) + lf = _load_logits(dir_off) if lo is None or lf is None: return None lo, lf = _logits_ensure_token_major(lo, lf) @@ -553,7 +593,43 @@ def _shape_of(dir_path: str, filename: str) -> str: return "(error)" -def _print_shapes(dir_on: str, dir_off: str, tag: str): +def _logits_shape(dir_path: str, manifest: dict | None) -> str: + """Shape string for logits, manifest-aware: tp>1 → show shard shape tagged. + + Under TP the file is sharded (``logits_tp{r}.pt``); report one shard's shape + prefixed with ``tp{N}×`` so the shapes table still flags mismatches without + pretending a plain ``logits.pt`` exists. + """ + tp_size = (manifest or {}).get("tp_size", 1) + if tp_size <= 1: + return _shape_of(dir_path, "logits.pt") + s0 = _shape_of(dir_path, "logits_tp0.pt") + if s0 in ("(missing)", "(error)"): + return s0 + return f"tp{tp_size}×{s0}" + + +def _print_topology(manifest_on: dict | None, manifest_off: dict | None) -> None: + """Print ON/OFF parallel topology from manifests; warn on mismatch.""" + def _topo(m): + if not m: + return "single-card (no manifest)" + return f"tp={m.get('tp_size', 1)} pp={m.get('pp_size', 1)} cp={m.get('cp_size', 1)}" + print(_SEP_SINGLE + "\n [topology] ON vs OFF parallel config") + print(_SEP_SINGLE) + print(f" ON : {_topo(manifest_on)}") + print(f" OFF: {_topo(manifest_off)}") + if manifest_on and manifest_off: + for key in ("tp_size", "pp_size", "cp_size"): + if manifest_on.get(key) != manifest_off.get(key): + print(f" {_CROSS} MISMATCH on {key}: ON={manifest_on.get(key)} " + f"OFF={manifest_off.get(key)} — comparison may be invalid") + print() + + +def _print_shapes(dir_on: str, dir_off: str, tag: str, + manifest_on: dict | None = None, + manifest_off: dict | None = None): """打印 ON/OFF 各 .pt 文件 shape —— 定位 shape mismatch 根因的第一手信息。""" print(_SEP_SINGLE + "\n [shapes] ON vs OFF dump shapes") print(_SEP_SINGLE) @@ -570,7 +646,12 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str): print(f" {'FILE':<28s} {'ON':<16s} {'OFF':<16s} {'STATUS'}") print(f" {'─' * 28} {'─' * 16} {'─' * 16} {'─' * 10}") for fname in files: - s_on, s_off = _shape_of(dir_on, fname), _shape_of(dir_off, fname) + if fname == "logits.pt": + # TP-sharded: per-rank logits_tp{r}.pt, not a plain logits.pt + s_on = _logits_shape(dir_on, manifest_on) + s_off = _logits_shape(dir_off, manifest_off) + else: + s_on, s_off = _shape_of(dir_on, fname), _shape_of(dir_off, fname) if s_on == "(missing)" or s_off == "(missing)": status = "—" elif s_on == s_off: @@ -807,8 +888,13 @@ def main(): _print_header(args.dir_on, args.dir_off, args.dir_off2, args.tag, args.mask, args.layer) + # ── parallel topology (manifest-driven: TP shards, future SP/PP) ── + manifest_on = _load_manifest(args.dir_on) + manifest_off = _load_manifest(args.dir_off) + _print_topology(manifest_on, manifest_off) + # ── shape diagnostics ── - _print_shapes(args.dir_on, args.dir_off, args.tag) + _print_shapes(args.dir_on, args.dir_off, args.tag, manifest_on, manifest_off) # ── resolve 2D mask ── ref = _load_tensor(args.dir_off, f"logprobs_{args.tag}.pt") @@ -848,8 +934,8 @@ def main(): a0 = (a.squeeze(1) if a.dim() == 3 else a)[0].cpu() b0 = (b.squeeze(1) if b.dim() == 3 else b)[0].cpu() _print_topk_vec(a0, b0, args.topk, "val", "first_token_attn") - lo = _load_tensor(args.dir_on, "logits.pt") - lf = _load_tensor(args.dir_off, "logits.pt") + lo = _load_logits(args.dir_on) + lf = _load_logits(args.dir_off) if lo is not None and lf is not None: lo_f, lf_f = _logits_first_token(lo, lf) _print_topk_vec(lo_f.cpu(), lf_f.cpu(), args.topk, diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py index 89386ba5..aa0541cc 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py @@ -65,19 +65,123 @@ def _rank0_only() -> bool: return True -# ── Generic helpers ───────────────────────────────────────────── +# ── Parallel-topology-aware dumping (TP / SP / PP) ────────────── +# +# Each dumped tensor has a *rank-scope* describing how (if at all) it is split +# across ranks. The scope decides WHO dumps and under WHAT filename: +# +# "global" — identical on every rank that holds it → dumped once by the +# representative rank (rank-0 under TP/DP; PP will need +# stage-aware routing, left as a per-tensor concern then). +# Plain filename, e.g. ``logprobs_old.pt``. +# "tp_vocab" — vocab-sharded by TP → EVERY tp rank dumps its own shard as +# ``_tp{r}.pt`` (no cross-rank comm on the dump path; +# cmp reassembles offline). tp_size==1 degrades to the plain +# ``logits.pt`` so single-card dumps are unchanged. +# +# Future scopes (reserved, wired when SP/PP land): +# "tp_seq" — sequence-sharded by SP (reduce-scatter) → ``_tp{r}``. +# "pp_stage" — layer-partitioned by PP → ``_pp{r}``. +# +# The static name→scope map is written to ``parallel_info.json`` (manifest) so +# cmp_diag knows how to gather each tensor without guessing. + +_TENSOR_SCOPES: dict[str, str] = { + "logits": "tp_vocab", # packed logits [N, V//tp] — vocab-sharded under TP +} + +_PARALLEL_INFO_CACHE: Any = None +_MANIFEST_WRITTEN: set[str] = set() + + +def _cached_parallel_info() -> Any: + """Read MegatronParallelInfo once and cache (tp/pp/cp ranks + sizes). + + Returns None when distributed/Megatron isn't initialized (single-rank), + which the callers treat as the single-card case. + """ + global _PARALLEL_INFO_CACHE + if _PARALLEL_INFO_CACHE is not None: + return _PARALLEL_INFO_CACHE + try: + from prefix_sharing.integrations.parallel_info import ( + get_megatron_parallel_info, + ) + _PARALLEL_INFO_CACHE = get_megatron_parallel_info() + except Exception: + _PARALLEL_INFO_CACHE = None + return _PARALLEL_INFO_CACHE + + +def _with_suffix(name: str, suffix: str) -> str: + """Insert a rank suffix before the extension: logits.pt → logits_tp0.pt.""" + if not suffix: + return name + stem, sep, ext = name.rpartition(".") + return f"{stem}{suffix}{sep}{ext}" if sep else f"{name}{suffix}" -def _save_tensor(name: str, tensor: torch.Tensor, dump_dir: str) -> bool: - """Save a tensor to dump_dir/name.pt, rank-0-only. Returns True on success.""" + +def _ensure_manifest(dump_dir: str, pi: Any) -> None: + """Write parallel_info.json once from rank-0: topology + scope map. + + Content is global/identical, so a single writer avoids filesystem races. + Read by cmp_diag to gather per-rank files (e.g. concat logits_tp{0..tp-1}.pt + back to full vocab) and to validate ON/OFF topology match. + """ + if dump_dir in _MANIFEST_WRITTEN: + return + _MANIFEST_WRITTEN.add(dump_dir) if not _rank0_only(): - return False + return + import json + manifest = { + "tp_size": getattr(pi, "tp_size", 1) if pi else 1, + "pp_size": getattr(pi, "pp_size", 1) if pi else 1, + "cp_size": getattr(pi, "cp_size", 1) if pi else 1, + "global_rank_of_dumper": getattr(pi, "global_rank", 0) if pi else 0, + "scopes": dict(_TENSOR_SCOPES), + } + try: + with open(os.path.join(dump_dir, "parallel_info.json"), + "w", encoding="utf-8") as f: + json.dump(manifest, f, indent=2, ensure_ascii=False) + _log.warning("parallel_info.json saved (tp=%d pp=%d)", + manifest["tp_size"], manifest["pp_size"]) + except Exception as e: + _log.warning("parallel_info.json save failed: %s", e) + + +# ── Generic helpers ───────────────────────────────────────────── + +def _save_tensor(name: str, tensor: torch.Tensor, dump_dir: str, + scope: str = "global") -> bool: + """Save a tensor to ``dump_dir/``, scope-aware. Returns True on success. + + See the ``_TENSOR_SCOPES`` block above for the scope semantics. ``scope`` + defaults to ``"global"`` (rank-0-only, plain filename) so existing callers + are unchanged; per-rank tensors opt in via ``scope="tp_vocab"`` (etc.). + """ + pi = _cached_parallel_info() + _ensure_manifest(dump_dir, pi) + + if scope == "tp_vocab" and pi is not None and pi.tp_size > 1: + # Every tp rank dumps its own vocab shard; no rank-0 gate, no comm. + fname = _with_suffix(name, f"_tp{pi.tp_rank}") + else: + # global scope, OR tp_vocab with tp_size==1 (single shard == full): + # rank-0 dumps once under the plain name (single-card compatible). + if scope == "tp_vocab": + scope = "global" # tp==1 → behaves as global for logging + if not _rank0_only(): + return False + fname = name try: - path = os.path.join(dump_dir, name) + path = os.path.join(dump_dir, fname) torch.save(tensor.detach().cpu().clone(), path) - _log.warning("%s saved (%s)", name, tensor.shape) + _log.warning("%s saved (%s, scope=%s)", fname, tensor.shape, scope) return True except Exception as e: - _log.warning("%s save failed: %s", name, e) + _log.warning("%s save failed: %s", fname, e) return False diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index e6c720f2..19838145 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -145,11 +145,17 @@ def dump_meta_verl080(prefix_lens: list[int], cu_seqlens: torch.Tensor) -> None: def dump_logits_verl080(logits: torch.Tensor) -> None: - """存 packed logits ``[N, V//tp]``（vocab 在最后一维，v070 约定）。""" + """存 packed logits，scope=``tp_vocab``：每个 tp rank 存自己的词表片。 + + 纯 TP 下 logits 是 ``[N, V//tp]``，各 tp rank 不同（vocab 切分）。每个 tp + rank 存自己的 shard（文件名 ``logits_tp{r}.pt``；``tp_size==1`` 时退化为 + ``logits.pt``，单卡零改动）。dump 路径不引入跨 rank 通信；cmp 侧按 + ``parallel_info.json`` 把 ``_tp{0..tp-1}`` 拼回 full vocab 再比。 + """ dump_dir = _get_dump_dir() if dump_dir is None: return - _save_tensor("logits.pt", logits, dump_dir) + _save_tensor("logits.pt", logits, dump_dir, scope="tp_vocab") def dump_logprobs_2d_verl080(logp_2d: torch.Tensor, tag: str) -> None: diff --git a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py index 20b51168..05d2f5c4 100644 --- a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py +++ b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py @@ -80,11 +80,9 @@ def test_build_prefix_sharing_micro_batch_verl070_trims_reuser_mask_and_context_ with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: assert current_prefix_sharing_context() is ctx - # 3 restore indices: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 0 # interior pos1 - assert ctx.prefix_last_restore_indices[2].provider_1d_pos == 2 # prefix-last - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel + # 1 prefix-last index (interior is bulk-sliced, not indexed) + assert len(ctx.prefix_last_restore_indices) == 1 + assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 2 # prefix-last, provider seq0 offset 2 assert current_prefix_sharing_context() is None @@ -153,20 +151,16 @@ def test_build_prefix_sharing_micro_batch_verl070_builds_common_tp_padded_layout assert layout.packed_position_ids.tolist() == expected_positions assert layout.valid_token_mask.tolist() == expected_mask with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: - # 3 restore indices: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 0 # interior pos1 - assert ctx.prefix_last_restore_indices[2].provider_1d_pos == 2 # prefix-last - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel + # 1 prefix-last index (interior is bulk-sliced, not indexed) + assert len(ctx.prefix_last_restore_indices) == 1 + assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 2 # prefix-last, provider seq0 offset 2 def test_restore_reuser_prefix_columns_2d_prefix_last_keeps_autograd(): - # prefix-last-only case (no interior response). - # input: [[1,2,3,10,11], [1,2,3,20,21]] → prefix_len=3, prompt_len=3 - # Planner emits 3 specs (2 interior + 1 prefix-last); here we focus on the - # prefix-last spec at index 2: provider_predict_pos=2, target_2d_pos=2. - # Reuser's first suffix token at target_2d_pos=2 differs from provider's, - # so logprob must be recomputed from saved provider logits. + # input: [[1,2,3,10,11], [1,2,3,20,21]] → prefix_len=3, prompt_len=3. + # Only the prefix-last spec is indexed (interior is bulk-sliced in the + # restore). Reuser's first suffix token at target_2d_pos=2 differs from + # provider's, so its logprob must be recomputed from saved provider logits. batch = { "input_ids": torch.tensor([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21]]), "attention_mask": torch.ones(2, 5, dtype=torch.bool), @@ -196,24 +190,20 @@ def gather_fn(provider_logits, reuse_label): ).squeeze(-1) with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: - # 3 restore specs: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - index = ctx.prefix_last_restore_indices[2] # prefix-last spec - assert not index.is_shared_prefix_interior + # 1 prefix-last spec (interior bulk-sliced, not indexed) + assert len(ctx.prefix_last_restore_indices) == 1 + index = ctx.prefix_last_restore_indices[0] # prefix-last spec # Simulate 2D postprocess: output dict with [B, L] log_probs. log_probs_2d = torch.zeros(2, 5) output = {"log_probs": log_probs_2d} - # 2D label: label[p] = token at p+1 (verl convention). - label_2d = torch.tensor([[0, 0, 0, 10, 11], [0, 0, 0, 20, 21]]) - # Saved provider packed logits for prefix-last recompute. saved_logits = torch.randn(1, 32, requires_grad=True) ctx.prefix_last_logits_saved[(index.reuse_idx_in_batch, index.target_2d_pos)] = saved_logits - output = restore_reuser_prefix_columns_2d(output, label_2d, gather_fn) - assert ctx.stats.actual_restore_count == ctx.stats.expected_restore_count == 3 + output = restore_reuser_prefix_columns_2d(output, gather_fn) + assert ctx.stats.actual_restore_count == ctx.stats.expected_restore_count == 1 restored_val = output["log_probs"][index.reuse_idx_in_batch, index.target_2d_pos] @@ -267,10 +257,8 @@ def test_build_prefix_sharing_micro_batch_verl070_keeps_global_layout_with_seque assert layout.cu_seqlens == expected_cu_seqlens assert layout.total_padded_length == expected_cu_seqlens[-1] with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: - # 3 restore specs: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 0 # interior pos1 - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel + # 1 prefix-last index (interior bulk-sliced, not indexed) + assert len(ctx.prefix_last_restore_indices) == 1 @pytest.mark.parametrize("pp_size", [2, 4, 8]) @@ -349,9 +337,8 @@ def test_build_prefix_sharing_micro_batch_verl070_combines_tp_padding_with_physi assert prefix_sharing_runtime_state.parallel_info.pp_size == pp_size assert prefix_sharing_runtime_state.packed_batch_layout.cu_seqlens == expected_cu_seqlens with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: - # 3 restore specs: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel + # 1 prefix-last index (interior bulk-sliced, not indexed) + assert len(ctx.prefix_last_restore_indices) == 1 def test_attention_hook_rejects_sp_local_shard_token_length(): diff --git a/prefix-sharing/tests/unit_test/test_observability.py b/prefix-sharing/tests/unit_test/test_observability.py index c6aa53b9..e181c7a3 100644 --- a/prefix-sharing/tests/unit_test/test_observability.py +++ b/prefix-sharing/tests/unit_test/test_observability.py @@ -33,7 +33,9 @@ def test_prefix_sharing_stats_from_plan_records_expected_reuse_summary(): assert stats.sharing_group_count == 1 assert stats.expected_reused_counts_per_layer == 1 assert stats.expected_reused_prefix_tokens_per_layer == 3 - assert stats.expected_restore_count == 3 + # expected_restore_count now counts reuser rows (interior is bulk-sliced, + # no longer one-index-per-position), so it equals the reuser count. + assert stats.expected_restore_count == 1 assert stats.actual_restore_count == 0 diff --git a/prefix-sharing/tests/unit_test/test_planner.py b/prefix-sharing/tests/unit_test/test_planner.py index cb7e1d58..f9f46589 100644 --- a/prefix-sharing/tests/unit_test/test_planner.py +++ b/prefix-sharing/tests/unit_test/test_planner.py @@ -35,31 +35,26 @@ def test_planner_builds_phase_one_prefix_sharing_plan_and_restore_specs(): assert prefix_sharing_plan.q_position_offsets == [0, 3, 5, 0] assert prefix_sharing_plan.kv_position_offsets == [0, 0, 0, 0] - # prefix_last_restore: reuser Q skips prefix; ALL prefix columns need restore. - # Row 1 (prefix_len=3): interior positions 1..2 (2 specs) + prefix-last (1 spec) = 3 - # Row 2 (prefix_len=5): interior positions 1..4 (4 specs) + prefix-last (1 spec) = 5 - # Total: 3 + 5 = 8 specs. + # prefix_last_restore: reuser Q skips prefix; only the prefix-last token + # needs a spec (interior columns are bulk-sliced by the 2D restore). + # Row 1 (prefix_len=3): 1 prefix-last spec. + # Row 2 (prefix_len=5): 1 prefix-last spec. + # Total: 2 specs. all_specs = prefix_sharing_plan.prefix_last_restore - assert len(all_specs) == 8 + assert len(all_specs) == 2 - # Row 1 prefix-last spec (index 2: after 2 interior specs) - spec1 = all_specs[2] - assert not spec1.is_shared_prefix_interior + # Row 1 prefix-last spec + spec1 = all_specs[0] assert spec1.reuse_idx_in_batch == 1 assert spec1.provider_idx_in_batch == 0 - assert spec1.provider_predict_pos == 2 # prefix_len - 1 = 2 - assert spec1.reuse_first_suffix_label_pos == 3 # prefix_len = 3 - assert spec1.target_2d_pos == 2 + assert spec1.target_2d_pos == 2 # prefix_len - 1 assert spec1.label_value == 20 # input_ids[1][3], the first suffix token - # Row 2 prefix-last spec (index 7: after 4 interior specs) - spec2 = all_specs[7] - assert not spec2.is_shared_prefix_interior + # Row 2 prefix-last spec + spec2 = all_specs[1] assert spec2.reuse_idx_in_batch == 2 assert spec2.provider_idx_in_batch == 0 - assert spec2.provider_predict_pos == 4 # prefix_len - 1 = 4 - assert spec2.reuse_first_suffix_label_pos == 5 # prefix_len = 5 - assert spec2.target_2d_pos == 4 + assert spec2.target_2d_pos == 4 # prefix_len - 1 assert spec2.label_value == 30 # input_ids[2][5], the first suffix token @@ -74,11 +69,12 @@ def test_planner_no_shared_prefix_keeps_original_shapes(): assert prefix_sharing_plan.prefix_last_restore == [] -def test_planner_generates_interior_response_restore_specs(): +def test_planner_emits_only_prefix_last_not_interior(): # seq1 (provider): [1,2,3 | A,B,C] prompt=[1,2,3], response=[A,B,C] # seq2 (reuser): [1,2,3 | A,D,E] prompt=[1,2,3], response=[A,D,E] - # Shared prefix: [1,2,3,A] (len 4). A is a response token in both, - # so it needs shared-prefix interior logprob restore. + # Shared prefix: [1,2,3,A] (len 4). A is a response token in both, so the + # prefix spans response positions — but interior columns are now bulk-sliced + # by the 2D restore; the planner only emits the single prefix-last spec. input_ids = [ [1, 2, 3, 4, 5, 6], # 1,2,3,A,B,C [1, 2, 3, 4, 7, 8], # 1,2,3,A,D,E @@ -88,30 +84,14 @@ def test_planner_generates_interior_response_restore_specs(): ) plan = planner.plan(input_ids, forward_id=1, micro_batch_id=1) - # Interior restore covers all prefix columns 1..prefix_len-1: - # prefix_label_pos=1: provider_predict_pos=0, target_2d_pos=0 - # prefix_label_pos=2: provider_predict_pos=1, target_2d_pos=1 - # prefix_label_pos=3: provider_predict_pos=2, target_2d_pos=2 - # + prefix-last - # Total: 4 specs. - assert len(plan.prefix_last_restore) == 4 - - interior_spec = plan.prefix_last_restore[2] # prefix_label_pos=3 (prompt_len area, was the old single interior) - assert interior_spec.is_shared_prefix_interior - assert interior_spec.reuse_idx_in_batch == 1 - assert interior_spec.provider_idx_in_batch == 0 - assert interior_spec.provider_predict_pos == 2 # logits[2] - assert interior_spec.reuse_first_suffix_label_pos == 3 # label pos 3 = A - assert interior_spec.target_2d_pos == 2 # label position prefix_label_pos-1 - assert interior_spec.label_value == 4 # input_ids[1][3] = token A - - prefix_last_spec = plan.prefix_last_restore[3] - assert not prefix_last_spec.is_shared_prefix_interior + # Only the prefix-last spec is emitted (interior is handled by the restore). + # Total: 1 spec. + assert len(plan.prefix_last_restore) == 1 + + prefix_last_spec = plan.prefix_last_restore[0] assert prefix_last_spec.reuse_idx_in_batch == 1 assert prefix_last_spec.provider_idx_in_batch == 0 - assert prefix_last_spec.provider_predict_pos == 3 # logits[3] = prefix-last - assert prefix_last_spec.reuse_first_suffix_label_pos == 4 # label pos 4 = D - assert prefix_last_spec.target_2d_pos == 3 # prefix_len-1 label position + assert prefix_last_spec.target_2d_pos == 3 # prefix_len - 1 assert prefix_last_spec.label_value == 7 # input_ids[1][4] = first suffix token D # Trimming still at prefix_len @@ -119,8 +99,8 @@ def test_planner_generates_interior_response_restore_specs(): assert plan.kept_lengths_q == [6, 2] -def test_planner_generates_interior_restore_with_minimal_args(): - """Shared-prefix interior restore covers all prefix columns using only input_ids.""" +def test_planner_emits_prefix_last_with_minimal_args(): + """Only the prefix-last spec is emitted; interior needs no spec.""" input_ids = [ [1, 2, 3, 4, 5], [1, 2, 3, 4, 6], @@ -130,8 +110,7 @@ def test_planner_generates_interior_restore_with_minimal_args(): ) plan = planner.plan(input_ids, forward_id=1, micro_batch_id=1) - # Interior positions 1..3 (3 specs) + prefix-last (1 spec) = 4 specs - assert len(plan.prefix_last_restore) == 4 - # prefix-last spec at index 3 (last one) - spec = plan.prefix_last_restore[3] - assert not spec.is_shared_prefix_interior + # Only the prefix-last spec is emitted (interior handled by the restore). + assert len(plan.prefix_last_restore) == 1 + spec = plan.prefix_last_restore[0] + assert spec.reuse_idx_in_batch == 1 diff --git a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py index 2900ea6d..38e2dd09 100644 --- a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py +++ b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py @@ -76,19 +76,18 @@ def test_non_nested_log_probs_returns_unchanged(): assert torch.equal(result["log_probs"], original) -def test_empty_restore_indices_returns_unchanged(): - """ctx.prefix_last_restore_indices 为空时 early return（无 restore 需求）。""" - state = _make_state([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21]]) +def test_no_sharing_returns_output_unchanged(): + """plan 无 reuser（has_sharing=False）时 early return，输出不变。""" + # 两条无公共前缀的序列 → 无 reuse → has_sharing=False + state = _make_state([[1, 2, 3, 10, 11], [4, 5, 6, 7, 8]]) nested = torch.nested.nested_tensor( [torch.tensor([1.0, 2.0]), torch.tensor([3.0])], layout=torch.jagged ) output = {"log_probs": nested} with prefix_sharing_runtime_context(state) as ctx: - saved_indices = ctx.prefix_last_restore_indices[:] - ctx.prefix_last_restore_indices.clear() + assert not ctx.prefix_sharing_plan.has_sharing result = restore_via_2d_unfold_verl080(output, _mock_vocab_log_probs_fn) assert result is output - ctx.prefix_last_restore_indices.extend(saved_indices) # ═══════════════════════════════════════ @@ -176,11 +175,9 @@ def test_restore_copies_interior_and_recomputes_prefix_last(): output = {"log_probs": nested_logp} with prefix_sharing_runtime_context(state) as ctx: - assert len(ctx.prefix_last_restore_indices) == 3 - prefix_last_idx = [ - i for i in ctx.prefix_last_restore_indices - if not i.is_shared_prefix_interior - ][0] + # 仅 1 条 prefix-last（interior 由 restore 侧 bulk 切片处理，不建索引） + assert len(ctx.prefix_last_restore_indices) == 1 + prefix_last_idx = ctx.prefix_last_restore_indices[0] assert prefix_last_idx.target_2d_pos == 2 assert prefix_last_idx.label_value == 20 # input_ids[1][3] @@ -229,10 +226,7 @@ def test_restore_with_entropy_copies_both_logp_and_entropy(): output = {"log_probs": nested_logp, "entropy": nested_ent} with prefix_sharing_runtime_context(state) as ctx: - prefix_last_idx = [ - i for i in ctx.prefix_last_restore_indices - if not i.is_shared_prefix_interior - ][0] + prefix_last_idx = ctx.prefix_last_restore_indices[0] ctx.prefix_last_logits_saved[ (prefix_last_idx.reuse_idx_in_batch, prefix_last_idx.target_2d_pos) ] = torch.tensor([[0.5, 0.3, 0.1, 0.1]]) @@ -265,10 +259,7 @@ def test_restore_clears_saved_logits_is_callers_responsibility(): output = {"log_probs": nested_logp} with prefix_sharing_runtime_context(state) as ctx: - prefix_last_idx = [ - i for i in ctx.prefix_last_restore_indices - if not i.is_shared_prefix_interior - ][0] + prefix_last_idx = ctx.prefix_last_restore_indices[0] key = (prefix_last_idx.reuse_idx_in_batch, prefix_last_idx.target_2d_pos) ctx.prefix_last_logits_saved[key] = torch.tensor([[0.5, 0.3, 0.1, 0.1]]) restore_via_2d_unfold_verl080(output, _mock_vocab_log_probs_fn) diff --git a/prefix-sharing/tests/unit_test/test_runtime_context.py b/prefix-sharing/tests/unit_test/test_runtime_context.py index 1de2ae92..89099073 100644 --- a/prefix-sharing/tests/unit_test/test_runtime_context.py +++ b/prefix-sharing/tests/unit_test/test_runtime_context.py @@ -27,11 +27,11 @@ def test_prefix_sharing_runtime_context_sets_and_clears_current_context(): with prefix_sharing_runtime_context(prefix_sharing_runtime_state) as ctx: assert current_prefix_sharing_context() is ctx assert ctx.prefix_sharing_plan is prefix_sharing_runtime_state.prefix_sharing_plan - # 3 restore specs: 2 interior (provider_predict_pos=0,1) + 1 prefix-last (pos=2) - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 0 # interior pos1 - assert ctx.prefix_last_restore_indices[2].provider_1d_pos == 2 # prefix-last - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel: no slot in reuser packed region + # 1 prefix-last spec (interior is bulk-sliced in the restore, not indexed). + assert len(ctx.prefix_last_restore_indices) == 1 + assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 2 # prefix-last, direct provider seq0 offset 2 + assert ctx.prefix_last_restore_indices[0].target_2d_pos == 2 + assert ctx.prefix_last_restore_indices[0].label_value == 20 # seq1[3] assert ctx.parallel_info is prefix_sharing_runtime_state.parallel_info assert ctx.parallel_info.pp_rank == 1 assert ctx.parallel_info.pp_size == 2 @@ -64,33 +64,27 @@ def test_prefix_sharing_runtime_context_uses_padded_layout_for_restore_indices() ) with prefix_sharing_runtime_context(runtime_state) as ctx: - # 3 restore specs: 2 interior + 1 prefix-last - assert len(ctx.prefix_last_restore_indices) == 3 - assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 0 # interior pos1 - assert ctx.prefix_last_restore_indices[2].provider_1d_pos == 2 # prefix-last - assert ctx.prefix_last_restore_indices[0].reuse_1d_pos == -1 # sentinel: no slot in reuser packed region + # 1 prefix-last spec (interior bulk-sliced, not indexed). + assert len(ctx.prefix_last_restore_indices) == 1 + assert ctx.prefix_last_restore_indices[0].provider_1d_pos == 2 # prefix-last, direct provider seq0 offset 2 assert ctx.stats.kept_padded_tokens == 8 -def test_chain_reuse_resolves_position_to_provider_with_matching_label(): - """Chain-reuse must route a deeper reuser to the provider whose - *continuation* matches, not the root. +def test_chain_reuse_prefix_last_resolves_to_direct_provider(): + """Chain reuse: a reuser's prefix-last index points at its **direct** + provider (not the root) and resolves to a real packed slot there. - Trie detection produces a natural chain here: + Trie detection produces a natural chain: seq0: [1, 2, 3, 4, 5] provider (root) seq1: [1, 2, 3, 7, 8] reuse(seq0), shared prefix [1,2,3] (len 3) seq2: [1, 2, 3, 7, 9] reuse(seq1), shared prefix [1,2,3,7] (len 4) - The trie routes seq2 to seq1 (not seq0): at position 3 both seq1 and seq2 - hold token 7, whereas seq0 holds 4. So seq2's restore for position 2 - (token "3", which predicts token 7) must reference seq1 — whose label at - that position is 7 — rather than seq0, whose label there is 4. - - The keep_start-1 extension in ``_resolve_provider_for_position`` keeps the - boundary position on the direct provider seq1 (seq1 keeps [3, 5), so - keep_start-1 = 2 covers position 2). Positions shared identically with - the root (0, 1) walk up to seq0. The prefix-last (position 3) lands on a - real packed slot of seq1 because seq1's kept range [3, 5) contains it. + seq2's prefix-last is position 3 (predicts token 9). Chain reuse forces + prefix_len_reuser (4) > prefix_len_provider (3), so position 3 lands at + seq1's keep_start — inside seq1's computed suffix region. The prefix-last + logits therefore live on the direct provider seq1 and need no walk up to + seq0. Interior positions are no longer indexed: the restore bulk-slices + them from the provider's already-restored 2D row. """ planner = PrefixSharingPlanner(PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)) plan = planner.plan( @@ -109,36 +103,26 @@ def test_chain_reuse_resolves_position_to_provider_with_matching_label(): parallel_info=MegatronParallelInfo(), ) with prefix_sharing_runtime_context(runtime_state) as ctx: - seq2 = { - idx.target_2d_pos: idx - for idx in ctx.prefix_last_restore_indices - if idx.reuse_idx_in_batch == 2 + # One prefix-last index per reuser-with-suffix (seq1, seq2); interior + # is no longer indexed. + by_row = { + idx.reuse_idx_in_batch: idx for idx in ctx.prefix_last_restore_indices } + assert set(by_row) == {1, 2} - # Positions 0, 1 (tokens 1, 2): labels 2, 3 are identical across all - # rows, so they resolve up the chain to the root seq0. - assert seq2[0].provider_idx_in_batch == 0 - assert seq2[0].label_value == 2 - assert seq2[1].provider_idx_in_batch == 0 - assert seq2[1].label_value == 3 - - # Position 2 (token "3"): seq2/seq1 predict 7 here, seq0 predicts 4. - # Must resolve to the direct provider seq1 (matching label 7), and - # stay there via the keep_start-1 extension — NOT walk up to seq0. - assert seq2[2].provider_idx_in_batch == 1, ( - "seq2 position 2 (token 3 -> label 7) must resolve to seq1 " - "(matching continuation 7), not seq0 (continuation 4)" - ) - assert seq2[2].label_value == 7 + # seq1 prefix-last: position 2 (predicts token 7), direct provider seq0. + assert by_row[1].provider_idx_in_batch == 0 + assert by_row[1].target_2d_pos == 2 + assert by_row[1].label_value == 7 + assert by_row[1].provider_1d_pos != -1 - # Prefix-last (position 3, predicts token 9): seq1 keeps [3, 5) which - # strictly contains position 3, so it resolves to a real packed slot - # on seq1 (not the -1 sentinel). - plast = seq2[3] - assert plast.is_shared_prefix_interior is False + # seq2 prefix-last: position 3 (predicts token 9), direct provider seq1 + # (NOT the root seq0) — the key chain-reuse property. + plast = by_row[2] assert plast.provider_idx_in_batch == 1 + assert plast.target_2d_pos == 3 assert plast.label_value == 9 assert plast.provider_1d_pos != -1, ( - "prefix-last must resolve to a real packed slot on the direct " - "provider seq1, not the -1 sentinel" + "seq2 prefix-last must resolve to a real packed slot on its direct " + "provider seq1, not walk up to the root seq0" ) From 3260dc6b6b3c4381694857e7d5eef2fe1764a45b Mon Sep 17 00:00:00 2001 From: Boundless Date: Thu, 25 Jun 2026 20:55:04 +0800 Subject: [PATCH 03/61] =?UTF-8?q?[refactor]=202D=E6=B3=A8=E5=85=A5?= =?UTF-8?q?=E6=B5=81=E7=A8=8B=E5=91=BD=E5=90=8D=E4=BC=98=E5=8C=96=E4=B8=8E?= =?UTF-8?q?docstring=E6=94=B6=E6=95=9B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit restore循环：i→reuser_idx、provider→provider_idx、idx→prefix_last_spec、prefix_last_by_row→prefix_last_spec_by_reuser、saved_key→saved_logits_key、provider_logits→saved_provider_logits；unfold/fold循环：row→seq_idx，lp_/ent_/logp_等隐晦缩写展开为log_probs_/entropy_；vocab_logprobs patch docstring收敛为prefix-last单语义。纯命名/注释调整，零行为变化。 --- .../prefix_sharing/integrations/verl_mcore.py | 104 +++++++++--------- .../vocab_logprobs.py | 10 +- 2 files changed, 57 insertions(+), 57 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index ef6490c4..7412bc43 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -246,21 +246,22 @@ def restore_reuser_prefix_columns_2d( off the **direct provider's already-restored 2D row** and only the single prefix-last logprob is recomputed. - Per reuser row ``i`` with direct provider ``p = provider_index[i]`` and - ``P = prefix_lens[i]`` (columns are identity-mapped in the unfolded 2D - tensor, so ``target_2d_pos`` == column): + Per ``reuser_idx`` with direct provider ``provider_idx = provider_index[reuser_idx]`` + and ``P = prefix_lens[reuser_idx]`` (columns are identity-mapped in the + unfolded 2D tensor, so ``target_2d_pos`` == column): - - **interior ``[0, P-2]``**: ``log_probs[i, 0:P-1] = log_probs[p, 0:P-1]`` + - **interior ``[0, P-2]``**: ``log_probs[reuser_idx, 0:P-1] = log_probs[provider_idx, 0:P-1]`` (bulk copy). Identical across the shared prefix (same logits + labels), - and ``p`` was restored earlier in the batch-order loop, so its row already - holds correct values — no per-position provider resolution needed. - - **prefix-last ``P-1``**: recompute ``log_probs[i, P-1]`` from the saved - provider logits + the reuser's own first-suffix label (differs from the - provider's). When the reuser has no suffix (``suffix_len == 0``) the + and ``provider_idx`` was restored earlier in the batch-order loop, so its + row already holds correct values — no per-position provider resolution + needed. + - **prefix-last ``P-1``**: recompute ``log_probs[reuser_idx, P-1]`` from the + saved provider logits + the reuser's own first-suffix label (differs from + the provider's). When the reuser has no suffix (``suffix_len == 0``) the planner emits no prefix-last spec; that column is masked downstream, so the provider's value is copied as a safe placeholder. - - **entropy ``[0, P-1]``**: ``entropy[i, 0:P] = entropy[p, 0:P]`` (whole - prefix copied, including prefix-last — entropy is label-independent). + - **entropy ``[0, P-1]``**: ``entropy[reuser_idx, 0:P] = entropy[provider_idx, 0:P]`` + (whole prefix copied, including prefix-last — entropy is label-independent). Rows are visited in ``range(B)`` order so a provider is always restored before any reuser that reads it (the same online-detector invariant @@ -299,45 +300,44 @@ def restore_reuser_prefix_columns_2d( provider_index = plan.provider_index prefix_lens = plan.prefix_lens - # prefix-last lookup keyed by reuser row. ``prefix_last_restore_indices`` - # now carries only prefix-last entries (one per reuser-with-suffix); - # interior is handled by the bulk slice below. - prefix_last_by_row = { - idx.reuse_idx_in_batch: idx for idx in ctx.prefix_last_restore_indices + # reuser row → its prefix-last restore spec (one per reuser-with-suffix; + # interior positions have no spec — they are bulk-sliced below). + prefix_last_spec_by_reuser = { + spec.reuse_idx_in_batch: spec for spec in ctx.prefix_last_restore_indices } restored_reusers = 0 # Row 0 is always a provider (nothing precedes it to reuse), so start at 1. # A reuser's provider always has a smaller batch index (online-detector - # invariant), so it is already restored when we reach row i. - for i in range(1, len(prefix_lens)): - prefix_len = prefix_lens[i] - if provider_index[i] == i or prefix_len <= 0: + # invariant), so it is already restored when we reach reuser_idx. + for reuser_idx in range(1, len(prefix_lens)): + prefix_len = prefix_lens[reuser_idx] + if provider_index[reuser_idx] == reuser_idx or prefix_len <= 0: continue # provider / non-reuser: row already complete - provider = provider_index[i] + provider_idx = provider_index[reuser_idx] # interior [0, prefix_len-2]: bulk-copy from the provider's restored row. if prefix_len - 1 > 0: - log_probs[i, 0:prefix_len - 1] = log_probs[provider, 0:prefix_len - 1] + log_probs[reuser_idx, 0:prefix_len - 1] = log_probs[provider_idx, 0:prefix_len - 1] # prefix-last (position prefix_len-1): recompute with the reuser's label. - idx = prefix_last_by_row.get(i) - if idx is not None: - saved_key = (i, idx.target_2d_pos) - provider_logits = ctx.prefix_last_logits_saved[saved_key] # [1, V//tp] + prefix_last_spec = prefix_last_spec_by_reuser.get(reuser_idx) + if prefix_last_spec is not None: + saved_logits_key = (reuser_idx, prefix_last_spec.target_2d_pos) + saved_provider_logits = ctx.prefix_last_logits_saved[saved_logits_key] # [1, V//tp] reuser_label = torch.tensor( - [idx.label_value], dtype=torch.long, device=log_probs.device, + [prefix_last_spec.label_value], dtype=torch.long, device=log_probs.device, ) # [1] - log_probs[i, prefix_len - 1] = vocab_parallel_log_probs_fn( - provider_logits, reuser_label, + log_probs[reuser_idx, prefix_len - 1] = vocab_parallel_log_probs_fn( + saved_provider_logits, reuser_label, ).reshape(()) else: # suffix_len == 0: no prefix-last spec; column is masked downstream. - log_probs[i, prefix_len - 1] = log_probs[provider, prefix_len - 1] + log_probs[reuser_idx, prefix_len - 1] = log_probs[provider_idx, prefix_len - 1] # entropy [0, prefix_len-1]: whole prefix copied (label-independent). if entropy is not None: - entropy[i, 0:prefix_len] = entropy[provider, 0:prefix_len] + entropy[reuser_idx, 0:prefix_len] = entropy[provider_idx, 0:prefix_len] restored_reusers += 1 @@ -415,7 +415,7 @@ def restore_via_2d_unfold_verl080( device = log_probs_nested.values().device # --- Step 1: 展开裁剪后 NestedTensor → 完整 2D [B, L_max] --- - logp_2d, entropy_2d = _unfold_trimmed_nested_to_2d( + log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d( log_probs_nested, entropy_nested if has_entropy else None, original_lengths, @@ -428,7 +428,7 @@ def restore_via_2d_unfold_verl080( # build_kv 式区间拼接：interior 整段从直接 provider 的已恢复 2D 行切片， # prefix-last 用 index.label_value + saved logits 重算。identity 列映射 # （target_2d_pos 即 2D 列号，无 left padding）。 - output_2d: dict[str, Any] = {"log_probs": logp_2d} + output_2d: dict[str, Any] = {"log_probs": log_probs_2d} if entropy_2d is not None: output_2d["entropy"] = entropy_2d output_2d = restore_reuser_prefix_columns_2d( @@ -442,10 +442,10 @@ def restore_via_2d_unfold_verl080( if entropy_2d is not None: output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths) - _n_prefix_last = len(ctx.prefix_last_restore_indices) + num_prefix_last = len(ctx.prefix_last_restore_indices) print( f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, " - f"restored reusers={_n_prefix_last} (prefix-last entries; " + f"restored reusers={num_prefix_last} (prefix-last entries; " f"interior handled by bulk slice)", flush=True, ) @@ -472,29 +472,29 @@ def _unfold_trimmed_nested_to_2d( """ import torch - lp_offsets = log_probs_nested.offsets() - lp_values = log_probs_nested.values() + log_probs_offsets = log_probs_nested.offsets() + log_probs_values = log_probs_nested.values() if entropy_nested is not None: - ent_offsets = entropy_nested.offsets() - ent_values = entropy_nested.values() + entropy_offsets = entropy_nested.offsets() + entropy_values = entropy_nested.values() - logp_rows: list[Any] = [] - ent_rows: list[Any] | None = [] if entropy_nested is not None else None + log_probs_rows: list[Any] = [] + entropy_rows: list[Any] | None = [] if entropy_nested is not None else None - for i in range(B): - orig_len = original_lengths[i] - prefix_len = input_keep_ranges[i][0] + for seq_idx in range(B): + orig_len = original_lengths[seq_idx] + prefix_len = input_keep_ranges[seq_idx][0] - lp_suffix = lp_values[lp_offsets[i]:lp_offsets[i + 1]] - logp_rows.append(_build_padded_row(lp_suffix, prefix_len, orig_len, L_max)) + log_probs_suffix = log_probs_values[log_probs_offsets[seq_idx]:log_probs_offsets[seq_idx + 1]] + log_probs_rows.append(_build_padded_row(log_probs_suffix, prefix_len, orig_len, L_max)) if entropy_nested is not None: - ent_suffix = ent_values[ent_offsets[i]:ent_offsets[i + 1]] - ent_rows.append(_build_padded_row(ent_suffix, prefix_len, orig_len, L_max)) + entropy_suffix = entropy_values[entropy_offsets[seq_idx]:entropy_offsets[seq_idx + 1]] + entropy_rows.append(_build_padded_row(entropy_suffix, prefix_len, orig_len, L_max)) - logp_2d = torch.stack(logp_rows, dim=0) - entropy_2d = torch.stack(ent_rows, dim=0) if ent_rows else None - return logp_2d, entropy_2d + log_probs_2d = torch.stack(log_probs_rows, dim=0) + entropy_2d = torch.stack(entropy_rows, dim=0) if entropy_rows else None + return log_probs_2d, entropy_2d def _build_padded_row( @@ -521,7 +521,7 @@ def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any: """完整 2D [B, L_max] → NestedTensor (jagged)，按各 original_lengths 切片。""" import torch - rows = [tensor_2d[i, :original_lengths[i]] for i in range(len(original_lengths))] + rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))] return torch.nested.nested_tensor(rows, layout=torch.jagged) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py index c4dcda99..65d7b746 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py @@ -8,9 +8,10 @@ 仍在、context 激活时，把 prefix-last 重算所需的 provider logits 保存到 ``ctx.prefix_last_logits_saved``，供后续 restore 使用。 -保存条件：``ctx.prefix_last_restore_indices`` 现仅含 prefix-last（每 reuser 一条）， -interior 由 restore 侧 build_kv 式 bulk 切片处理，不读 logits。逐条保存 provider -直接对应位置的 vocab 维 logits。 +保存条件：``ctx.prefix_last_restore_indices`` 每条对应一个 reuser 的 prefix-last +（每 reuser 一条）。逐条保存其 provider 直接对应位置的 vocab 维 logits，供 restore +侧重算 prefix-last logprob 使用（interior 区段由 restore 侧 build_kv 式 bulk 切片处理， +不读 logits）。 """ from __future__ import annotations @@ -69,8 +70,7 @@ def patched_fn(logits, labels): # ##### [PS-diag] 验证 packed 坐标对齐 end ##### for index in ctx.prefix_last_restore_indices: - # prefix_last_restore_indices 现在只含 prefix-last（interior 由 - # restore 侧 bulk 切片处理），逐条保存其 provider 的 vocab 维 logits。 + # 每条对应一个 reuser 的 prefix-last，逐条保存其 provider 的 vocab 维 logits。 pos = index.provider_1d_pos key = (index.reuse_idx_in_batch, index.target_2d_pos) if pos < 0: From f8bded16f23f8f91bd9a728a9741b814bc15e100 Mon Sep 17 00:00:00 2001 From: Boundless Date: Thu, 25 Jun 2026 20:55:29 +0800 Subject: [PATCH 04/61] =?UTF-8?q?[refactor]=20diag=20dump=E5=A4=B1?= =?UTF-8?q?=E8=B4=A5=E6=8F=90=E7=A4=BA=E6=94=B9=E7=94=A8print?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit last-attn/rope_freqs的ON dump失败提示从prefix_log.warning改为print，与同区块其它diag输出一致（直接走stdout）。 --- .../prefix_sharing/integrations/megatron_runtime.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 8756a03b..01ddb36d 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -126,7 +126,7 @@ def prefix_attention( attention_module.layer_number, attention_module.config.num_layers) except Exception as e: - prefix_log.warning(f"last-attn dump (ON) failed: {e}") + print(f"last-attn dump (ON) failed: {e}") ######### prefix-sharing diag: ON attention_output (per-layer) ######### # --- @@ -216,7 +216,7 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: dump_rope_freqs_on(q_freqs, attention_module.layer_number, attention_module.config.num_layers) except Exception as e: - prefix_log.warning(f"rope_freqs_on dump failed: {e}") + print(f"rope_freqs_on dump failed: {e}") ######### prefix-sharing diag: ON rope_freqs (per-layer) ######### query = apply_rotary_pos_emb( query.unsqueeze(1), From ab4943f25edb0a055232305ee1b6b657881ef121 Mon Sep 17 00:00:00 2001 From: Boundless Date: Thu, 25 Jun 2026 22:12:12 +0800 Subject: [PATCH 05/61] =?UTF-8?q?[test]=20NPU=20flash-attn:=20sparse=5Fmod?= =?UTF-8?q?e=3D3=20TND=20varlen=20=E5=8F=AF=E8=A1=8C=E6=80=A7=E5=AE=9E?= =?UTF-8?q?=E9=AA=8C?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 flash_atten_npu_test.py，验证 sparse_mode=3 (rightDownCausal) 在 TND varlen 下对 reuser(Q2048 - ground truth = TorchReferenceBackend；NPU-only，非NPU 干净 skip；__main__ 自带决策树 - 配套 plan: note 目录 npu-flash-attn-mode3-tnd-plan.md --- .../backends/flash_atten_npu_test.py | 468 ++++++++++++++++++ 1 file changed, 468 insertions(+) create mode 100644 prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py diff --git a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py new file mode 100644 index 00000000..1c09d708 --- /dev/null +++ b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py @@ -0,0 +1,468 @@ +"""TND (varlen) + sparse_mode=3 (rightDownCausal) 可行性实验（NPU only）。 + +背景 +---- +当前 ON 路径 ([flash_atten_npu.py](prefix_sharing/backends/flash_atten_npu.py)) 用 +BSH + per-sample B1SS mask + sparse_mode=1，在 reuser 上结果错；改 sparse_mode=3 又崩 +（"attenmask compression requires [2048,2048]"，因为 mode 3 要的是压缩 [2048,2048]， +不是 B1SS）。 + +调研 Ascend 60RC2 官方文档后发现：**sparse_mode=3 (rightDownCausal) 原生支持 Q Any: + """构造一个结构可控的 PrefixSharingPlan。 + + - batch_sizes[i] = 样本 i 的原始长度 + - prefix_lens[i] = 样本 i 的共享前缀长度（0 表示 provider / 无共享） + """ + config = PrefixSharingConfig(enable_prefix_sharing=True, backend="flash_atten_npu") + planner = PrefixSharingPlanner(config) + input_ids = [list(range(s)) for s in batch_sizes] + plan = planner.plan(input_ids) + + object.__setattr__(plan, "batch_size", len(batch_sizes)) + object.__setattr__(plan, "original_lengths", batch_sizes) + object.__setattr__(plan, "prefix_lens", prefix_lens) + object.__setattr__(plan, "kept_lengths_q", [b - p for b, p in zip(batch_sizes, prefix_lens)]) + object.__setattr__(plan, "expanded_lengths_kv", list(batch_sizes)) + object.__setattr__(plan, "q_position_offsets", prefix_lens) + object.__setattr__(plan, "kv_position_offsets", [0] * len(batch_sizes)) + + cu_seqlens_q = [0] + cu_seqlens_kv = [0] + max_seqlen_q = 0 + max_seqlen_kv = 0 + for b, p in zip(batch_sizes, prefix_lens): + q_len = b - p + kv_len = b + cu_seqlens_q.append(cu_seqlens_q[-1] + q_len) + cu_seqlens_kv.append(cu_seqlens_kv[-1] + kv_len) + max_seqlen_q = max(max_seqlen_q, q_len) + max_seqlen_kv = max(max_seqlen_kv, kv_len) + + object.__setattr__(plan, "cu_seqlens_q", cu_seqlens_q) + object.__setattr__(plan, "cu_seqlens_kv", cu_seqlens_kv) + object.__setattr__(plan, "max_seqlen_q", max_seqlen_q) + object.__setattr__(plan, "max_seqlen_kv", max_seqlen_kv) + object.__setattr__(plan, "provider_index", [0] * len(batch_sizes)) + object.__setattr__(plan, "is_provider", [p == 0 for p in prefix_lens]) + object.__setattr__(plan, "reuse_specs", ()) + object.__setattr__(plan, "prefix_last_restore", []) + return plan + + +def _make_layout(kept_lengths_q: list[int]) -> PackedBatchLayout: + return PackedBatchLayout.from_valid_lengths(kept_lengths_q) + + +def _random_qkv(total_q: int, total_kv: int, seed: int = 42): + """随机生成 THD Q / 原始(未展开) K/V。原始 K/V 长度 = kept_lengths_q 之和。""" + torch.manual_seed(seed) + q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 + k = torch.randn(total_kv, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 + v = torch.randn(total_kv, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 + return q, k, v + + +# ═══════════════════════════════════════════════════════════════════════ +# 本实验新增 helper +# ═══════════════════════════════════════════════════════════════════════ +_COMPRESSED_MASK: torch.Tensor | None = None + + +def _compressed_causal_mask() -> torch.Tensor: + """压缩 [2048,2048] 下三角 mask（True=masked），与 baseline get_attention_mask 一致。 + + sparse_mode 2/3/4 共用这一张压缩 mask；区别只在 kernel 锚定方式。 + 构建一次缓存。理论上覆盖任意 seq 长（压缩模式由 actual_seq 重建每段 causal）。 + """ + global _COMPRESSED_MASK + if _COMPRESSED_MASK is None or _COMPRESSED_MASK.device != DEVICE: + _COMPRESSED_MASK = torch.triu( + torch.ones([2048, 2048], dtype=torch.bool, device=DEVICE), diagonal=1 + ) + return _COMPRESSED_MASK + + +def _build_inputs(plan: Any, seed: int = 42): + """构造一次实验所需的全部张量 + ground truth。 + + 返回: + q: [总Q, NUM_HEADS, HEAD_DIM] —— reuser suffix-only + ek, ev: [总KV展开, NUM_KV_HEADS, HEAD_DIM] —— build_kv 展开后的 K/V（reuser 是 prefix+suffix） + ref_out: torch_ref 在相同 (q, ek, ev, plan) 上的输出，作为 ground truth + layout: PackedBatchLayout + """ + layout = _make_layout(plan.kept_lengths_q) + total_q = int(sum(plan.kept_lengths_q)) + # 原始(未展开) K/V 长度 = kept_lengths_q 之和（裁剪后 provider 全长、reuser suffix） + q, k_raw, v_raw = _random_qkv(total_q, total_q, seed=seed) + + store = PrefixAttentionStore() + ref = TorchReferenceBackend() + ek, ev = ref.build_kv( + k_raw, v_raw, store, plan, + packed_batch_layout=layout, layer_id=0, tp_rank=0, + ) + ref_out = ref.attention(q, ek, ev, plan, packed_batch_layout=layout) + store.close() + return q, ek, ev, ref_out, layout + + +def _varlen_tnd_call( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + *, + cu_q: list[int], + cu_kv: list[int], + sparse_mode: int, +) -> torch.Tensor: + """TND varlen npu_fusion_attention 调用。 + + - input_layout="TND" + - actual_seq_qlen/kvlen 用 cu_seqlens（**带前导 0**，对齐 verl util.py:69 的 zeros(batch+1)） + - atten_mask = 压缩 [2048,2048] 下三角 + - scale / keep_prob = 1/√d / 1.0 + - mode 2/3 下 pre/next_tokens 不生效，走默认 + """ + fn = _import_npu_fusion_attention() + result = fn( + q, k, v, + NUM_HEADS, + "TND", + atten_mask=_compressed_causal_mask(), + scale=SCALE, + keep_prob=1.0, + sparse_mode=sparse_mode, + actual_seq_qlen=list(cu_q), + actual_seq_kvlen=list(cu_kv), + ) + return result[0] if isinstance(result, (tuple, list)) else result + + +def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float: + return float((a.float() - b.float()).abs().max().item()) + + +# ═══════════════════════════════════════════════════════════════════════ +# Probe 矩阵。每个返回 dict: {name, pass, detail, out_diff, grad_diff} +# ═══════════════════════════════════════════════════════════════════════ +def probe_a_provider_only_mode2() -> dict: + """A. provider-only + mode 2（baseline sanity）。 + + 两个 provider，无共享，Q==KV。验证 TND varlen + 压缩 mask + mode 2 的接线正确 + （与 baseline 同款），并与 torch_ref 对齐。A 不过 = 环境/接线坏，后续 probe 都不可信。 + """ + plan = _make_plan(batch_sizes=[4, 6], prefix_lens=[0, 0]) + q, ek, ev, ref_out, _ = _build_inputs(plan) + try: + out = _varlen_tnd_call( + q, ek, ev, + cu_q=plan.cu_seqlens_q, cu_kv=plan.cu_seqlens_kv, + sparse_mode=2, + ) + diff = _max_abs_diff(out, ref_out) + ok = diff < _ATOL_FP16 + return {"name": "A", "pass": ok, "detail": f"provider-only mode2 out_diff={diff:.4e}", + "out_diff": diff, "grad_diff": None} + except Exception as e: # noqa: BLE001 + return {"name": "A", "pass": False, "detail": f"exception: {type(e).__name__}: {e}", + "out_diff": None, "grad_diff": None} + + +def probe_b_single_reuser_mode3() -> dict: + """B. 单 reuser + mode 3（核心 probe）。 + + 从 [8,8]/[0,4] 的 batch 里取 reuser(index=1)：Q=suffix(4)、KV=prefix+suffix(8)、 + prefix_len=4。单独喂给 TND varlen + mode 3，与 torch_ref 的 reuser 段输出对齐。 + + B 过 = mode 3 在 varlen 对 Q dict: + """C. 全 batch（provider+reuser）mode 3 单次调用。 + + 一次 npu_fusion_attention 覆盖 provider 段（Q==KV）和 reuser 段（Q dict: + """D. C 的反向（128-tile 约束）。 + + 全 batch mode 3 forward + sum().backward()，对齐 torch_ref 的 Q/K/V 梯度。 + D 过 = varlen 反向通，128-tile 约束不挡路（或被满足）。 + D 崩 tiling → 需要把 max_q/max_kv 补到 128 倍数重试（生产改写时处理）。 + """ + plan = _make_plan(batch_sizes=[8, 8], prefix_lens=[0, 4]) + + def _run(forward_fn): + q, ek, ev, _, _ = _build_inputs(plan) + q = q.clone().detach().requires_grad_(True) + ek = ek.clone().detach().requires_grad_(True) + ev = ev.clone().detach().requires_grad_(True) + out = forward_fn(q, ek, ev) + out.sum().backward() + return {"q": q.grad, "k": ek.grad, "v": ev.grad} + + # FA 反向 + def fa_fwd(qq, kk, vv): + return _varlen_tnd_call(qq, kk, vv, cu_q=plan.cu_seqlens_q, + cu_kv=plan.cu_seqlens_kv, sparse_mode=3) + # ref 反向（torch_ref） + layout = _make_layout(plan.kept_lengths_q) + ref_backend = TorchReferenceBackend() + + def ref_fwd(qq, kk, vv): + return ref_backend.attention(qq, kk, vv, plan, packed_batch_layout=layout) + + try: + grads_fa = _run(fa_fwd) + grads_ref = _run(ref_fwd) + max_grad_diff = 0.0 + for name in ("q", "k", "v"): + d = _max_abs_diff(grads_fa[name], grads_ref[name]) + max_grad_diff = max(max_grad_diff, d) + ok = max_grad_diff < _ATOL_GRAD_FP16 + return {"name": "D", "pass": ok, + "detail": f"full-batch mode3 backward grad_diff={max_grad_diff:.4e}", + "out_diff": None, "grad_diff": max_grad_diff} + except Exception as e: # noqa: BLE001 + return {"name": "D", "pass": False, + "detail": f"backward exception (可能 128-tile): {type(e).__name__}: {e}", + "out_diff": None, "grad_diff": None} + + +def probe_e_long_seq_gt_2048() -> dict: + """E.（可选）mode 3 + 某段 seq>2048。 + + 把 provider 段拉到 >2048，确认压缩 [2048,2048] mask 不限 seq 长（压缩模式由 + actual_seq 重建每段 causal，理论上支持任意 seq）。这是用户关心的点。 + """ + plan = _make_plan(batch_sizes=[2100, 2100], prefix_lens=[0, 100]) + q, ek, ev, ref_out, _ = _build_inputs(plan) + try: + out = _varlen_tnd_call( + q, ek, ev, + cu_q=plan.cu_seqlens_q, cu_kv=plan.cu_seqlens_kv, + sparse_mode=3, + ) + diff = _max_abs_diff(out, ref_out) + ok = diff < _ATOL_FP16 + return {"name": "E", "pass": ok, + "detail": f"seq>2048 mode3 out_diff={diff:.4e}", + "out_diff": diff, "grad_diff": None} + except Exception as e: # noqa: BLE001 + return {"name": "E", "pass": False, "detail": f"exception: {type(e).__name__}: {e}", + "out_diff": None, "grad_diff": None} + + +_PROBES = [ + probe_a_provider_only_mode2, + probe_b_single_reuser_mode3, + probe_c_full_batch_mode3, + probe_d_full_batch_backward, + probe_e_long_seq_gt_2048, +] + + +# ═══════════════════════════════════════════════════════════════════════ +# pytest 包装（每个 probe 一个 test，非 NPU skip） +# ═══════════════════════════════════════════════════════════════════════ +@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") +def test_probe_a_provider_only_mode2(): + r = probe_a_provider_only_mode2() + assert r["pass"], r["detail"] + + +@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") +def test_probe_b_single_reuser_mode3(): + r = probe_b_single_reuser_mode3() + assert r["pass"], r["detail"] + + +@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") +def test_probe_c_full_batch_mode3(): + r = probe_c_full_batch_mode3() + assert r["pass"], r["detail"] + + +@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") +def test_probe_d_full_batch_backward(): + r = probe_d_full_batch_backward() + assert r["pass"], r["detail"] + + +@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") +def test_probe_e_long_seq_gt_2048(): + r = probe_e_long_seq_gt_2048() + assert r["pass"], r["detail"] + + +# ═══════════════════════════════════════════════════════════════════════ +# 决策树（__main__ 用） +# ═══════════════════════════════════════════════════════════════════════ +def _decide(results: dict[str, dict]) -> str: + a, b, c, d = results["A"], results["B"], results["C"], results["D"] + if not a["pass"]: + return ("A 失败 → 环境/接线坏（input_layout 字符串/cu_seqlens 格式/scale 等）。" + "先修 A，后续 probe 都不可信。") + if not b["pass"]: + return ("B 失败 → mode 3 在本 CANN 版本的 varlen 下没对 Q None: + if not _HAS_NPU: + print("[skip] 无 NPU 设备或 mindspeed 内核，本实验只能在 NPU 上跑。") + return + print(f"[env] DEVICE={DEVICE} NUM_HEADS={NUM_HEADS} NUM_KV_HEADS={NUM_KV_HEADS} " + f"HEAD_DIM={HEAD_DIM} SCALE={SCALE:.4f}") + print("=" * 70) + results: dict[str, dict] = {} + for probe in _PROBES: + r = probe() + results[r["name"]] = r + tag = "PASS" if r["pass"] else "FAIL" + diff_str = [] + if r["out_diff"] is not None: + diff_str.append(f"out_diff={r['out_diff']:.4e}") + if r["grad_diff"] is not None: + diff_str.append(f"grad_diff={r['grad_diff']:.4e}") + diff_txt = f" [{', '.join(diff_str)}]" if diff_str else "" + print(f"Probe {r['name']}: {tag}{diff_txt}") + print(f" {r['detail']}") + print("=" * 70) + print("[决策] " + _decide(results)) + + +if __name__ == "__main__": + _run_all() From 87941800f03ba69aac744a725c17db9d82991093 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 09:30:47 +0800 Subject: [PATCH 06/61] =?UTF-8?q?[feat]=20=E6=96=B0=E5=A2=9E=20NPU=20flash?= =?UTF-8?q?-attn=20TND=20varlen=20+=20sparse=5Fmode=3D3=20=E5=90=8E?= =?UTF-8?q?=E7=AB=AF?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新后端 flash_atten_npu_tnd (NpuFlashAttentionTndBackend)：TND varlen + sparse_mode=3 (rightDownCausal)，一次调用覆盖全 batch，reuser Q 仍 suffix-only。 - mode 3 原生支持 reuser(Q2048） - 老 BSH 后端 flash_atten_npu 保留作回退；config 里 backend='flash_atten_npu_tnd' 切换 - 注册到 factory/config/__init__，加 test_factory_flash_atten_npu_tnd --- .../prefix_sharing/backends/__init__.py | 2 + .../prefix_sharing/backends/factory.py | 15 +- .../backends/flash_atten_npu_tnd.py | 255 ++++++++++++++++++ prefix-sharing/prefix_sharing/core/config.py | 6 +- .../tests/unit_test/test_backend_factory.py | 11 +- 5 files changed, 280 insertions(+), 9 deletions(-) create mode 100644 prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py diff --git a/prefix-sharing/prefix_sharing/backends/__init__.py b/prefix-sharing/prefix_sharing/backends/__init__.py index 0996d446..63cc533c 100644 --- a/prefix-sharing/prefix_sharing/backends/__init__.py +++ b/prefix-sharing/prefix_sharing/backends/__init__.py @@ -6,6 +6,7 @@ from prefix_sharing.backends.flash_atten_base import FlashAttentionMixin, FlashBackendValidationError from prefix_sharing.backends.flash_atten_gpu import GpuFlashAttentionBackend from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend +from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend from prefix_sharing.backends.torch_ref import TorchReferenceBackend __all__ = [ @@ -14,6 +15,7 @@ "FlashBackendValidationError", "GpuFlashAttentionBackend", "NpuFlashAttentionBackend", + "NpuFlashAttentionTndBackend", "PrefixAttentionBackend", "PrefixDeltanetBackend", "TorchReferenceBackend", diff --git a/prefix-sharing/prefix_sharing/backends/factory.py b/prefix-sharing/prefix_sharing/backends/factory.py index a98ed65b..2359a73b 100644 --- a/prefix-sharing/prefix_sharing/backends/factory.py +++ b/prefix-sharing/prefix_sharing/backends/factory.py @@ -17,23 +17,28 @@ def get_backend_instance( * ``"torch_ref"`` -> :class:`~prefix_sharing.backends.torch_ref.TorchReferenceBackend` * ``"flash_atten_gpu"`` -> :class:`~prefix_sharing.backends.flash_atten_gpu.GpuFlashAttentionBackend` * ``"flash_atten_npu"`` -> :class:`~prefix_sharing.backends.flash_atten_npu.NpuFlashAttentionBackend` + * ``"flash_atten_npu_tnd"`` -> :class:`~prefix_sharing.backends.flash_atten_npu_tnd.NpuFlashAttentionTndBackend` (recommended for NPU) """ if backend is not None: return backend - + if config.backend == "torch_ref": from prefix_sharing.backends.torch_ref import TorchReferenceBackend return TorchReferenceBackend() - + if config.backend == "flash_atten_gpu": from prefix_sharing.backends.flash_atten_gpu import GpuFlashAttentionBackend return GpuFlashAttentionBackend() - + if config.backend == "flash_atten_npu": from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend return NpuFlashAttentionBackend() - + + if config.backend == "flash_atten_npu_tnd": + from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend + return NpuFlashAttentionTndBackend() + raise ValueError( f"Unknown backend '{config.backend}'. " - f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu" + f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu, flash_atten_npu_tnd" ) diff --git a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py new file mode 100644 index 00000000..033735fb --- /dev/null +++ b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py @@ -0,0 +1,255 @@ +"""Ascend NPU Flash Attention backend — TND varlen + sparse_mode=3 (rightDownCausal). + +这是**推荐的 NPU 后端**，与 OFF baseline（MindSpeed ``dot_product_attention``）使用 +相同的 TND varlen 约定，配合 ``sparse_mode=3``。老的 BSH 后端 +(:mod:`prefix_sharing.backends.flash_atten_npu`) 保留作对照/回退，不要删。 + +为什么用 sparse_mode=3（核心） +------------------------------- +``sparse_mode=3`` 即 **rightDownCausal**："以右下顶点划分的下三角"。对一个 Q 比 +KV 短的 segment（Sq < Skv），它把 Q **右对齐到 KV 末端**：query i（局部）见 kv j +iff ``j <= (Skv - Sq) + i``。 + +代到 prefix-sharing 的 reuser（Q=suffix、KV=prefix+suffix、Skv-Sq=prefix_len）： +query i 见 kv j iff ``j <= prefix_len + i`` = **全部 prefix KV 可见 + suffix KV causal**， +正是 reuser 的正确语义。而对 provider（Sq==Skv）退化成 ``j <= i`` = 标准 causal， +与 baseline（``sparse_mode=2`` leftUpCausal）一致。 + +因此**整个 batch（providers + reusers）一次 varlen 调用即可**：kernel 按每个 segment +的 actual_seq 自动判 provider（标准 causal）还是 reuser（右对齐 causal）。reuser 的 +Q 仍是 suffix-only（省算力核心收益不变），mask 用现成的压缩 ``[2048,2048]`` 下三角 +（与 baseline ``get_attention_mask`` 完全一样），**无需自建 mask、无需按 prefix_len +分组拆调用**。 + +为什么不用老 backend 的 BSH + sparse_mode=1 +------------------------------------------- +老 backend 用 BSH + per-sample B1SS mask + ``sparse_mode=1``，实测在 reuser 上结果 +错（mode 1=allMask 对 B1SS 自建 mask 的处理不对）。改 ``sparse_mode=3`` 又崩 +（"attenmask compression requires [2048,2048]"，因为 mode 3 要压缩 ``[2048,2048]``， +不是 B1SS）。TND varlen + 压缩 mask 才是 mode 3 的正确用法。 + +曾经担心 varlen 反向有 128-tile 约束（见老 backend 的 docstring），实测 +（``flash_atten_npu_test.py`` Probe D）已证伪——本 CANN 版本 varlen 反向正常。 + +选用方式 +-------- +配置里设 ``backend="flash_atten_npu_tnd"`` 即可指向本后端。 + +Select via config: ``backend="flash_atten_npu_tnd"``. +""" + +from __future__ import annotations + +import importlib +import math +from functools import lru_cache +from typing import Any + +from prefix_sharing.backends.base import BackendCapabilities +from prefix_sharing.backends.flash_atten_base import ( + FlashAttentionMixin, + FlashBackendValidationError, +) +from prefix_sharing.backends.torch_ref import TorchReferenceBackend +from prefix_sharing.core.config import PrefixSharingConfig +from prefix_sharing.core.planner import PrefixSharingPlan + + +_CANDIDATES = [ + ("mindspeed.ops.fusion_attention_v2", "npu_fusion_attention"), + ("mindspeed.ops", "npu_fusion_attention"), +] + + +@lru_cache(maxsize=None) +def _import_npu_fusion_attention(): + last_err = None + for module_name, attr in _CANDIDATES: + try: + module = importlib.import_module(module_name) + return getattr(module, attr) + except ImportError as e: + last_err = e + raise RuntimeError( + "NpuFlashAttentionTndBackend requires MindSpeed (mindspeed.ops). " + "Install MindSpeed matching your CANN version." + ) from last_err + + +def _torch() -> Any: + try: + import torch + except ModuleNotFoundError as exc: + raise RuntimeError("NpuFlashAttentionTndBackend requires PyTorch") from exc + return torch + + +# 压缩 [2048,2048] 下三角 mask（True=masked）按 device 缓存。 +# sparse_mode 2/3/4 共用这张压缩 mask；kernel 拿 actual_seq 重建每段 causal， +# 故不限 seq 长（实测 seq>2048 正常，见 flash_atten_npu_test.py Probe E）。 +_COMPRESSED_MASK: dict[Any, Any] = {} + + +def _compressed_causal_mask(device: Any) -> Any: + torch = _torch() + cached = _COMPRESSED_MASK.get(device) + if cached is None: + cached = torch.triu( + torch.ones([2048, 2048], dtype=torch.bool, device=device), diagonal=1 + ) + _COMPRESSED_MASK[device] = cached + return cached + + +class NpuFlashAttentionTndBackend(FlashAttentionMixin): + """Ascend NPU backend via ``npu_fusion_attention`` (TND varlen, sparse_mode=3). + + 单次 varlen 调用覆盖整个 batch（providers + reusers）： + + - **Provider 段**（Sq==Skv）：``sparse_mode=3`` 退化成标准 causal。 + - **Reuser 段**（Sq None: + self._torch_ref = TorchReferenceBackend() + + def validate(self, config: PrefixSharingConfig, model_config: Any | None = None) -> None: + config.validate(model_config=model_config) + _import_npu_fusion_attention() + + def apply_rope( + self, + query: Any, + key: Any, + prefix_sharing_plan: PrefixSharingPlan, + **kwargs: Any, + ) -> tuple[Any, Any]: + return self._torch_ref.apply_rope(query, key, prefix_sharing_plan, **kwargs) + + def build_kv( + self, + key: Any, + value: Any, + store: Any, + prefix_sharing_plan: PrefixSharingPlan, + *, + packed_batch_layout: Any | None = None, + layer_id: int, + tp_rank: int = 0, + stats: Any | None = None, + ) -> tuple[Any, Any]: + """KV 展开委托给 torch reference（与其它后端一致）。""" + return self._torch_ref.build_kv( + key, + value, + store, + prefix_sharing_plan, + packed_batch_layout=packed_batch_layout, + layer_id=layer_id, + tp_rank=tp_rank, + stats=stats, + ) + + # ------------------------------------------------------------------ + # attention — TND varlen + sparse_mode=3，单次调用全 batch + # ------------------------------------------------------------------ + def attention( + self, + query: Any, + key: Any, + value: Any, + prefix_sharing_plan: PrefixSharingPlan, + **kwargs: Any, + ) -> Any: + """Run prefix-sharing attention via TND-varlen ``npu_fusion_attention``. + + 流程： + 1. ``_prepare_flash_inputs`` 剥 Q 的 TP padding、产出 cu_seqlens_q/kv + （batch+1，带前导 0，取自 plan）。 + 2. 一次 ``npu_fusion_attention``，``input_layout="TND"``、 + ``sparse_mode=3``（rightDownCausal）、压缩 ``[2048,2048]`` mask。 + 3. 必要时把输出按 pad_layout 回填 TP padding，恢复原 Q 形状。 + + K/V 来自 ``build_kv``，已是展开后的 TND（reuser 是 prefix+suffix）， + 无 padding，跟随 ``plan.expanded_lengths_kv``。 + """ + layer_id = kwargs.get("layer_id", "?") + packed_batch_layout = kwargs.get("packed_batch_layout") + if packed_batch_layout is None: + raise FlashBackendValidationError( + "flash_atten_npu_tnd.attention requires packed_batch_layout kwarg." + ) + + print( + f"[PS][backend] flash_atten_npu_tnd attention: " + f"layer={layer_id}, " + f"q_shape={tuple(query.shape)}, k_shape={tuple(key.shape)}, " + f"v_shape={tuple(value.shape)}" + ) + + npu_fusion_attention = _import_npu_fusion_attention() + + # Step 1: 剥 Q 的 TP padding + 取 cu_seqlens_q/kv（plan 语义长度，带前导 0）。 + q, k, v, cu_seqlens_q, cu_seqlens_kv, _, _, _, pad_layout = ( + self._prepare_flash_inputs( + query, + key, + value, + prefix_sharing_plan, + attention_mask=kwargs.get("attention_mask"), + packed_batch_layout=packed_batch_layout, + ) + ) + + num_q_heads = q.shape[1] + head_dim = q.shape[-1] + scale = kwargs.get("softmax_scale") or (1.0 / math.sqrt(head_dim)) + dropout_p = kwargs.get("dropout_p", 0.0) + keep_prob = kwargs.get("keep_prob", 1.0 - dropout_p) + + # Step 2: 单次 varlen 调用，sparse_mode=3（rightDownCausal）。 + # mode 3 下 pre/next_tokens 不生效，走默认；atten_mask 用压缩 [2048,2048]。 + try: + result = npu_fusion_attention( + q, + k, + v, + num_q_heads, + "TND", + atten_mask=_compressed_causal_mask(q.device), + scale=scale, + keep_prob=keep_prob, + sparse_mode=3, + actual_seq_qlen=cu_seqlens_q.tolist(), + actual_seq_kvlen=cu_seqlens_kv.tolist(), + ) + except Exception as exc: + raise FlashBackendValidationError( + f"npu_fusion_attention (TND, sparse_mode=3) failed: " + f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}, " + f"cu_seqlens_q={cu_seqlens_q.tolist()}, " + f"cu_seqlens_kv={cu_seqlens_kv.tolist()}" + ) from exc + + output = result[0] if isinstance(result, (tuple, list)) else result + + # Step 3: 回填 TP padding，恢复原 Q 形状（TP=1 时 pad_layout 为 None，no-op）。 + if pad_layout is not None: + output = self._repad_output(output, pad_layout) + + return output diff --git a/prefix-sharing/prefix_sharing/core/config.py b/prefix-sharing/prefix_sharing/core/config.py index 6d85047c..737f4034 100644 --- a/prefix-sharing/prefix_sharing/core/config.py +++ b/prefix-sharing/prefix_sharing/core/config.py @@ -114,7 +114,7 @@ def validate(self, model_config: Any | None = None, integrate_mode: str | None = return if self.detector != "trie": raise PrefixSharingConfigError("phase 1 supports only detector='trie'") - supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu"} + supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"} if self.backend not in supported_backends: raise PrefixSharingConfigError( f"backend='{self.backend}' is not supported. " @@ -217,10 +217,10 @@ def validate_for_engine( # 基础校验 if self.detector != "trie": raise PrefixSharingConfigError("phase 1 supports only detector='trie'") - if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu"}: + if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"}: raise PrefixSharingConfigError( f"backend='{self.backend}' is not supported. " - f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu" + f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu, flash_atten_npu_tnd" ) if self.boundary_strategy != "prefix_last_restore": raise PrefixSharingConfigError( diff --git a/prefix-sharing/tests/unit_test/test_backend_factory.py b/prefix-sharing/tests/unit_test/test_backend_factory.py index 1ef5128d..455d6317 100644 --- a/prefix-sharing/tests/unit_test/test_backend_factory.py +++ b/prefix-sharing/tests/unit_test/test_backend_factory.py @@ -7,6 +7,7 @@ from prefix_sharing.backends.factory import get_backend_instance from prefix_sharing.backends.flash_atten_gpu import GpuFlashAttentionBackend from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend +from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend from prefix_sharing.backends.torch_ref import TorchReferenceBackend from prefix_sharing.core.config import PrefixSharingConfig @@ -33,6 +34,14 @@ def test_factory_flash_atten_npu(): assert backend.capabilities.name == "flash_atten_npu" +def test_factory_flash_atten_npu_tnd(): + config = PrefixSharingConfig(enable_prefix_sharing=True, backend="flash_atten_npu_tnd") + backend = get_backend_instance(config) + assert isinstance(backend, NpuFlashAttentionTndBackend) + assert backend.capabilities.name == "flash_atten_npu_tnd" + assert backend.capabilities.supports_cann + + def test_factory_unknown_backend() -> None: config = PrefixSharingConfig(enable_prefix_sharing=True, backend="unknown") with pytest.raises(ValueError, match="Unknown backend"): @@ -62,7 +71,7 @@ def test_config_validates_backends() -> None: def test_config_accepts_supported_backends(): - for name in ("torch_ref", "flash_atten_gpu", "flash_atten_npu"): + for name in ("torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"): cfg = PrefixSharingConfig(enable_prefix_sharing=True, backend=name) cfg.validate() # should not raise \ No newline at end of file From f6a5dd516998a1a4c145cfda40236bfbd28728a2 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 09:55:37 +0800 Subject: [PATCH 07/61] =?UTF-8?q?[fix]=20flash=5Fatten=5Fnpu=5Ftnd:=20=5Fp?= =?UTF-8?q?repare=5Fflash=5Finputs=20=E8=A7=A3=E5=8C=85=E6=95=B0=E9=87=8F?= =?UTF-8?q?=E4=BF=AE=E6=AD=A3=EF=BC=888=20=E5=85=83=E7=BB=84=E4=B8=8D?= =?UTF-8?q?=E6=98=AF=209=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit attention() 里多写了一个 _，导致 ValueError: not enough values to unpack。 _prepare_flash_inputs 返回 8 元组 (q,k,v,cu_seqlens_q,cu_seqlens_kv,max_seqlen_q,max_seqlen_kv,pad_layout)。 实验文件走自己的 _varlen_tnd_call 没踩到，真实 backend 路径才暴露。 --- prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py index 033735fb..29abb8fb 100644 --- a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py +++ b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py @@ -205,7 +205,9 @@ def attention( npu_fusion_attention = _import_npu_fusion_attention() # Step 1: 剥 Q 的 TP padding + 取 cu_seqlens_q/kv（plan 语义长度，带前导 0）。 - q, k, v, cu_seqlens_q, cu_seqlens_kv, _, _, _, pad_layout = ( + # _prepare_flash_inputs 返回 8 元组：q, k, v, cu_seqlens_q, cu_seqlens_kv, + # max_seqlen_q, max_seqlen_kv, pad_layout（max_seqlen_* 这里不用）。 + q, k, v, cu_seqlens_q, cu_seqlens_kv, _, _, pad_layout = ( self._prepare_flash_inputs( query, key, From d9279f6deebf2156a95794674981622dba93de4b Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 11:21:26 +0800 Subject: [PATCH 08/61] =?UTF-8?q?[refactor]=20cmp=5Fdiag=5Fverl080:=20firs?= =?UTF-8?q?t=5Ftoken=20=E6=94=B9=E4=B8=BA=E5=8F=AF=E6=8C=87=E5=AE=9A=20pac?= =?UTF-8?q?ked=20=E4=BD=8D=E7=BD=AE=E7=9A=84=20packed=5Ftoken?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 新增 --token POS（单个 int，默认 0），比较 packed[POS] 而非写死 packed[0] - cmp_first_token → cmp_packed_token(dir_on, dir_off, pos, layer)；_logits_first_token → _logits_at_pos；_first_token_metrics → _vec_metrics；_print_first_token → _print_packed_token - 输出名/section 头不再有 first_token 字样：[packed_token] attn_L{layer}_pos{pos} / logits_pos{pos} - --layer 同时控制 packed_token 的 attn 层（默认最后一层）；logits 永远最后一层 - logits top-K 表按 val 排序、去掉 REL_ERR 列；attn 表保持原样 --- .../prefix_sharing/tools/cmp_diag_verl080.py | 116 +++++++++++------- 1 file changed, 69 insertions(+), 47 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 259ffb86..44b61da4 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -7,7 +7,7 @@ **packed（suffix 对齐）** - attention_output per-layer cos 每层 attention 输出余弦相似度 - - first_token packed[0]（attn[0] + logits[0]） + - packed_token packed[pos]（attn[pos] + logits[pos]，--token 指定 pos，默认 0） - logits packed 全 packed logits suffix 对齐对比 **2D（v080 特有，restore 后 ``[B, L_max]``）** @@ -32,14 +32,14 @@ cu_seqlens_q_logits.pt [B+1] logits packed 边界（同上） Usage: - # 完整对比（attn per-layer + first_token + logits + logprobs + entropy） + # 完整对比（attn per-layer + packed_token + logits + logprobs + entropy） python cmp_diag_verl080.py --dir-on ./dump_on --dir-off ./dump_off --tag old # 只看某一层 attention（1-indexed） python cmp_diag_verl080.py --dir-on ./dump_on --dir-off ./dump_off \\ --tag old --layer 12 - # top-K 误差最大位置（2D + first_token） + # top-K 误差最大位置（2D + packed_token） python cmp_diag_verl080.py --dir-on ./dump_on --dir-off ./dump_off \\ --tag old --topk 20 @@ -132,7 +132,7 @@ def _pearson_r(t1: torch.Tensor, t2: torch.Tensor, return float("nan") if sx == 0 or sy == 0 else float(cov / (sx * sy)) -def _first_token_metrics(a_vec: torch.Tensor, b_vec: torch.Tensor) -> dict: +def _vec_metrics(a_vec: torch.Tensor, b_vec: torch.Tensor) -> dict: err = _error_abs_rel(a_vec, b_vec) cos = float(_cosine_sim(a_vec, b_vec, dim=-1)) pr = _pearson_r(a_vec, b_vec) @@ -268,12 +268,12 @@ def _align_packed(on_tensor: torch.Tensor, off_tensor: torch.Tensor, # Logits helpers # ════════════════════════════════════════════════════════════════ -def _logits_first_token(lo: torch.Tensor, lf: torch.Tensor - ) -> tuple[torch.Tensor, torch.Tensor]: - """取 packed[0] 的 full-vocab 向量（logits 词表恒在最后一维）。""" +def _logits_at_pos(lo: torch.Tensor, lf: torch.Tensor, + pos: int = 0) -> tuple[torch.Tensor, torch.Tensor]: + """取 packed[pos] 的 full-vocab 向量（logits 词表恒在最后一维）。""" lo_2d = lo.reshape(-1, lo.size(-1)) lf_2d = lf.reshape(-1, lf.size(-1)) - return lo_2d[0, :].contiguous(), lf_2d[0, :].contiguous() + return lo_2d[pos, :].contiguous(), lf_2d[pos, :].contiguous() def _logits_ensure_token_major(lo: torch.Tensor, lf: torch.Tensor @@ -284,7 +284,7 @@ def _logits_ensure_token_major(lo: torch.Tensor, lf: torch.Tensor # ════════════════════════════════════════════════════════════════ -# Packed compare: attention_output / first_token / logits +# Packed compare: attention_output / packed_token / logits # ════════════════════════════════════════════════════════════════ def _cos_for_layer(a: torch.Tensor, b: torch.Tensor, @@ -357,32 +357,36 @@ def cmp_attn_layer(dir_on: str, dir_off: str, metrics={"layers": results}) -def cmp_first_token(dir_on: str, dir_off: str) -> list[CheckResult]: - """packed[0] 对比：最后一层 attn[0] + logits[0]。 +def cmp_packed_token(dir_on: str, dir_off: str, + pos: int = 0, layer: int | None = None) -> list[CheckResult]: + """packed[pos] 对比：attn[pos]（可指定层）+ logits[pos]（仅最后一层）。 - 第一个序列（row 0）永远是 provider（完整序列），packed[0] 是完整 suffix token， - 可直接对比无需对齐。 + - pos：packed 里的 token 位置（单个 int，默认 0）。不是 2D 位置。 + - attn：用 *layer*（默认最后一层）。对比第 1 层可区分 + "结构错（位置/RoPE/mask，第 1 层就偏）" vs "数值累积（第 1 层完美、深层才偏）"。 + - logits：永远最后一层。logits 只在最后产生，不受 *layer* 影响。 """ results: list[CheckResult] = [] - last = _get_num_layers(dir_on) or _get_num_layers(dir_off) + attn_layer = layer if layer is not None else ( + _get_num_layers(dir_on) or _get_num_layers(dir_off)) - if last: - a = _load_attn_output(dir_on, last) - b = _load_attn_output(dir_off, last) + if attn_layer: + a = _load_attn_output(dir_on, attn_layer) + b = _load_attn_output(dir_off, attn_layer) if a is not None and b is not None: a0 = a.squeeze(1) if a.dim() == 3 else a b0 = b.squeeze(1) if b.dim() == 3 else b results.append(CheckResult( - name="first_token_attn", - metrics=_first_token_metrics(a0[0], b0[0]))) + name=f"attn_L{attn_layer}_pos{pos}", + metrics=_vec_metrics(a0[pos], b0[pos]))) lo = _load_logits(dir_on) lf = _load_logits(dir_off) if lo is not None and lf is not None: - lo_first, lf_first = _logits_first_token(lo, lf) + lo_p, lf_p = _logits_at_pos(lo, lf, pos) results.append(CheckResult( - name="first_token_logits", - metrics=_first_token_metrics(lo_first, lf_first))) + name=f"logits_pos{pos}", + metrics=_vec_metrics(lo_p, lf_p))) return results @@ -736,8 +740,8 @@ def _print_per_layer(r: CheckResult): print() -def _print_first_token(r: CheckResult): - print(_SEP_SINGLE + f"\n [first_token] {r.name} (packed position [0])") +def _print_packed_token(r: CheckResult): + print(_SEP_SINGLE + f"\n [packed_token] {r.name}") print(_SEP_SINGLE) m = r.metrics for k in ["mean_abs", "max_abs", "rel_max", "rel_mean", "cos", "pearson"]: @@ -778,8 +782,11 @@ def _print_2d_result(r: CheckResult): def _print_topk_vec(on_vec: torch.Tensor, off_vec: torch.Tensor, - topk: int, sort_by: str, label: str): - """1D 向量 top-K（first_token per-dim）。""" + topk: int, sort_by: str, label: str, show_rel: bool = True): + """1D 向量 top-K（packed_token per-dim）。 + + show_rel=False 时省略 REL_ERR 列（用于 logits 表，只看 val/abs）。 + """ abs_err = (on_vec - off_vec).abs() rel_err = abs_err / torch.maximum(on_vec.abs(), off_vec.abs()).clamp(min=1e-8) if sort_by == "abs": @@ -791,10 +798,16 @@ def _print_topk_vec(on_vec: torch.Tensor, off_vec: torch.Tensor, _, idx = sort_key.topk(min(topk, sort_key.numel())) idx = idx.to(torch.long) print(f"\n [{label}] top-{topk} dims (sort by {sort_by})") - print(f" {'DIM':>6s} {'ON':>14s} {'OFF':>14s} {'ABS_ERR':>12s} {'REL_ERR':>12s}") - for i in idx.tolist(): - print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" - f" {float(abs_err[i]):>12.6e} {float(rel_err[i]):>12.6e}") + if show_rel: + print(f" {'DIM':>6s} {'ON':>14s} {'OFF':>14s} {'ABS_ERR':>12s} {'REL_ERR':>12s}") + for i in idx.tolist(): + print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" + f" {float(abs_err[i]):>12.6e} {float(rel_err[i]):>12.6e}") + else: + print(f" {'DIM':>6s} {'ON':>14s} {'OFF':>14s} {'ABS_ERR':>12s}") + for i in idx.tolist(): + print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" + f" {float(abs_err[i]):>12.6e}") def _print_topk_2d(on_t: torch.Tensor, off_t: torch.Tensor, @@ -875,7 +888,11 @@ def main(): ap.add_argument("--mask", choices=["label", "attention", "none"], default="label", help="2D mask type (default: label)") ap.add_argument("--layer", type=int, default=None, - help="Compare specific attn layer 1-indexed (default: all)") + help="Compare specific attn layer 1-indexed (default: all). " + "Also used by packed_token attn (default: last layer).") + ap.add_argument("--token", type=int, default=0, + help="Packed token position for packed_token compare " + "(single int index, default: 0)") ap.add_argument("--atol", type=float, default=1e-5, help="Absolute tolerance for 2D (default: 1e-5)") ap.add_argument("--topk", type=int, default=0, @@ -919,27 +936,32 @@ def main(): all_results.append(r) _print_per_layer(r) - # ── packed: first_token（attn[0] + logits[0]） ── - ft_results = cmp_first_token(args.dir_on, args.dir_off) - for r in ft_results: + # ── packed: packed_token（attn[pos] + logits[pos]） ── + # pos 由 --token 指定（默认 0）；attn 用 --layer 指定的层（默认最后一层）； + # logits 永远最后一层（logits 只在最后产生）。 + pos = args.token + pt_results = cmp_packed_token(args.dir_on, args.dir_off, pos, args.layer) + for r in pt_results: all_results.append(r) - _print_first_token(r) - # first_token top-K（per-dim） - if args.topk > 0 and ft_results: - last = _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off) - if last: - a = _load_attn_output(args.dir_on, last) - b = _load_attn_output(args.dir_off, last) + _print_packed_token(r) + # packed_token top-K（per-dim） + if args.topk > 0 and pt_results: + attn_layer = args.layer if args.layer is not None else ( + _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) + if attn_layer: + a = _load_attn_output(args.dir_on, attn_layer) + b = _load_attn_output(args.dir_off, attn_layer) if a is not None and b is not None: - a0 = (a.squeeze(1) if a.dim() == 3 else a)[0].cpu() - b0 = (b.squeeze(1) if b.dim() == 3 else b)[0].cpu() - _print_topk_vec(a0, b0, args.topk, "val", "first_token_attn") + ap = (a.squeeze(1) if a.dim() == 3 else a)[pos].cpu() + bp = (b.squeeze(1) if b.dim() == 3 else b)[pos].cpu() + _print_topk_vec(ap, bp, args.topk, "val", + f"attn_L{attn_layer}_pos{pos}") lo = _load_logits(args.dir_on) lf = _load_logits(args.dir_off) if lo is not None and lf is not None: - lo_f, lf_f = _logits_first_token(lo, lf) - _print_topk_vec(lo_f.cpu(), lf_f.cpu(), args.topk, - args.sort_err, "first_token_logits") + lo_p, lf_p = _logits_at_pos(lo, lf, pos) + _print_topk_vec(lo_p.cpu(), lf_p.cpu(), args.topk, + "val", f"logits_pos{pos}", show_rel=False) # ── packed: logits（suffix 对齐） ── r = cmp_logits_packed(args.dir_on, args.dir_off) From d7aab2643cef5d030ee45c7b076f70e666451e5a Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 11:34:42 +0800 Subject: [PATCH 09/61] =?UTF-8?q?[fix]=20cmp=5Fdiag=5Fverl080:=20packed=5F?= =?UTF-8?q?token=20=E6=AF=94=E8=BE=83=E5=89=8D=E5=85=88=E5=81=9A=20suffix?= =?UTF-8?q?=20=E5=AF=B9=E9=BD=90?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 之前 packed[pos] 直接比 ON[pos] vs OFF[pos]，但 ON 是裁剪后的 suffix-only packed、OFF 是完整 packed，两者 token 不对应——比较的是不同 token，结论无意义。 修复：新增 _aligned_vec_at_pos，先用 align_mask（OFF cu_seqlens + ON prefix_lens）把 OFF 的 suffix 段抽出来与 ON 对齐，再取 [pos]。pos 索引对齐后的 suffix-packed 空间（ON/OFF 一致，指向同一 token）。attn 和 logits 都走这条对齐路径，top-K 表同样先对齐。 --- .../prefix_sharing/tools/cmp_diag_verl080.py | 108 +++++++++++++----- 1 file changed, 78 insertions(+), 30 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 44b61da4..a4b97fc0 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -268,12 +268,41 @@ def _align_packed(on_tensor: torch.Tensor, off_tensor: torch.Tensor, # Logits helpers # ════════════════════════════════════════════════════════════════ -def _logits_at_pos(lo: torch.Tensor, lf: torch.Tensor, - pos: int = 0) -> tuple[torch.Tensor, torch.Tensor]: - """取 packed[pos] 的 full-vocab 向量（logits 词表恒在最后一维）。""" - lo_2d = lo.reshape(-1, lo.size(-1)) - lf_2d = lf.reshape(-1, lf.size(-1)) - return lo_2d[pos, :].contiguous(), lf_2d[pos, :].contiguous() +def _aligned_vec_at_pos( + on_tensor: torch.Tensor | None, + off_tensor: torch.Tensor | None, + is_attn: bool, + pos: int, + align_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor] | None: + """ON(suffix-only)/OFF(full-packed) 的 packed 张量 **suffix 对齐后** 取 [pos]。 + + ON 物理裁剪后只含 suffix，OFF 含完整序列，两者 token 不直接对应——必须先用 + align_mask 把 OFF 的 suffix 段抽出来与 ON 对齐，再取 [pos]。pos 索引的是 + 对齐后的 suffix-packed 空间（ON/OFF 一致，指向同一个 token）。 + + - is_attn=True：attn_output ``[T,1,hidden]`` → ``[T,hidden]``。 + - is_attn=False：logits → ``[N, V]``（vocab 恒在最后一维）。 + 返回 (on_vec, off_vec)（同 token、同向量长度），或 None（数据缺失 / pos 越界 / + 对齐失败）。 + """ + if on_tensor is None or off_tensor is None: + return None + if is_attn: + on = on_tensor.squeeze(1) if on_tensor.dim() == 3 else on_tensor + off = off_tensor.squeeze(1) if off_tensor.dim() == 3 else off_tensor + else: + on = on_tensor.reshape(-1, on_tensor.size(-1)) + off = off_tensor.reshape(-1, off_tensor.size(-1)) + if align_mask is not None and on.shape[0] != off.shape[0]: + try: + on, off = _align_packed(on, off, align_mask) + except ValueError: + return None + n = min(on.shape[0], off.shape[0]) + if pos < 0 or pos >= n: + return None + return on[pos].contiguous(), off[pos].contiguous() def _logits_ensure_token_major(lo: torch.Tensor, lf: torch.Tensor @@ -358,14 +387,22 @@ def cmp_attn_layer(dir_on: str, dir_off: str, def cmp_packed_token(dir_on: str, dir_off: str, - pos: int = 0, layer: int | None = None) -> list[CheckResult]: - """packed[pos] 对比：attn[pos]（可指定层）+ logits[pos]（仅最后一层）。 + pos: int = 0, layer: int | None = None, + align_mask: torch.Tensor | None = None) -> list[CheckResult]: + """packed[pos] 对比（**suffix 对齐后**）：attn[pos]（可指定层）+ logits[pos]（仅最后一层）。 - - pos：packed 里的 token 位置（单个 int，默认 0）。不是 2D 位置。 + ON 是裁剪后的 suffix-only packed，OFF 是完整 packed，两者 token **不直接对应**—— + 必须先用 align_mask（OFF cu_seqlens + ON prefix_lens）把 OFF 的 suffix 段抽出来 + 与 ON 对齐，再取 [pos]。pos 索引的是对齐后的 suffix-packed 空间（ON/OFF 一致）。 + + - pos：对齐后 suffix-packed 里的位置（单个 int，默认 0）。 - attn：用 *layer*（默认最后一层）。对比第 1 层可区分 - "结构错（位置/RoPE/mask，第 1 层就偏）" vs "数值累积（第 1 层完美、深层才偏）"。 - - logits：永远最后一层。logits 只在最后产生，不受 *layer* 影响。 + "结构错（第 1 层就偏）" vs "数值累积（第 1 层完美、深层才偏）"。 + - logits：永远最后一层。 + - align_mask：可选，复用调用方已构建的；None 则内部构建。 """ + if align_mask is None: + align_mask = _build_attn_align_mask(dir_on, dir_off) results: list[CheckResult] = [] attn_layer = layer if layer is not None else ( _get_num_layers(dir_on) or _get_num_layers(dir_off)) @@ -373,20 +410,27 @@ def cmp_packed_token(dir_on: str, dir_off: str, if attn_layer: a = _load_attn_output(dir_on, attn_layer) b = _load_attn_output(dir_off, attn_layer) - if a is not None and b is not None: - a0 = a.squeeze(1) if a.dim() == 3 else a - b0 = b.squeeze(1) if b.dim() == 3 else b + vecs = _aligned_vec_at_pos(a, b, True, pos, align_mask) + if vecs is None: results.append(CheckResult( name=f"attn_L{attn_layer}_pos{pos}", - metrics=_vec_metrics(a0[pos], b0[pos]))) + metrics={"error": f"无法对齐或 pos {pos} 越界"})) + else: + results.append(CheckResult( + name=f"attn_L{attn_layer}_pos{pos}", + metrics=_vec_metrics(vecs[0], vecs[1]))) lo = _load_logits(dir_on) lf = _load_logits(dir_off) - if lo is not None and lf is not None: - lo_p, lf_p = _logits_at_pos(lo, lf, pos) + vecs = _aligned_vec_at_pos(lo, lf, False, pos, align_mask) + if vecs is None: results.append(CheckResult( name=f"logits_pos{pos}", - metrics=_vec_metrics(lo_p, lf_p))) + metrics={"error": f"无法对齐或 pos {pos} 越界"})) + else: + results.append(CheckResult( + name=f"logits_pos{pos}", + metrics=_vec_metrics(vecs[0], vecs[1]))) return results @@ -744,6 +788,9 @@ def _print_packed_token(r: CheckResult): print(_SEP_SINGLE + f"\n [packed_token] {r.name}") print(_SEP_SINGLE) m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n") + return for k in ["mean_abs", "max_abs", "rel_max", "rel_mean", "cos", "pearson"]: v = m.get(k) if v is not None: @@ -936,31 +983,32 @@ def main(): all_results.append(r) _print_per_layer(r) - # ── packed: packed_token（attn[pos] + logits[pos]） ── - # pos 由 --token 指定（默认 0）；attn 用 --layer 指定的层（默认最后一层）； - # logits 永远最后一层（logits 只在最后产生）。 + # ── packed: packed_token（attn[pos] + logits[pos]，suffix 对齐后） ── + # pos 由 --token 指定（默认 0，索引对齐后的 suffix-packed 空间）； + # attn 用 --layer 指定的层（默认最后一层）；logits 永远最后一层。 pos = args.token - pt_results = cmp_packed_token(args.dir_on, args.dir_off, pos, args.layer) + align_mask = _build_attn_align_mask(args.dir_on, args.dir_off) + pt_results = cmp_packed_token(args.dir_on, args.dir_off, pos, args.layer, + align_mask=align_mask) for r in pt_results: all_results.append(r) _print_packed_token(r) - # packed_token top-K（per-dim） + # packed_token top-K（per-dim，同样先对齐再取 [pos]） if args.topk > 0 and pt_results: attn_layer = args.layer if args.layer is not None else ( _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) if attn_layer: a = _load_attn_output(args.dir_on, attn_layer) b = _load_attn_output(args.dir_off, attn_layer) - if a is not None and b is not None: - ap = (a.squeeze(1) if a.dim() == 3 else a)[pos].cpu() - bp = (b.squeeze(1) if b.dim() == 3 else b)[pos].cpu() - _print_topk_vec(ap, bp, args.topk, "val", + vecs = _aligned_vec_at_pos(a, b, True, pos, align_mask) + if vecs is not None: + _print_topk_vec(vecs[0].cpu(), vecs[1].cpu(), args.topk, "val", f"attn_L{attn_layer}_pos{pos}") lo = _load_logits(args.dir_on) lf = _load_logits(args.dir_off) - if lo is not None and lf is not None: - lo_p, lf_p = _logits_at_pos(lo, lf, pos) - _print_topk_vec(lo_p.cpu(), lf_p.cpu(), args.topk, + vecs = _aligned_vec_at_pos(lo, lf, False, pos, align_mask) + if vecs is not None: + _print_topk_vec(vecs[0].cpu(), vecs[1].cpu(), args.topk, "val", f"logits_pos{pos}", show_rel=False) # ── packed: logits（suffix 对齐） ── From 7f80253bf8d1f540ff8736bb5fed0876733de164 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 11:38:50 +0800 Subject: [PATCH 10/61] =?UTF-8?q?[fix]=20cmp=5Fdiag=5Fverl080:=20logits=20?= =?UTF-8?q?=E7=9A=84=20val=20=E6=8E=92=E5=BA=8F=E7=94=A8=E5=B8=A6=E7=AC=A6?= =?UTF-8?q?=E5=8F=B7=E5=AE=9E=E9=99=85=E5=80=BC=EF=BC=8C=E4=B8=8D=E7=94=A8?= =?UTF-8?q?=E7=BB=9D=E5=AF=B9=E5=80=BC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit abs 排序会把大负值 logit（softmax 后概率极低、不会被选中的 token）顶到表头，掩盖真正的高 logit 候选。改为 max(on, off) 带符号值排序，让候选 token 排前面。 仅改 1D _print_topk_vec 的 val 分支（logits/attn 表）；2D logp/entropy 表的 val 仍用 abs（logp 本身为负，abs 有意义）。 --- prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index a4b97fc0..1105c399 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -840,8 +840,11 @@ def _print_topk_vec(on_vec: torch.Tensor, off_vec: torch.Tensor, sort_key = abs_err elif sort_by == "rel": sort_key = rel_err - else: # "val" - sort_key = torch.maximum(on_vec.abs(), off_vec.abs()) + else: # "val" —— 带符号的实际值，不是绝对值 + # 对 logits：绝对值大但符号为负的 logit，softmax 后概率极低、不会被选中。 + # 按 abs 排会把这种"必不选"的 token 顶到表头，掩盖真正的高 logit 候选。 + # 改用 max(on, off) 带符号值，让真正的高 logit（候选 token）排前面。 + sort_key = torch.maximum(on_vec, off_vec) _, idx = sort_key.topk(min(topk, sort_key.numel())) idx = idx.to(torch.long) print(f"\n [{label}] top-{topk} dims (sort by {sort_by})") From c0de0972651f078805dbd5ca113531b733cc8f7e Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 15:28:36 +0800 Subject: [PATCH 11/61] =?UTF-8?q?[diag]=20PREFIX=5FSHARING=5FFORCE=5FZERO?= =?UTF-8?q?=5FPREFIX:=20=E5=BC=BA=E5=88=B6=200-prefix=20=E9=9A=94=E7=A6=BB?= =?UTF-8?q?=20kernel-mode=20=E5=B7=AE=E5=BC=82?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增环境变量 PREFIX_SHARING_FORCE_ZERO_PREFIX： - prefix_detector.detect() 命中即早返回 0-prefix（所有 provider、无复用） - verl_mcore build 旁路 has_sharing 早返回，让 ON pipeline 仍跑全路径： trim(0-prefix→不裁)/build_kv(全 provider→无注入)/attention(mode 3 全序列)/restore(no-op) - 用途：跑 ON(0-prefix, mode 3, 全序列, 无裁剪/注入/位置补偿) vs OFF(mode 2, 全序列) ≈ → kernel mode 不是根因；偏差大 → kernel mode 是根因。 --- .../prefix_sharing/core/prefix_detector.py | 14 ++++++++++++++ .../prefix_sharing/integrations/verl_mcore.py | 13 ++++++++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/prefix-sharing/prefix_sharing/core/prefix_detector.py b/prefix-sharing/prefix_sharing/core/prefix_detector.py index 3d40a37d..fdf9edc8 100644 --- a/prefix-sharing/prefix_sharing/core/prefix_detector.py +++ b/prefix-sharing/prefix_sharing/core/prefix_detector.py @@ -180,6 +180,20 @@ def detect(self, input_ids: Sequence[TokenSequence]) -> PrefixDetectionResult: group_key_to_id: dict[tuple[int, int], int] = {} group_members: dict[int, list[int]] = {} + # 诊断开关：跳过 trie 检测，直接返回 0-prefix（所有序列当 provider，无复用）。 + # 用法见 PREFIX_SHARING_FORCE_ZERO_PREFIX 说明（配合 build 层的 has_sharing 旁路）。 + import os as _os + if _os.environ.get("PREFIX_SHARING_FORCE_ZERO_PREFIX"): + return PrefixDetectionResult( + batch_size=batch_size, + reuse_specs=(), + groups=(), + group_ids=tuple(group_ids), + provider_index=tuple(provider_index), + prefix_lens=tuple(prefix_lens), + is_provider=tuple(is_provider), + ) + for index, seq in enumerate(input_ids): node = root matched = 0 diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index 7412bc43..d089efb6 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -641,7 +641,18 @@ def build_prefix_sharing_micro_batch_verl080( # ── 阶段 4: 前缀共享规划 ── plan = PrefixSharingPlanner(ps_config).plan(sequences) - if not plan.has_sharing: + + # 诊断开关 PREFIX_SHARING_FORCE_ZERO_PREFIX：detect() 被旁路返回 0-prefix，此处 + # plan.has_sharing=False。仍要让 ON pipeline 跑（prefix_attention 用 mode-3 跑全 + # 序列、无裁剪/注入），故旁路 has_sharing 早返回。用于隔离 "kernel-mode 差异"： + # ON(0-prefix, mode 3, 全序列) vs OFF(mode 2, 全序列) + # ≈ → kernel mode 不是根因，偏差来自裁剪/注入；偏差大 → kernel mode 是根因。 + import os as _os_force_zero + _force_zero_prefix = bool(_os_force_zero.environ.get("PREFIX_SHARING_FORCE_ZERO_PREFIX")) + if _force_zero_prefix: + print("[PS][prepare] FORCE_ZERO_PREFIX: 跳过 has_sharing 早返回，" + "ON pipeline 用 mode-3 跑全序列（无裁剪/注入）") + elif not plan.has_sharing: print("[PS][prepare] no prefix sharing detected") return batch, None From 6b122afa6f3b003140e6d80fac25535e0973c258 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 16:52:25 +0800 Subject: [PATCH 12/61] [diag] prefix-sharing deviation probes: RoPE extrapolation log + ground-truth + skip-restore MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit S1.1: RoPE 外推自动诊断日志（step_min/max/mean/std + step_is_uniform） - 不需要 env var，每次外推触发时在 megatron_runtime.py 自动打印 - step_is_uniform=False → 非线性 RoPE 调度，外推可能引入误差 S1.2: PREFIX_SHARING_DIAG_ROPE_GROUND_TRUTH=1 - 从 inv_freq 直接计算完整频率表，跳过线性外推 - 启用后 ON vs OFF 一致 → RoPE 外推是根因 S4.3: PREFIX_SHARING_DIAG_SKIP_RESTORE=1 - 跳过 logprobs restore，直接返回 forward 原始输出 - suffix 一致 → restore 是根因；不一致 → 偏差在 forward Co-Authored-By: Claude Fable 5 --- .../integrations/megatron_runtime.py | 84 +++++++++++++++++++ .../prefix_sharing/integrations/verl_mcore.py | 10 +++ 2 files changed, 94 insertions(+) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 01ddb36d..e72699e0 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -173,6 +173,27 @@ def _apply_positioned_rope( if q_pos_emb is not None and max_needed > q_pos_emb.shape[0]: dim_half = q_pos_emb.shape[-1] // 2 step = q_pos_emb[1:2, :, :, :dim_half] - q_pos_emb[0:1, :, :, :dim_half] + + # ── [PS-diag] RoPE extrapolation diagnostic log ── + _layer = getattr(attention_module, 'layer_number', -1) + _num_extra = int(max_needed - q_pos_emb.shape[0]) + _step_vals = step.detach().flatten() + print( + f"[PS][RoPE-extrapolate-Q] layer={_layer} " + f"max_needed={max_needed} precomputed={q_pos_emb.shape[0]} extra={_num_extra} " + f"step_min={_step_vals.min().item():.8f} step_max={_step_vals.max().item():.8f} " + f"step_mean={_step_vals.mean().item():.8f} step_std={_step_vals.std().item():.8f} " + f"step_is_uniform={bool((_step_vals.max() - _step_vals.min()).abs() < 1e-8)}", + flush=True, + ) + # ── [PS-diag] end ── + + # RoPE 具有线性性质：freqs[p] = p * inv_freq。 + # 因此可以通过 pos_emb[1] - pos_emb[0] 恢复出 step（即 inv_freq）， + # 从而生成缺失的高位置频率向量。 + # 注意：这对标准 sin/cos RoPE 成立，但对 NTK-aware/YaRN 等非线性 + # 频率调度不成立。如果 step_std > 0 且 step_is_uniform=False， + # 说明使用了非线性 RoPE 调度，线性外推会引入数值误差。 extra_positions = torch.arange( q_pos_emb.shape[0], max_needed, device=q_pos_emb.device, dtype=q_pos_emb.dtype, @@ -183,6 +204,21 @@ def _apply_positioned_rope( if k_pos_emb is not None and max_needed > k_pos_emb.shape[0]: dim_half = k_pos_emb.shape[-1] // 2 step = k_pos_emb[1:2, :, :, :dim_half] - k_pos_emb[0:1, :, :, :dim_half] + + # ── [PS-diag] RoPE extrapolation diagnostic log ── + _layer = getattr(attention_module, 'layer_number', -1) + _num_extra = int(max_needed - k_pos_emb.shape[0]) + _step_vals = step.detach().flatten() + print( + f"[PS][RoPE-extrapolate-K] layer={_layer} " + f"max_needed={max_needed} precomputed={k_pos_emb.shape[0]} extra={_num_extra} " + f"step_min={_step_vals.min().item():.8f} step_max={_step_vals.max().item():.8f} " + f"step_mean={_step_vals.mean().item():.8f} step_std={_step_vals.std().item():.8f} " + f"step_is_uniform={bool((_step_vals.max() - _step_vals.min()).abs() < 1e-8)}", + flush=True, + ) + # ── [PS-diag] end ── + extra_positions = torch.arange( k_pos_emb.shape[0], max_needed, device=k_pos_emb.device, dtype=k_pos_emb.dtype, @@ -191,6 +227,54 @@ def _apply_positioned_rope( extra_emb = torch.cat([extra_angles, extra_angles], dim=-1) k_pos_emb = torch.cat([k_pos_emb, extra_emb], dim=0) + # ── [PS-diag] RoPE ground-truth probe ── + # 当 PREFIX_SHARING_DIAG_ROPE_GROUND_TRUTH=1 时，从 inv_freq 直接计算完整 + # 频率表（跳过线性外推），用于验证外推是否引入数值偏差。 + # 如果启用此开关后 ON vs OFF 结果一致，RoPE 外推就是根因。 + import os as _os_gt + if _os_gt.environ.get("PREFIX_SHARING_DIAG_ROPE_GROUND_TRUTH"): + _layer_gt = getattr(attention_module, 'layer_number', -1) + # 尝试从 attention_module 获取 inv_freq + _inv_freq = None + _rotary_emb = getattr(attention_module, 'rotary_pos_emb', None) + if _rotary_emb is not None: + _inv_freq = getattr(_rotary_emb, 'inv_freq', None) + if _inv_freq is None: + # fallback: 从 config 读取 RoPE 参数 + _cfg = attention_module.config + _dim = getattr(_cfg, 'hidden_size', 4096) // getattr(_cfg, 'num_attention_heads', 32) + _base = getattr(_cfg, 'rope_theta', 10000.0) + _inv_freq = 1.0 / (_base ** (torch.arange( + 0, _dim, 2, device=q_pos_emb.device if q_pos_emb is not None + else k_pos_emb.device).float() / _dim)) + + _device = q_pos_emb.device if q_pos_emb is not None else k_pos_emb.device + _dtype = q_pos_emb.dtype if q_pos_emb is not None else k_pos_emb.dtype + _all_positions = torch.arange(0, max_needed, device=_device, dtype=torch.float) + _freqs = torch.outer(_all_positions, _inv_freq.to(_device).float()) # [max_needed, dim/2] + _emb_gt = torch.cat([_freqs, _freqs], dim=-1) # [max_needed, dim] + # reshape 匹配 pos_emb 维度 [max_needed, 1, 1, dim] + _emb_gt = _emb_gt.unsqueeze(1).unsqueeze(1).to(_dtype) + + _extra_old = max(0, int(max_needed - ( + q_pos_emb.shape[0] if q_pos_emb is not None else max_needed))) + _is_q_truncated = q_pos_emb is not None and q_pos_emb.shape[0] < max_needed + _is_k_truncated = k_pos_emb is not None and k_pos_emb.shape[0] < max_needed + + if q_pos_emb is not None: + q_pos_emb = _emb_gt + if k_pos_emb is not None: + k_pos_emb = _emb_gt + + print( + f"[PS][RoPE-ground-truth] layer={_layer_gt} " + f"max_needed={max_needed} emb_shape={_emb_gt.shape} " + f"was_extrapolated_q={_is_q_truncated} was_extrapolated_k={_is_k_truncated} " + f"old_extra_count={_extra_old}", + flush=True, + ) + # ── [PS-diag] end ── + # Build kwargs for apply_rotary_pos_emb. # Only include version-specific params when they're provided, # to maintain backward compat with v070 (mcore <= 0.15.x). diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index d089efb6..2743e1c7 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -393,6 +393,16 @@ def restore_via_2d_unfold_verl080( if ctx is None: return output plan = ctx.prefix_sharing_plan + # ── [PS-diag] skip-restore probe ── + # PREFIX_SHARING_DIAG_SKIP_RESTORE=1 时跳过 restore，直接返回 forward 原始输出。 + # 用于隔离：如果 skip restore 后 suffix logprobs 与 OFF 一致，偏差在 restore； + # 如果 suffix logprobs 仍不一致，偏差在 forward（attention/build_kv）。 + import os as _os_skip + if _os_skip.environ.get("PREFIX_SHARING_DIAG_SKIP_RESTORE"): + print("[PS][diag] SKIP_RESTORE: 跳过 restore，直接返回 forward 原始输出") + return output + # ── [PS-diag] end ── + # Guard on reuser presence, not on prefix_last_restore_indices: a batch # whose reusers all have suffix_len == 0 has no prefix-last spec but still # needs interior prefix columns restored. From 717d9940f62715ef4a9727772fb9c97f9ac371be Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 17:21:16 +0800 Subject: [PATCH 13/61] [diag] RoPE extrapolation: add step01 vs step12 linearity check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 新增 step01_vs_step12_maxdiff / is_linear 字段： - 计算 pos_emb[0→1] 和 pos_emb[1→2] 的 step 差值 - maxdiff ≈ 0 → pos_emb[p] 对 p 严格线性 → 外推数学正确 - maxdiff > 0 → pos_emb[p] 对 p 非线性 → 外推引入误差 Co-Authored-By: Claude Fable 5 --- .../integrations/megatron_runtime.py | 26 ++++++++++++++----- 1 file changed, 19 insertions(+), 7 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index e72699e0..afb53b19 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -174,16 +174,24 @@ def _apply_positioned_rope( dim_half = q_pos_emb.shape[-1] // 2 step = q_pos_emb[1:2, :, :, :dim_half] - q_pos_emb[0:1, :, :, :dim_half] - # ── [PS-diag] RoPE extrapolation diagnostic log ── + # ── [PS-diag] RoPE extrapolation: 验证相邻位置 step 是否恒定 ── _layer = getattr(attention_module, 'layer_number', -1) _num_extra = int(max_needed - q_pos_emb.shape[0]) _step_vals = step.detach().flatten() + # 关键诊断: step01 == step12 ? (pos_emb[p] 对 p 是否线性) + _step01_vs_step12_diff = None + if q_pos_emb.shape[0] >= 3: + _step12 = q_pos_emb[2:3, :, :, :dim_half] - q_pos_emb[1:2, :, :, :dim_half] + _step12_vals = _step12.detach().flatten() + _diff = (_step_vals - _step12_vals).abs() + _step01_vs_step12_diff = _diff.max().item() + _is_linear = (_step01_vs_step12_diff is not None and _step01_vs_step12_diff < 1e-8) print( f"[PS][RoPE-extrapolate-Q] layer={_layer} " f"max_needed={max_needed} precomputed={q_pos_emb.shape[0]} extra={_num_extra} " f"step_min={_step_vals.min().item():.8f} step_max={_step_vals.max().item():.8f} " f"step_mean={_step_vals.mean().item():.8f} step_std={_step_vals.std().item():.8f} " - f"step_is_uniform={bool((_step_vals.max() - _step_vals.min()).abs() < 1e-8)}", + f"step01_vs_step12_maxdiff={_step01_vs_step12_diff} is_linear={_is_linear}", flush=True, ) # ── [PS-diag] end ── @@ -191,9 +199,6 @@ def _apply_positioned_rope( # RoPE 具有线性性质：freqs[p] = p * inv_freq。 # 因此可以通过 pos_emb[1] - pos_emb[0] 恢复出 step（即 inv_freq）， # 从而生成缺失的高位置频率向量。 - # 注意：这对标准 sin/cos RoPE 成立，但对 NTK-aware/YaRN 等非线性 - # 频率调度不成立。如果 step_std > 0 且 step_is_uniform=False， - # 说明使用了非线性 RoPE 调度，线性外推会引入数值误差。 extra_positions = torch.arange( q_pos_emb.shape[0], max_needed, device=q_pos_emb.device, dtype=q_pos_emb.dtype, @@ -205,16 +210,23 @@ def _apply_positioned_rope( dim_half = k_pos_emb.shape[-1] // 2 step = k_pos_emb[1:2, :, :, :dim_half] - k_pos_emb[0:1, :, :, :dim_half] - # ── [PS-diag] RoPE extrapolation diagnostic log ── + # ── [PS-diag] RoPE extrapolation: 验证相邻位置 step 是否恒定 ── _layer = getattr(attention_module, 'layer_number', -1) _num_extra = int(max_needed - k_pos_emb.shape[0]) _step_vals = step.detach().flatten() + _step01_vs_step12_diff = None + if k_pos_emb.shape[0] >= 3: + _step12 = k_pos_emb[2:3, :, :, :dim_half] - k_pos_emb[1:2, :, :, :dim_half] + _step12_vals = _step12.detach().flatten() + _diff = (_step_vals - _step12_vals).abs() + _step01_vs_step12_diff = _diff.max().item() + _is_linear = (_step01_vs_step12_diff is not None and _step01_vs_step12_diff < 1e-8) print( f"[PS][RoPE-extrapolate-K] layer={_layer} " f"max_needed={max_needed} precomputed={k_pos_emb.shape[0]} extra={_num_extra} " f"step_min={_step_vals.min().item():.8f} step_max={_step_vals.max().item():.8f} " f"step_mean={_step_vals.mean().item():.8f} step_std={_step_vals.std().item():.8f} " - f"step_is_uniform={bool((_step_vals.max() - _step_vals.min()).abs() < 1e-8)}", + f"step01_vs_step12_maxdiff={_step01_vs_step12_diff} is_linear={_is_linear}", flush=True, ) # ── [PS-diag] end ── From c86ea1242d089e4c6ece64167504f763b0066431 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 17:32:48 +0800 Subject: [PATCH 14/61] [diag] add post-RoPE Q/K dump for both ON and OFF paths ON path: dump_rope_emb_layer() called in _apply_positioned_rope after RoPE OFF path: extract QKV, reconstruct standard 0..seg_len-1 position IDs, apply_rotary_pos_emb manually, dump via dump_rope_emb_layer() Both paths write to rope_emb.pt for cmp_diag comparison. Co-Authored-By: Claude Fable 5 --- .../integrations/megatron_runtime.py | 13 ++++++ .../verl080_mcore0161_ms0160/attention.py | 45 ++++++++++++++----- 2 files changed, 47 insertions(+), 11 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index afb53b19..942e2236 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -326,6 +326,19 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: k_freqs, **_rope_kwargs(cu_seqlens_kv), ).squeeze(1) + + ######### prefix-sharing diag: ON post-RoPE Q/K dump (per-layer) ######### + try: + from prefix_sharing.tools.diagnostic_dump import dump_rope_emb_layer + dump_rope_emb_layer( + attention_module.layer_number, + query, key, + attention_module.config.num_layers, + ) + except Exception as e: + print(f"rope_emb_layer dump failed: {e}") + ######### prefix-sharing diag: ON post-RoPE Q/K dump end ######### + return query, key diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 97ef0c1e..0a909c61 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -51,20 +51,14 @@ def patched_forward( sequence_len_offset=sequence_len_offset, inference_params=inference_params, ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs_off dump ##### - # OFF 走原始 forward，不经 prefix_attention/_apply_positioned_rope， - # 所以 ON 路径里的 dump_attn_on/dump_rope_freqs_on 不会触发。 - # 这里在 OFF 分支补 dump，让 cmp_diag 的 attn/RoPE 对比有 OFF ground truth。 - # v070 是直接改 megatron attention 源码在 forward 内部 dump；v080 用 patch - # wrapper 在 forward 返回后 dump output + 入参 rotary_pos_emb 解包出 angle table， - # 语义等价（唯一拿不到的是 rope_emb rotated q/k，在 forward 内部，但 rope_freqs - # angle table 已够验证 RoPE）。 + # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump ##### import os as _os if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump import ( - dump_attn_off, dump_rope_freqs_off, + dump_attn_off, dump_rope_freqs_off, dump_rope_emb_layer, ) from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb + from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result _bs = ( len(packed_seq_params.cu_seqlens_q_padded) - 1 @@ -75,9 +69,38 @@ def patched_forward( dump_attn_off(_attn_out, packed_seq_params, self.layer_number, _bs, self.config.num_layers) if rotary_pos_emb is not None: - _q_pos_emb, _ = _unpack_rotary_pos_emb(rotary_pos_emb) + _q_pos_emb, _k_pos_emb = _unpack_rotary_pos_emb(rotary_pos_emb) dump_rope_freqs_off(_q_pos_emb, self.layer_number, self.config.num_layers) - # ##### [PS-diag] OFF attn_outputs + rope_freqs_off dump end ##### + # OFF post-RoPE Q/K dump: 提取 QKV + 手动 apply RoPE 后 dump，供与 ON 对比 + # OFF 路径 position IDs 是标准的 0..seg_len-1（每 segment 内连续） + if (packed_seq_params is not None + and packed_seq_params.qkv_format == "thd" + and hasattr(packed_seq_params, "cu_seqlens_q_padded")): + _off_q, _off_k, _off_v = self.get_query_key_value_tensors( + hidden_states, key_value_states, + split_qkv=True, output_gate=False, + ) + _off_q = _off_q.squeeze(1) + _off_k = _off_k.squeeze(1) + _cu = packed_seq_params.cu_seqlens_q_padded + _pos_list = [torch.arange(_cu[i+1] - _cu[i], device=_off_q.device) + for i in range(len(_cu) - 1)] + _off_positions = torch.cat(_pos_list, dim=0).long() + _off_q_freqs = _q_pos_emb.index_select(0, _off_positions) + _off_k_freqs = (_k_pos_emb or _q_pos_emb).index_select(0, _off_positions) + _off_q_rope = apply_rotary_pos_emb( + _off_q.unsqueeze(1), _off_q_freqs, + config=self.config, cu_seqlens=None, + ).squeeze(1) + _off_k_rope = apply_rotary_pos_emb( + _off_k.unsqueeze(1), _off_k_freqs, + config=self.config, cu_seqlens=None, + ).squeeze(1) + dump_rope_emb_layer( + self.layer_number, _off_q_rope, _off_k_rope, + self.config.num_layers, + ) + # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump end ##### return _result # ── prefix-sharing path ── From f5e9d09a283665fd18901c0bbfd964c72bb09549 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 17:49:22 +0800 Subject: [PATCH 15/61] [diag] rope_emb: add position IDs to dump + rewrite cmp_rope_emb for position-aligned comparison MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1. rope_emb.pt 新增 positions 字段（per-token absolute position IDs） 2. cmp_rope_emb 重写: - first packed token Q/K 直接对比（row 0, pos 0，provider 必相等） - first row Q/K 按 position ID 对齐后逐位对比（0..L-1，两路径同） - reuser row 不直接对比（ON absolute pos vs OFF relative pos，设计差异） 3. ON 调用点传入 packed_position_ids 4. OFF 调用点传入重建的 segment-relative positions Co-Authored-By: Claude Fable 5 --- .../integrations/megatron_runtime.py | 1 + .../verl080_mcore0161_ms0160/attention.py | 1 + .../prefix_sharing/tools/cmp_diag.py | 102 ++++++++++++++++-- .../prefix_sharing/tools/diagnostic_dump.py | 9 +- 4 files changed, 102 insertions(+), 11 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 942e2236..b13f2e89 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -334,6 +334,7 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: attention_module.layer_number, query, key, attention_module.config.num_layers, + positions=packed_position_ids, ) except Exception as e: print(f"rope_emb_layer dump failed: {e}") diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 0a909c61..20488b34 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -99,6 +99,7 @@ def patched_forward( dump_rope_emb_layer( self.layer_number, _off_q_rope, _off_k_rope, self.config.num_layers, + positions=_off_positions, ) # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump end ##### return _result diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag.py b/prefix-sharing/prefix_sharing/tools/cmp_diag.py index 7feadafe..e9c04062 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag.py @@ -387,9 +387,20 @@ def cmp_position_ids(dir_a: str, dir_b: str) -> CheckResult | None: # 2. RoPE encoding — absolute equality # ══════════════════════════════════════════════════════════════════ -def cmp_rope_emb(dir_a: str, dir_b: str) -> CheckResult | None: - fa = os.path.join(dir_a, "rope_emb.pt") - fb = os.path.join(dir_b, "rope_emb.pt") +def cmp_rope_emb(dir_on: str, dir_off: str) -> CheckResult | None: + """Compare post-RoPE Q/K between ON and OFF, aligned by position IDs. + + ON positions are absolute (preserved from original input). + OFF positions are per-segment relative (0..L-1 for each row). + + Comparison strategy: + 1. First packed token (row 0, pos 0): always a provider, must match exactly. + 2. Within the first row: match tokens by position ID (0..L-1 for both paths). + 3. For reuser rows: ON uses absolute positions (prefix_len..), OFF uses + relative (0..). Positions differ by design — skip direct comparison. + """ + fa = os.path.join(dir_on, "rope_emb.pt") + fb = os.path.join(dir_off, "rope_emb.pt") if not os.path.exists(fa) or not os.path.exists(fb): return None a = torch.load(fa, weights_only=True) @@ -400,12 +411,87 @@ def cmp_rope_emb(dir_a: str, dir_b: str) -> CheckResult | None: if la != lb: return CheckResult(name="rope_emb", passed=False, metrics={"error": "layer set mismatch"}) - md = 0.0 + + max_diff_q = 0.0 + max_diff_k = 0.0 + first_token_q_diff = 0.0 + first_token_k_diff = 0.0 + first_row_q_diff = 0.0 + first_row_k_diff = 0.0 + first_row_len = 0 + num_layers = len(la) + for lyr in sorted(la): - for k in ("query", "key"): - md = max(md, float((a[lyr][k] - b[lyr][k]).abs().max())) - return CheckResult(name="rope_emb", passed=md == 0.0, - metrics={"max_diff": md, "num_layers": len(la)}) + on_entry = a[lyr] + off_entry = b[lyr] + on_q, on_k = on_entry["query"], on_entry["key"] + off_q, off_k = off_entry["query"], off_entry["key"] + on_pos = on_entry.get("positions") + off_pos = off_entry.get("positions") + + # ── Check 1: first packed token (row 0, position 0) ── + first_token_q_diff = max(first_token_q_diff, + float((on_q[0] - off_q[0]).abs().max())) + first_token_k_diff = max(first_token_k_diff, + float((on_k[0] - off_k[0]).abs().max())) + + # ── Check 2: first row — match by ON position IDs ── + if on_pos is not None and off_pos is not None: + on_pos_t = on_pos.long() + off_pos_t = off_pos.long() + # Find first-row extent in ON: tokens before the first position reset + # (position decreases or jumps to prefix_start) + _on_row1_end = 1 + for _i in range(1, len(on_pos_t)): + if on_pos_t[_i] <= on_pos_t[_i - 1]: + break + _on_row1_end = _i + 1 + _on_row1_len = _on_row1_end # positions 0..L-1 + + # First row in OFF: positions go 0..L-1 (find matching extent) + _off_row1_end = 1 + for _i in range(1, len(off_pos_t)): + if off_pos_t[_i] <= off_pos_t[_i - 1]: + break + _off_row1_end = _i + 1 + _off_row1_len = _off_row1_end + + _cmp_len = min(_on_row1_len, _off_row1_len) + if _cmp_len > 0: + first_row_len = max(first_row_len, _cmp_len) + # Match by position ID within the first row + for _pos_id in range(_cmp_len): + _on_idx = _pos_id # ON row 1 starts at packed index 0 + _off_idx = _pos_id # OFF row 1 starts at packed index 0 + first_row_q_diff = max(first_row_q_diff, + float((on_q[_on_idx] - off_q[_off_idx]).abs().max())) + first_row_k_diff = max(first_row_k_diff, + float((on_k[_on_idx] - off_k[_off_idx]).abs().max())) + max_diff_q = max(max_diff_q, first_row_q_diff) + max_diff_k = max(max_diff_k, first_row_k_diff) + else: + # Fallback: no positions available, compare first row by direct offset + _cmp_len = min(on_q.shape[0], off_q.shape[0], 128) + first_row_len = _cmp_len + first_row_q_diff = float((on_q[:_cmp_len] - off_q[:_cmp_len]).abs().max()) + first_row_k_diff = float((on_k[:_cmp_len] - off_k[:_cmp_len]).abs().max()) + max_diff_q = first_row_q_diff + max_diff_k = first_row_k_diff + + # Pass if all diffs are near zero + threshold = 1e-7 + passed = (first_token_q_diff < threshold and first_token_k_diff < threshold + and first_row_q_diff < threshold and first_row_k_diff < threshold) + return CheckResult(name="rope_emb", passed=passed, metrics={ + "num_layers": num_layers, + "first_row_len": first_row_len, + "first_token_q_maxdiff": first_token_q_diff, + "first_token_k_maxdiff": first_token_k_diff, + "first_row_q_maxdiff": first_row_q_diff, + "first_row_k_maxdiff": first_row_k_diff, + "max_diff_q": max_diff_q, + "max_diff_k": max_diff_k, + }) def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py index aa0541cc..308a2401 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py @@ -250,7 +250,8 @@ def _flush_attn_buffer(dump_dir: str) -> None: def _add_to_rope_buffer(layer_number: int, rotated_query: torch.Tensor, - rotated_key: torch.Tensor) -> None: + rotated_key: torch.Tensor, + positions: torch.Tensor | None = None) -> None: """Accumulate one layer's RoPE encoding into the global buffer.""" global _ROPE_BUFFER if _ROPE_BUFFER is None: @@ -258,6 +259,7 @@ def _add_to_rope_buffer(layer_number: int, rotated_query: torch.Tensor, _ROPE_BUFFER[layer_number] = { "query": rotated_query.detach().cpu().clone(), "key": rotated_key.detach().cpu().clone(), + "positions": positions.detach().cpu().clone() if positions is not None else None, } @@ -350,7 +352,8 @@ def dump_position_ids(position_ids: torch.Tensor) -> None: # ── RoPE encoding dump (per layer) ────────────────────────────── def dump_rope_emb_layer(layer_number: int, rotated_query: torch.Tensor, - rotated_key: torch.Tensor, num_layers: int) -> None: + rotated_key: torch.Tensor, num_layers: int, + positions: torch.Tensor | None = None) -> None: """Accumulate one layer's post-RoPE query/key. Auto-flush on last layer. Call after RoPE is applied in each attention layer, for both ON and OFF modes. @@ -358,7 +361,7 @@ def dump_rope_emb_layer(layer_number: int, rotated_query: torch.Tensor, dump_dir = _get_dump_dir() if dump_dir is None: return - _add_to_rope_buffer(layer_number, rotated_query, rotated_key) + _add_to_rope_buffer(layer_number, rotated_query, rotated_key, positions) if layer_number == num_layers: _flush_rope_buffer(dump_dir) From 64fc4e4e293a521069ae4a1664dcb93a67131c01 Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 19:24:45 +0800 Subject: [PATCH 16/61] =?UTF-8?q?[diag]=20verl080:=20post-RoPE=20Q/K=20dum?= =?UTF-8?q?p=20+=20cmp=20(=E5=A4=8D=E5=88=BB=20attn=5Foutput=20=E6=A8=A1?= =?UTF-8?q?=E5=BC=8F)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dump: diagnostic_dump_verl080.py - dump_rope_emb_verl080(layer, query, key, num_layers) - 格式: rope_emb.pt = {layer: {"query": [T,H,D], "key": [T,H,D]}} - ON suffix-only / OFF 完整 packed，不需要 positions（suffix 对齐用已有元数据） cmp: cmp_diag_verl080.py - cmp_rope_emb_layer(): per-layer Q/K cosine（suffix 对齐，同 cmp_attn_layer） - cmp_rope_emb_token(): 指定位置 Q/K 向量对比（同 cmp_packed_token） - _print_rope_emb_per_layer(): 每层 Q_cos + K_cos 报告 调用点: - ON: megatron_runtime.py → dump_rope_emb_verl080 - OFF: attention.py → dump_rope_emb_verl080 Co-Authored-By: Claude Fable 5 --- .../integrations/megatron_runtime.py | 4 +- .../verl080_mcore0161_ms0160/attention.py | 5 +- .../prefix_sharing/tools/cmp_diag_verl080.py | 232 ++++++++++++++++++ .../tools/diagnostic_dump_verl080.py | 37 +++ 4 files changed, 274 insertions(+), 4 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index b13f2e89..d2c523cd 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -329,8 +329,8 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: ######### prefix-sharing diag: ON post-RoPE Q/K dump (per-layer) ######### try: - from prefix_sharing.tools.diagnostic_dump import dump_rope_emb_layer - dump_rope_emb_layer( + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_emb_verl080 + dump_rope_emb_verl080( attention_module.layer_number, query, key, attention_module.config.num_layers, diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 20488b34..3407abd0 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -55,8 +55,9 @@ def patched_forward( import os as _os if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump import ( - dump_attn_off, dump_rope_freqs_off, dump_rope_emb_layer, + dump_attn_off, dump_rope_freqs_off, ) + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_emb_verl080 from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result @@ -96,7 +97,7 @@ def patched_forward( _off_k.unsqueeze(1), _off_k_freqs, config=self.config, cu_seqlens=None, ).squeeze(1) - dump_rope_emb_layer( + dump_rope_emb_verl080( self.layer_number, _off_q_rope, _off_k_rope, self.config.num_layers, positions=_off_positions, diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 1105c399..0e833087 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -434,6 +434,171 @@ def cmp_packed_token(dir_on: str, dir_off: str, return results +# ══════════════════════════════════════════════════════════════════ +# Post-RoPE Q/K compare: per-layer + packed_token +# ══════════════════════════════════════════════════════════════════ + +def _load_rope_emb(dir_path: str, layer: int + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Load post-RoPE Q/K for a single layer from rope_emb.pt. + + Returns ``(query, key)`` or ``(None, None)``. + """ + fp = os.path.join(dir_path, "rope_emb.pt") + if not os.path.exists(fp): + return None, None + d = torch.load(fp, weights_only=True) + if not isinstance(d, dict): + return None, None + entry = d.get(layer) + if entry is None: + return None, None + return entry.get("query"), entry.get("key") + + +def _rope_emb_cos_for_layer(qa: torch.Tensor, ka: torch.Tensor, + qb: torch.Tensor, kb: torch.Tensor, + align_mask: torch.Tensor | None = None) -> dict: + """单层 Q/K suffix 对齐 + per-token cosine（Q 和 K 分别算）。""" + # Q/K shape: [T, H, D] → 压平 head*dim 维度算 cosine + qa_flat = qa.reshape(qa.shape[0], -1) + qb_flat = qb.reshape(qb.shape[0], -1) + ka_flat = ka.reshape(ka.shape[0], -1) + kb_flat = kb.reshape(kb.shape[0], -1) + + if align_mask is not None and qa.shape[0] != qb.shape[0]: + qa_flat, qb_flat = _align_packed(qa_flat, qb_flat, align_mask) + ka_flat, kb_flat = _align_packed(ka_flat, kb_flat, align_mask) + + q_cos = _cosine_sim(qa_flat, qb_flat, dim=-1) + k_cos = _cosine_sim(ka_flat, kb_flat, dim=-1) + return { + "n_tokens": qa_flat.shape[0], + "Q_cos_avg": float(q_cos.mean()), "Q_cos_min": float(q_cos.min()), + "K_cos_avg": float(k_cos.mean()), "K_cos_min": float(k_cos.min()), + } + + +def cmp_rope_emb_layer(dir_on: str, dir_off: str, + layer: int | None) -> CheckResult | None: + """Post-RoPE Q/K per-layer cosine（suffix 对齐）。 + + 单层模式（layer 给定）：返回该层 Q/K cos。全层模式：返回所有层汇总。 + """ + align_mask = _build_attn_align_mask(dir_on, dir_off) + + if layer is not None: + q_on, k_on = _load_rope_emb(dir_on, layer) + q_off, k_off = _load_rope_emb(dir_off, layer) + if q_on is None or q_off is None: + return None + need = align_mask is not None and q_on.shape[0] != q_off.shape[0] + try: + d = _rope_emb_cos_for_layer(q_on, k_on, q_off, k_off, + align_mask if need else None) + except ValueError as e: + return CheckResult(name=f"rope_emb_L{layer}", passed=False, + metrics={"error": str(e)}) + d["layer"] = layer + ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS + and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) + return CheckResult(name=f"rope_emb_L{layer}", passed=ok, metrics=d) + + # All layers + fa = os.path.join(dir_on, "rope_emb.pt") + fb = os.path.join(dir_off, "rope_emb.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None + da = torch.load(fa, weights_only=True) + db = torch.load(fb, weights_only=True) + if not isinstance(da, dict) or not isinstance(db, dict): + return None + + results = {} + for lyr in sorted(set(da.keys()) & set(db.keys())): + ea, eb = da[lyr], db[lyr] + qa, ka = ea.get("query"), ea.get("key") + qb, kb = eb.get("query"), eb.get("key") + if qa is None or qb is None: + continue + need = align_mask is not None and qa.shape[0] != qb.shape[0] + try: + results[lyr] = _rope_emb_cos_for_layer(qa, ka, qb, kb, + align_mask if need else None) + except ValueError as e: + results[lyr] = {"error": str(e)} + return CheckResult(name="rope_emb_per_layer", passed=True, + metrics={"layers": results}) + + +def _rope_emb_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, + q_off: torch.Tensor | None, k_off: torch.Tensor | None, + pos: int, + align_mask: torch.Tensor | None + ) -> tuple[torch.Tensor, torch.Tensor, + torch.Tensor, torch.Tensor] | None: + """Q/K suffix 对齐后取 [pos]，返回 (q_on, q_off, k_on, k_off) 四向量。 + + 每个向量压平 [H*D]，可直接做 vec_metrics 对比。 + """ + if q_on is None or q_off is None: + return None + q_on_f = q_on.reshape(q_on.shape[0], -1) + q_off_f = q_off.reshape(q_off.shape[0], -1) + k_on_f = k_on.reshape(k_on.shape[0], -1) if k_on is not None else None + k_off_f = k_off.reshape(k_off.shape[0], -1) if k_off is not None else None + + if align_mask is not None and q_on.shape[0] != q_off.shape[0]: + try: + q_on_f, q_off_f = _align_packed(q_on_f, q_off_f, align_mask) + if k_on_f is not None: + k_on_f, k_off_f = _align_packed(k_on_f, k_off_f, align_mask) + except ValueError: + return None + n = min(q_on_f.shape[0], q_off_f.shape[0]) + if pos < 0 or pos >= n: + return None + qo = q_on_f[pos].contiguous() + qf = q_off_f[pos].contiguous() + ko = k_on_f[pos].contiguous() if k_on_f is not None else None + kf = k_off_f[pos].contiguous() if k_off_f is not None else None + return qo, qf, ko, kf + + +def cmp_rope_emb_token(dir_on: str, dir_off: str, + pos: int = 0, layer: int | None = None, + align_mask: torch.Tensor | None = None + ) -> list[CheckResult]: + """Q/K packed[pos] 对比（**suffix 对齐后**）。 + + 对 Q 和 K 分别输出 attn_L{lyr}_Q_pos{pos} / attn_L{lyr}_K_pos{pos}。 + """ + if align_mask is None: + align_mask = _build_attn_align_mask(dir_on, dir_off) + results: list[CheckResult] = [] + rope_layer = layer if layer is not None else ( + _get_num_layers(dir_on) or _get_num_layers(dir_off)) + + if rope_layer: + q_on, k_on = _load_rope_emb(dir_on, rope_layer) + q_off, k_off = _load_rope_emb(dir_off, rope_layer) + vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) + if vecs is None: + results.append(CheckResult( + name=f"rope_emb_L{rope_layer}_pos{pos}", + metrics={"error": f"无法对齐或 pos {pos} 越界"})) + else: + qo, qf, ko, kf = vecs + results.append(CheckResult( + name=f"rope_emb_L{rope_layer}_Q_pos{pos}", + metrics=_vec_metrics(qo, qf))) + if ko is not None and kf is not None: + results.append(CheckResult( + name=f"rope_emb_L{rope_layer}_K_pos{pos}", + metrics=_vec_metrics(ko, kf))) + return results + + def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: """全 packed logits suffix 对齐 + per-token cosine。""" lo = _load_logits(dir_on) @@ -784,6 +949,45 @@ def _print_per_layer(r: CheckResult): print() +def _print_rope_emb_per_layer(r: CheckResult): + print(_SEP_SINGLE + "\n [rope_emb] Post-RoPE Q/K Per-Layer Cosine Similarity") + print(_SEP_SINGLE) + layers = r.metrics.get("layers") + if isinstance(layers, dict): + print(f" {'LAYER':>6s} {'Q_COS_AVG':>14s} {'Q_COS_MIN':>14s} " + f"{'K_COS_AVG':>14s} {'K_COS_MIN':>14s} " + f"{'TOKENS':>8s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 14} {'─' * 14} {'─' * 14} {'─' * 14} " + f"{'─' * 8} {'─' * 8}") + bad = [] + for lyr in sorted(layers.keys()): + d = layers[lyr] + if "error" in d: + print(f" {lyr:>6d} {d['error']}") + bad.append(lyr) + continue + ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS + and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) + print(f" {lyr:>6d} {d['Q_cos_avg']:>14.6e} {d['Q_cos_min']:>14.6e} " + f"{d['K_cos_avg']:>14.6e} {d['K_cos_min']:>14.6e} " + f"{d['n_tokens']:>8d} {'PASS' if ok else 'WARN':>8s}") + if not ok: + bad.append(lyr) + if bad: + print(f"\n ⚠ First deviating layer: {bad[0]}") + elif "Q_cos_avg" in r.metrics: + d = r.metrics + ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS + and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) + print(f" L{d['layer']} Q_cos_avg={d['Q_cos_avg']:.6e} " + f"Q_cos_min={d['Q_cos_min']:.6e} " + f"K_cos_avg={d['K_cos_avg']:.6e} K_cos_min={d['K_cos_min']:.6e} " + f"{'PASS' if ok else 'WARN'}") + elif "error" in r.metrics: + print(f" {_CROSS} {r.metrics['error']}") + print() + + def _print_packed_token(r: CheckResult): print(_SEP_SINGLE + f"\n [packed_token] {r.name}") print(_SEP_SINGLE) @@ -986,6 +1190,12 @@ def main(): all_results.append(r) _print_per_layer(r) + # ── packed: post-RoPE Q/K per-layer cos ── + r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer) + if r: + all_results.append(r) + _print_rope_emb_per_layer(r) + # ── packed: packed_token（attn[pos] + logits[pos]，suffix 对齐后） ── # pos 由 --token 指定（默认 0，索引对齐后的 suffix-packed 空间）； # attn 用 --layer 指定的层（默认最后一层）；logits 永远最后一层。 @@ -1020,6 +1230,28 @@ def main(): all_results.append(r) _print_logits_packed(r) + # ── packed: post-RoPE Q/K packed_token ── + rope_pt_results = cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, + align_mask=align_mask) + for r in rope_pt_results: + all_results.append(r) + _print_packed_token(r) + # rope_emb packed_token top-K + if args.topk > 0 and rope_pt_results: + rope_layer = args.layer if args.layer is not None else ( + _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) + if rope_layer: + q_on, k_on = _load_rope_emb(args.dir_on, rope_layer) + q_off, k_off = _load_rope_emb(args.dir_off, rope_layer) + vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) + if vecs is not None: + qo, qf, ko, kf = vecs + _print_topk_vec(qo.cpu(), qf.cpu(), args.topk, "val", + f"rope_emb_L{rope_layer}_Q_pos{pos}") + if ko is not None and kf is not None: + _print_topk_vec(ko.cpu(), kf.cpu(), args.topk, "val", + f"rope_emb_L{rope_layer}_K_pos{pos}") + # ── 2D: logprobs + entropy ── for fname, cname in [("logprobs", "logp"), ("entropy", "entropy")]: fn = f"{fname}_{args.tag}.pt" diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 19838145..020b97a0 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -203,3 +203,40 @@ def dump_entropy_2d_verl080(ent_2d: torch.Tensor | None, tag: str) -> None: if dump_dir is None or ent_2d is None: return _save_tensor(f"entropy_{tag}.pt", ent_2d, dump_dir) + + +# ════════════════════════════════════════════════════════════════ +# Post-RoPE Q/K dump (per layer) +# ════════════════════════════════════════════════════════════════ + +_ROPE_EMB_BUFFER: dict[int, dict] | None = None + + +def dump_rope_emb_verl080(layer_number: int, + rotated_query: torch.Tensor, + rotated_key: torch.Tensor, + num_layers: int) -> None: + """Accumulate one layer's post-RoPE Q/K. Auto-flush to ``rope_emb.pt`` on last layer. + + Format: ``{layer_idx: {"query": [T, H, D], "key": [T, H, D]}}`` + ON packed 只含 suffix（裁剪后），OFF packed 含完整序列。 + cmp 侧用 prefix_lens + cu_seqlens 做 suffix 对齐后对比（同 attn_output 模式）。 + """ + global _ROPE_EMB_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _ROPE_EMB_BUFFER is None: + _ROPE_EMB_BUFFER = {} + _ROPE_EMB_BUFFER[layer_number] = { + "query": rotated_query.detach().cpu().clone(), + "key": rotated_key.detach().cpu().clone(), + } + if layer_number == num_layers: + from prefix_sharing.tools.diagnostic_dump import _rank0_only + if _rank0_only(): + try: + torch.save(_ROPE_EMB_BUFFER, os.path.join(dump_dir, "rope_emb.pt")) + except Exception: + pass + _ROPE_EMB_BUFFER = None From f4be6b5a2de1bb5e634063f2782e2474dc00226c Mon Sep 17 00:00:00 2001 From: Boundless Date: Fri, 26 Jun 2026 19:27:17 +0800 Subject: [PATCH 17/61] [fix] dump_rope_emb_verl080: accept optional positions param Co-Authored-By: Claude Fable 5 --- .../prefix_sharing/tools/diagnostic_dump_verl080.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 020b97a0..79e02047 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -215,12 +215,14 @@ def dump_entropy_2d_verl080(ent_2d: torch.Tensor | None, tag: str) -> None: def dump_rope_emb_verl080(layer_number: int, rotated_query: torch.Tensor, rotated_key: torch.Tensor, - num_layers: int) -> None: + num_layers: int, + positions: torch.Tensor | None = None) -> None: """Accumulate one layer's post-RoPE Q/K. Auto-flush to ``rope_emb.pt`` on last layer. - Format: ``{layer_idx: {"query": [T, H, D], "key": [T, H, D]}}`` + Format: ``{layer_idx: {"query": [T, H, D], "key": [T, H, D], "positions": [T] or None}}`` ON packed 只含 suffix（裁剪后），OFF packed 含完整序列。 cmp 侧用 prefix_lens + cu_seqlens 做 suffix 对齐后对比（同 attn_output 模式）。 + positions 可选，用于手动排查时的位置回溯。 """ global _ROPE_EMB_BUFFER dump_dir = _get_dump_dir() @@ -228,10 +230,13 @@ def dump_rope_emb_verl080(layer_number: int, return if _ROPE_EMB_BUFFER is None: _ROPE_EMB_BUFFER = {} - _ROPE_EMB_BUFFER[layer_number] = { + entry = { "query": rotated_query.detach().cpu().clone(), "key": rotated_key.detach().cpu().clone(), } + if positions is not None: + entry["positions"] = positions.detach().cpu().clone() + _ROPE_EMB_BUFFER[layer_number] = entry if layer_number == num_layers: from prefix_sharing.tools.diagnostic_dump import _rank0_only if _rank0_only(): From badaba33bb913f2101117c958c7868bd61ec275a Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 10:59:07 +0800 Subject: [PATCH 18/61] [fix] attention OFF diag: local import torch (NameError on torch.arange/cat) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OFF 路径 diag dump 用 torch.arange/torch.cat 构造每 segment 的 position IDs， 但 attention.py 顶部未导入 torch（与 forward_step.py:170 惯例一致）， 导致 OFF 路径报 NameError: name 'torch' is not defined。 在 diag block 内局部 import torch，保持顶部零 torch 导入且非 diag 路径零开销。 --- .../setup/patches/verl080_mcore0161_ms0160/attention.py | 1 + 1 file changed, 1 insertion(+) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 3407abd0..c70bf37a 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -54,6 +54,7 @@ def patched_forward( # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump ##### import os as _os if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + import torch # OFF diag block 用 torch.arange/torch.cat（顶部未导入 torch，与 forward_step.py 惯例一致，按需局部导入） from prefix_sharing.tools.diagnostic_dump import ( dump_attn_off, dump_rope_freqs_off, ) From da396612450e5f429214477bbc866ddee1f9c24c Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 11:41:49 +0800 Subject: [PATCH 19/61] =?UTF-8?q?[fix]=20OFF=20post-RoPE=20Q/K:=20hook=20?= =?UTF-8?q?=E7=9C=9F=E5=AE=9E=E5=BC=A0=E9=87=8F=E6=9B=BF=E4=BB=A3=E9=87=8D?= =?UTF-8?q?=E7=AE=97=20+=20=E4=BF=AE=20rope=5Femb=20=E5=86=99=E7=9B=98=20b?= =?UTF-8?q?ug?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit OFF 路径 post-RoPE Q/K 之前用 get_query_key_value_tensors + 重建 positions + index_select + 重算 apply_rotary_pos_emb，既脆又有 bug（_k_pos_emb or _q_pos_emb 对 tensor 求 bool 崩溃）。post-RoPE Q/K 是 original_forward 内部中间变量，唯一 拿到真实张量的办法是 hook。 - diagnostic_dump_verl080: 新增 capture_post_rope_qk() context manager， monkey-patch megatron.core.transformer.attention.apply_rotary_pos_emb， 捕获 mcore THD prefill 每层 Q→K 两次调用的返回值，finally 还原。 ON 不受影响（apply_rotary_pos_emb 在 megatron_runtime 自有命名空间，不走 attention 模块全局）。 - attention.py OFF 分支：original_forward 外套 capture_post_rope_qk， 直接取 captures[0]/[1] dump，删掉重算逻辑。 - 修潜伏 bug：dump_rope_emb_verl080 原手写 os.path.join 但未 import os， NameError 被 except 静默吞掉导致 rope_emb.pt 从未写盘；改用 _save_tensor。 → ON/OFF 都需用本提交重跑一次。 --- .../verl080_mcore0161_ms0160/attention.py | 97 ++++++++++--------- .../tools/diagnostic_dump_verl080.py | 45 +++++++-- 2 files changed, 90 insertions(+), 52 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index c70bf37a..afeaa577 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -36,31 +36,41 @@ def patched_forward( ctx = current_prefix_sharing_context() if ctx is None: # ── normal path: 调用原始 forward ── - _result = original_forward( - self, - hidden_states, - attention_mask, - key_value_states=key_value_states, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - rotary_pos_cos_sin=rotary_pos_cos_sin, - attention_bias=attention_bias, - packed_seq_params=packed_seq_params, - sequence_len_offset=sequence_len_offset, - inference_params=inference_params, - ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump ##### + # diag: hook 截获 original_forward 内部 apply_rotary_pos_emb 的真实 post-RoPE + # Q/K（mcore Attention.forward THD prefill 每层调两次：先 Q 后 K）。post-RoPE Q/K + # 是 forward 内部中间变量，唯一能拿到真实张量的办法就是 hook 那个模块级 rotary 函数。 import os as _os - if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: - import torch # OFF diag block 用 torch.arange/torch.cat（顶部未导入 torch，与 forward_step.py 惯例一致，按需局部导入） + from contextlib import nullcontext + _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None + if _diag_on and rotary_pos_emb is not None: + from prefix_sharing.tools.diagnostic_dump_verl080 import capture_post_rope_qk + _rope_cm = capture_post_rope_qk() + else: + _rope_cm = nullcontext() + with _rope_cm as _rope_caps: + _result = original_forward( + self, + hidden_states, + attention_mask, + key_value_states=key_value_states, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + inference_params=inference_params, + ) + # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump ##### + if _diag_on: + import torch # 仅用于构造 debug 用的 positions（cmp 不依赖） from prefix_sharing.tools.diagnostic_dump import ( dump_attn_off, dump_rope_freqs_off, ) from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_emb_verl080 from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb - from megatron.core.models.common.embeddings.rope_utils import apply_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result _bs = ( len(packed_seq_params.cu_seqlens_q_padded) - 1 @@ -73,35 +83,30 @@ def patched_forward( if rotary_pos_emb is not None: _q_pos_emb, _k_pos_emb = _unpack_rotary_pos_emb(rotary_pos_emb) dump_rope_freqs_off(_q_pos_emb, self.layer_number, self.config.num_layers) - # OFF post-RoPE Q/K dump: 提取 QKV + 手动 apply RoPE 后 dump，供与 ON 对比 - # OFF 路径 position IDs 是标准的 0..seg_len-1（每 segment 内连续） - if (packed_seq_params is not None - and packed_seq_params.qkv_format == "thd" - and hasattr(packed_seq_params, "cu_seqlens_q_padded")): - _off_q, _off_k, _off_v = self.get_query_key_value_tensors( - hidden_states, key_value_states, - split_qkv=True, output_gate=False, - ) - _off_q = _off_q.squeeze(1) - _off_k = _off_k.squeeze(1) - _cu = packed_seq_params.cu_seqlens_q_padded - _pos_list = [torch.arange(_cu[i+1] - _cu[i], device=_off_q.device) - for i in range(len(_cu) - 1)] - _off_positions = torch.cat(_pos_list, dim=0).long() - _off_q_freqs = _q_pos_emb.index_select(0, _off_positions) - _off_k_freqs = (_k_pos_emb or _q_pos_emb).index_select(0, _off_positions) - _off_q_rope = apply_rotary_pos_emb( - _off_q.unsqueeze(1), _off_q_freqs, - config=self.config, cu_seqlens=None, - ).squeeze(1) - _off_k_rope = apply_rotary_pos_emb( - _off_k.unsqueeze(1), _off_k_freqs, - config=self.config, cu_seqlens=None, - ).squeeze(1) + # OFF post-RoPE Q/K：直接用 hook 截获的真实张量 + # （original_forward 内部 apply_rotary_pos_emb 的返回值）， + # captures[0]=Q, captures[1]=K。不再 get_query_key_value_tensors + 重算 RoPE。 + if _rope_caps is not None and len(_rope_caps) >= 2: + _off_q_rope, _off_k_rope = _rope_caps[0], _rope_caps[1] + # positions 仅作 debug 记录；cmp 按 cu_seqlens+prefix_lens 对齐，不用它 + _off_positions = None + if (packed_seq_params is not None + and hasattr(packed_seq_params, "cu_seqlens_q_padded")): + _cu = packed_seq_params.cu_seqlens_q_padded + _off_positions = torch.cat([ + torch.arange(_cu[i + 1] - _cu[i]) + for i in range(len(_cu) - 1) + ]).long() dump_rope_emb_verl080( self.layer_number, _off_q_rope, _off_k_rope, - self.config.num_layers, - positions=_off_positions, + self.config.num_layers, positions=_off_positions, + ) + else: + print( + f"[PS-diag] OFF rope_emb L{self.layer_number}: " + f"expected 2 captures (Q,K), got " + f"{len(_rope_caps) if _rope_caps is not None else 'None'}; skip", + flush=True, ) # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump end ##### return _result diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 79e02047..2c9d028c 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -22,6 +22,7 @@ from __future__ import annotations +import contextlib from typing import Any import torch @@ -238,10 +239,42 @@ def dump_rope_emb_verl080(layer_number: int, entry["positions"] = positions.detach().cpu().clone() _ROPE_EMB_BUFFER[layer_number] = entry if layer_number == num_layers: - from prefix_sharing.tools.diagnostic_dump import _rank0_only - if _rank0_only(): - try: - torch.save(_ROPE_EMB_BUFFER, os.path.join(dump_dir, "rope_emb.pt")) - except Exception: - pass + # 用 _save_tensor 写盘（修潜伏 bug：原手写 os.path.join 但本文件未 import os， + # NameError 被外层 except 静默吞掉 → rope_emb.pt 从未真正写盘）。 + _save_tensor("rope_emb.pt", _ROPE_EMB_BUFFER, dump_dir) _ROPE_EMB_BUFFER = None + + +@contextlib.contextmanager +def capture_post_rope_qk(): + """Hook mcore 的 ``apply_rotary_pos_emb``，截获 original_forward 内部真实算出的 post-RoPE Q/K。 + + mcore ``Attention.forward`` 的 THD prefill 路径每层恰好调用模块级 + ``apply_rotary_pos_emb`` 两次——先 Q 后 K（见 megatron/core/transformer/ + attention.py:1097,1110）。这里 monkey-patch 该模块全局函数，把每次调用的 + 返回值（旋转后的张量）追加到 yield 的 list；``finally`` 还原原函数。 + + ON 路径不受影响：它在 megatron_runtime.py 里把 ``apply_rotary_pos_emb`` + import 进了自有命名空间，不经 attention 模块全局解析，patch 不到它。 + + 用法（OFF 分支，包住 original_forward 调用）:: + + with capture_post_rope_qk() as caps: + result = original_forward(...) + # caps[0]=rotated_query, caps[1]=rotated_key（单层 Attention.forward 内的调用顺序） + """ + import megatron.core.transformer.attention as _attn_mod + + _orig = _attn_mod.apply_rotary_pos_emb + captures: list = [] + + def _capturing(t, *args, **kwargs): + out = _orig(t, *args, **kwargs) + captures.append(out) + return out + + _attn_mod.apply_rotary_pos_emb = _capturing + try: + yield captures + finally: + _attn_mod.apply_rotary_pos_emb = _orig From eb9ebc32b03e51e214854f0abaed3f1e41d6ec02 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 12:48:06 +0800 Subject: [PATCH 20/61] =?UTF-8?q?[diag]=20cmp=5Fverl080:=20rope=5Femb=20?= =?UTF-8?q?=E5=85=A5=20shapes=20=E8=A1=A8=20+=20=E7=BB=86=E5=8C=96=20pos?= =?UTF-8?q?=20=E5=A4=B1=E8=B4=A5=E8=AF=8A=E6=96=AD?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _print_shapes 加入 rope_emb.pt；_shape_of 支持嵌套 dict（每层 {query,key, positions}），显示 (dict,NL,Q(T,H,D))，一眼看出 ON(suffix-only)/OFF(full) token 数差异。 - cmp_rope_emb_token 的 pos 失败从笼统「无法对齐或越界」改为 _diag_rope_pos_fail 三分：缺失 / 对齐失败(带 n_on,n_off,align_mask len,sum) / pos 越界(带对齐后 token 数)，定位 pos16 失败根因。 --- .../prefix_sharing/tools/cmp_diag_verl080.py | 27 +++++++++++++++++-- 1 file changed, 25 insertions(+), 2 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 0e833087..25bba365 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -565,6 +565,21 @@ def _rope_emb_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, return qo, qf, ko, kf +def _diag_rope_pos_fail(q_on: torch.Tensor | None, q_off: torch.Tensor | None, + pos: int, align_mask: torch.Tensor | None) -> str: + """rope_emb packed_token 取 [pos] 失败时的诊断串：区分 缺失 / 对齐失败 / pos 越界。""" + if q_on is None or q_off is None: + return f"rope_emb 该层在 {'ON' if q_on is None else 'OFF'} 侧缺失" + n_on, n_off = q_on.shape[0], q_off.shape[0] + if align_mask is not None and n_on != n_off: + msum = int(align_mask.sum()) + return (f"对齐失败: n_on={n_on} n_off={n_off} " + f"align_mask(len={align_mask.shape[0]}, sum={msum}); " + f"需 ON tokens==sum({msum}) 且 mask_len==n_off({n_off})") + post = min(n_on, n_off) + return f"pos {pos} 越界: 对齐后 token 数={post} (n_on={n_on}, n_off={n_off})" + + def cmp_rope_emb_token(dir_on: str, dir_off: str, pos: int = 0, layer: int | None = None, align_mask: torch.Tensor | None = None @@ -586,7 +601,7 @@ def cmp_rope_emb_token(dir_on: str, dir_off: str, if vecs is None: results.append(CheckResult( name=f"rope_emb_L{rope_layer}_pos{pos}", - metrics={"error": f"无法对齐或 pos {pos} 越界"})) + metrics={"error": _diag_rope_pos_fail(q_on, q_off, pos, align_mask)})) else: qo, qf, ko, kf = vecs results.append(CheckResult( @@ -799,7 +814,14 @@ def _shape_of(dir_path: str, filename: str) -> str: if isinstance(obj, dict): # per-layer dict（attn_outputs / rope_freqs_*）：显示层数 + 首层 shape sample = next(iter(obj.values())) if obj else None - sample_shape = f",{tuple(sample.shape)}" if sample is not None else "" + # rope_emb.pt：每层值是 {"query","key"[,"positions"]} dict，取 query 的 shape 代表 + if isinstance(sample, dict): + _q = sample.get("query") + sample_shape = f",Q{tuple(_q.shape)}" if _q is not None else "" + elif sample is not None: + sample_shape = f",{tuple(sample.shape)}" + else: + sample_shape = "" return f"(dict,{len(obj)}L{sample_shape})" return str(tuple(obj.shape)) except Exception: @@ -853,6 +875,7 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, f"attention_mask_{tag}.pt", "logits.pt", "attn_outputs.pt", + "rope_emb.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] From ea764668ed05ab082d7a36da5f4f0c7a2cc8ab36 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 13:04:41 +0800 Subject: [PATCH 21/61] =?UTF-8?q?[fix]=20rope=5Femb.pt=20=E5=86=99?= =?UTF-8?q?=E7=9B=98=EF=BC=9A=5Fsave=5Ftensor=20=E5=AF=B9=20dict=20?= =?UTF-8?q?=E5=A4=B1=E8=B4=A5=20=E2=86=92=20=E7=9B=B4=E6=8E=A5=20torch.sav?= =?UTF-8?q?e?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dump_rope_emb_verl080 的 flush 之前用 _save_tensor，但 _save_tensor 内部对入参 做 .detach().cpu().clone()，dict 没 .detach() → AttributeError 被其 except 静默吞掉， rope_emb.pt 永不写盘（ON/OFF 都没写，cmp 显示 missing）。 da396612 原本想修 os 未导入的 bug，换成 _save_tensor 反而引入这个 dict 回归。 现仿 _flush_attn_buffer（写 attn_outputs.pt 的方式）：rank0 + 直接 torch.save(dict) + 本地 import os，失败 print 不再静默。entries 插入时已 cpu clone，无需再搬。 注：用户目录里的 rope_emb_on.pt / rope_emb_off.pt 是旧版本残留，当前代码不产生 （grep 零命中），真正名字是 rope_emb.pt（dump 与 cmp 一致），需清掉残留重跑。 --- .../tools/diagnostic_dump_verl080.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 2c9d028c..e2a13ac7 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -239,9 +239,17 @@ def dump_rope_emb_verl080(layer_number: int, entry["positions"] = positions.detach().cpu().clone() _ROPE_EMB_BUFFER[layer_number] = entry if layer_number == num_layers: - # 用 _save_tensor 写盘（修潜伏 bug：原手写 os.path.join 但本文件未 import os， - # NameError 被外层 except 静默吞掉 → rope_emb.pt 从未真正写盘）。 - _save_tensor("rope_emb.pt", _ROPE_EMB_BUFFER, dump_dir) + # 不能用 _save_tensor：它对入参做 .detach().cpu().clone()，dict 没 .detach() → + # AttributeError 被其 except 吞掉，rope_emb.pt 永不写盘（这是 da396612 引入的回归）。 + # entries 在插入时已 detach().cpu().clone()，这里 rank0 直接 torch.save（仿 + # _flush_attn_buffer 写 attn_outputs.pt 的方式）。 + import os as _os + from prefix_sharing.tools.diagnostic_dump import _rank0_only + if _rank0_only(): + try: + torch.save(_ROPE_EMB_BUFFER, _os.path.join(dump_dir, "rope_emb.pt")) + except Exception as _e: + print(f"[PS-diag] rope_emb.pt save failed: {_e}", flush=True) _ROPE_EMB_BUFFER = None From a4f3495f0ff0f229b88652379c911102e8267f0f Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 14:35:34 +0800 Subject: [PATCH 22/61] =?UTF-8?q?[diag]=20RoPE=20=E5=85=A8=E6=B5=81?= =?UTF-8?q?=E6=B0=B4=E7=BA=BF=E5=AF=B9=E6=AF=94=EF=BC=9Apre=20Q/K=20?= =?UTF-8?q?=E2=86=92=20rope=5Ffreqs=20=E2=86=92=20post=20Q/K=EF=BC=88layer?= =?UTF-8?q?/token=20=E7=BA=A7=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit dump 侧（rope_preqk.pt，旋转前 Q/K）： - diagnostic_dump_verl080: 抽 _flush_dict_buffer 公用写盘 helper（避开 _save_tensor 对 dict 失败的 bug）；新增 dump_rope_preqk_verl080；capture_post_rope_qk → capture_rope_qk， 每次 apply_rotary_pos_emb 同时存 {"pre":输入, "post":返回}。 - attention.py(OFF): hook 解包 pre/post，分别 dump rope_preqk.pt / rope_emb.pt。 - megatron_runtime.py(ON): _apply_positioned_rope 入口 dump 旋转前 Q/K。 cmp 侧（cmp_diag_verl080）： - _ROPE_STAGES=[pre,post]；cmp_rope_emb_layer/token 改 stage 参数，支持单 stage 调用。 - rope_freqs 加 layer 过滤（cmp_rope_freqs(layer=...)）；新增 cmp_rope_freqs_token （指定 pos 的角度向量对比，max_abs 应为 0）。抽 _align_rope_freqs_layer helper。 - shapes 表加 rope_preqk.pt / rope_emb.pt；_shape_of 支持嵌套 dict。 - main 按 RoPE 计算顺序重排：per-layer 与 token 都是 pre Q/K → rope_freqs → post Q/K， 定位分歧出现在哪一步（pre 偏=上游；freqs 偏=角度表；post 才偏=旋转应用）。 - pos 失败诊断细化（_diag_rope_pos_fail 三分：缺失/对齐失败/越界）。 注：layer 1-indexed（megatron self.layer_number 从 1 起），第一层用 --layer 1。 --- .../integrations/megatron_runtime.py | 9 + .../verl080_mcore0161_ms0160/attention.py | 19 +- .../prefix_sharing/tools/cmp_diag_verl080.py | 294 ++++++++++++------ .../tools/diagnostic_dump_verl080.py | 74 +++-- 4 files changed, 273 insertions(+), 123 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index d2c523cd..6d657325 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -161,6 +161,15 @@ def _apply_positioned_rope( positions = packed_position_ids.to(device=query.device, dtype=torch.long) max_needed = positions.max().item() + 1 + ######### prefix-sharing diag: pre-RoPE Q/K dump（旋转前，尚未加位置编码）######### + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_preqk_verl080 + dump_rope_preqk_verl080(attention_module.layer_number, query, key, + attention_module.config.num_layers) + except Exception as _e: + print(f"rope_preqk (pre-RoPE) dump failed: {_e}", flush=True) + ######### prefix-sharing diag: pre-RoPE Q/K dump end ######### + # 当 packed_position_ids 所需要的最大 position id 超过了 q_pos_emb / k_pos_emb 的当前长度时， # 就需要对 q_pos_emb / k_pos_emb 进行扩展。 # THD 模式下生成的 pos_emb 仅覆盖 positions 0 .. max_seqlen_q-1 这段范围， diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index afeaa577..9017238b 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -43,8 +43,8 @@ def patched_forward( from contextlib import nullcontext _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None if _diag_on and rotary_pos_emb is not None: - from prefix_sharing.tools.diagnostic_dump_verl080 import capture_post_rope_qk - _rope_cm = capture_post_rope_qk() + from prefix_sharing.tools.diagnostic_dump_verl080 import capture_rope_qk + _rope_cm = capture_rope_qk() else: _rope_cm = nullcontext() with _rope_cm as _rope_caps: @@ -69,7 +69,9 @@ def patched_forward( from prefix_sharing.tools.diagnostic_dump import ( dump_attn_off, dump_rope_freqs_off, ) - from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_emb_verl080 + from prefix_sharing.tools.diagnostic_dump_verl080 import ( + dump_rope_emb_verl080, dump_rope_preqk_verl080, + ) from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result _bs = ( @@ -87,7 +89,8 @@ def patched_forward( # （original_forward 内部 apply_rotary_pos_emb 的返回值）， # captures[0]=Q, captures[1]=K。不再 get_query_key_value_tensors + 重算 RoPE。 if _rope_caps is not None and len(_rope_caps) >= 2: - _off_q_rope, _off_k_rope = _rope_caps[0], _rope_caps[1] + # 每个捕获是 {"pre": 旋转前, "post": 旋转后} + _q_cap, _k_cap = _rope_caps[0], _rope_caps[1] # positions 仅作 debug 记录；cmp 按 cu_seqlens+prefix_lens 对齐，不用它 _off_positions = None if (packed_seq_params is not None @@ -97,10 +100,16 @@ def patched_forward( torch.arange(_cu[i + 1] - _cu[i]) for i in range(len(_cu) - 1) ]).long() + # post-RoPE（旋转后） dump_rope_emb_verl080( - self.layer_number, _off_q_rope, _off_k_rope, + self.layer_number, _q_cap["post"], _k_cap["post"], self.config.num_layers, positions=_off_positions, ) + # pre-RoPE（旋转前，纯 QKV 投影） + dump_rope_preqk_verl080( + self.layer_number, _q_cap["pre"], _k_cap["pre"], + self.config.num_layers, + ) else: print( f"[PS-diag] OFF rope_emb L{self.layer_number}: " diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 25bba365..32b0d30e 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -438,13 +438,21 @@ def cmp_packed_token(dir_on: str, dir_off: str, # Post-RoPE Q/K compare: per-layer + packed_token # ══════════════════════════════════════════════════════════════════ -def _load_rope_emb(dir_path: str, layer: int +# RoPE 对比阶段：**先 pre（旋转前，rope_preqk.pt）后 post（旋转后，rope_emb.pt）**。 +# (stage, fname, label) — label 用作结果名前缀与打印 section 头。 +_ROPE_STAGES: list[tuple[str, str, str]] = [ + ("pre", "rope_preqk.pt", "rope_preqk"), + ("post", "rope_emb.pt", "rope_emb"), +] + + +def _load_rope_emb(dir_path: str, layer: int, fname: str = "rope_emb.pt" ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """Load post-RoPE Q/K for a single layer from rope_emb.pt. + """Load Q/K for a single layer from ``fname`` (rope_emb.pt=post, rope_preqk.pt=pre). Returns ``(query, key)`` or ``(None, None)``. """ - fp = os.path.join(dir_path, "rope_emb.pt") + fp = os.path.join(dir_path, fname) if not os.path.exists(fp): return None, None d = torch.load(fp, weights_only=True) @@ -479,17 +487,14 @@ def _rope_emb_cos_for_layer(qa: torch.Tensor, ka: torch.Tensor, } -def cmp_rope_emb_layer(dir_on: str, dir_off: str, - layer: int | None) -> CheckResult | None: - """Post-RoPE Q/K per-layer cosine（suffix 对齐）。 - - 单层模式（layer 给定）：返回该层 Q/K cos。全层模式：返回所有层汇总。 - """ +def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, + fname: str, label: str) -> CheckResult | None: + """单 stage（fname/label）的 Q/K per-layer cosine（suffix 对齐）。""" align_mask = _build_attn_align_mask(dir_on, dir_off) if layer is not None: - q_on, k_on = _load_rope_emb(dir_on, layer) - q_off, k_off = _load_rope_emb(dir_off, layer) + q_on, k_on = _load_rope_emb(dir_on, layer, fname) + q_off, k_off = _load_rope_emb(dir_off, layer, fname) if q_on is None or q_off is None: return None need = align_mask is not None and q_on.shape[0] != q_off.shape[0] @@ -497,16 +502,16 @@ def cmp_rope_emb_layer(dir_on: str, dir_off: str, d = _rope_emb_cos_for_layer(q_on, k_on, q_off, k_off, align_mask if need else None) except ValueError as e: - return CheckResult(name=f"rope_emb_L{layer}", passed=False, + return CheckResult(name=f"{label}_L{layer}", passed=False, metrics={"error": str(e)}) d["layer"] = layer ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) - return CheckResult(name=f"rope_emb_L{layer}", passed=ok, metrics=d) + return CheckResult(name=f"{label}_L{layer}", passed=ok, metrics=d) # All layers - fa = os.path.join(dir_on, "rope_emb.pt") - fb = os.path.join(dir_off, "rope_emb.pt") + fa = os.path.join(dir_on, fname) + fb = os.path.join(dir_off, fname) if not os.path.exists(fa) or not os.path.exists(fb): return None da = torch.load(fa, weights_only=True) @@ -527,10 +532,22 @@ def cmp_rope_emb_layer(dir_on: str, dir_off: str, align_mask if need else None) except ValueError as e: results[lyr] = {"error": str(e)} - return CheckResult(name="rope_emb_per_layer", passed=True, + return CheckResult(name=f"{label}_per_layer", passed=True, metrics={"layers": results}) +def cmp_rope_emb_layer(dir_on: str, dir_off: str, layer: int | None, + stage: str = "post") -> CheckResult | None: + """Q/K per-layer cosine（suffix 对齐），单 stage。 + + stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_emb.pt（旋转后）。 + 调用方按 pre → rope_freqs → post 顺序分别调用，便于定位分歧出现在 RoPE 哪一步。 + """ + if stage == "pre": + return _cmp_rope_stage_layer(dir_on, dir_off, layer, "rope_preqk.pt", "rope_preqk") + return _cmp_rope_stage_layer(dir_on, dir_off, layer, "rope_emb.pt", "rope_emb") + + def _rope_emb_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, q_off: torch.Tensor | None, k_off: torch.Tensor | None, pos: int, @@ -580,40 +597,52 @@ def _diag_rope_pos_fail(q_on: torch.Tensor | None, q_off: torch.Tensor | None, return f"pos {pos} 越界: 对齐后 token 数={post} (n_on={n_on}, n_off={n_off})" -def cmp_rope_emb_token(dir_on: str, dir_off: str, - pos: int = 0, layer: int | None = None, - align_mask: torch.Tensor | None = None - ) -> list[CheckResult]: - """Q/K packed[pos] 对比（**suffix 对齐后**）。 - - 对 Q 和 K 分别输出 attn_L{lyr}_Q_pos{pos} / attn_L{lyr}_K_pos{pos}。 - """ +def _cmp_rope_stage_token(dir_on: str, dir_off: str, pos: int, layer: int | None, + align_mask: torch.Tensor | None, fname: str, + label: str) -> list[CheckResult]: + """单 stage（fname/label）的 Q/K packed[pos]（suffix 对齐后）。""" if align_mask is None: align_mask = _build_attn_align_mask(dir_on, dir_off) results: list[CheckResult] = [] rope_layer = layer if layer is not None else ( _get_num_layers(dir_on) or _get_num_layers(dir_off)) - if rope_layer: - q_on, k_on = _load_rope_emb(dir_on, rope_layer) - q_off, k_off = _load_rope_emb(dir_off, rope_layer) + q_on, k_on = _load_rope_emb(dir_on, rope_layer, fname) + q_off, k_off = _load_rope_emb(dir_off, rope_layer, fname) vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) if vecs is None: results.append(CheckResult( - name=f"rope_emb_L{rope_layer}_pos{pos}", + name=f"{label}_L{rope_layer}_pos{pos}", metrics={"error": _diag_rope_pos_fail(q_on, q_off, pos, align_mask)})) else: qo, qf, ko, kf = vecs results.append(CheckResult( - name=f"rope_emb_L{rope_layer}_Q_pos{pos}", + name=f"{label}_L{rope_layer}_Q_pos{pos}", metrics=_vec_metrics(qo, qf))) if ko is not None and kf is not None: results.append(CheckResult( - name=f"rope_emb_L{rope_layer}_K_pos{pos}", + name=f"{label}_L{rope_layer}_K_pos{pos}", metrics=_vec_metrics(ko, kf))) return results +def cmp_rope_emb_token(dir_on: str, dir_off: str, + pos: int = 0, layer: int | None = None, + align_mask: torch.Tensor | None = None, + stage: str = "post") -> list[CheckResult]: + """Q/K packed[pos] 对比（**suffix 对齐后**），单 stage。 + + stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_emb.pt（旋转后）。 + 对 Q、K 分别输出 {label}_L{lyr}_Q_pos{pos} / {label}_L{lyr}_K_pos{pos}。 + 调用方按 pre → rope_freqs → post 顺序分别调用。 + """ + if stage == "pre": + return _cmp_rope_stage_token(dir_on, dir_off, pos, layer, align_mask, + "rope_preqk.pt", "rope_preqk") + return _cmp_rope_stage_token(dir_on, dir_off, pos, layer, align_mask, + "rope_emb.pt", "rope_emb") + + def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: """全 packed logits suffix 对齐 + per-token cosine。""" lo = _load_logits(dir_on) @@ -646,17 +675,34 @@ def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: "cos_avg": cos_avg, "cos_min": cos_min}) -def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: - """对比 pre-RoPE 角度表（angle table，非 cos/sin）— suffix 对齐。 +def _align_rope_freqs_layer(on_dict: dict, off_dict: dict, layer: int, + seqlens: list[int], + align_mask: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor] | None: + """单层 rope_freqs：OFF raw 表按 seqlens 重建 per-token + suffix 对齐。 - ON ``rope_freqs_on.pt``: per-token 角度 dict {layer: [T_on, 1, 1, D]} - （已 index_select 到 packed_position_ids，每 token 实际旋转角度） - OFF ``rope_freqs_off.pt``: raw 角度表 dict {layer: [L0, 1, 1, D]} - （freqs[p] = p * inv_freq，未切片） + 返回 (on_aligned, off_aligned)，shape [N, 1, 1, D]；layer 缺失或对齐失败返回 None。 + 供 cmp_rope_freqs（per-layer max_diff）与 cmp_rope_freqs_token（[pos] 角度向量）复用。 + """ + if layer not in on_dict or layer not in off_dict: + return None + on_freqs = on_dict[layer] # [T_on,1,1,D] + off_freqs = torch.cat([off_dict[layer][:s, :, :, :] for s in seqlens], dim=0) # [T_off,1,1,D] + try: + return _align_packed(on_freqs, off_freqs, align_mask) + except ValueError: + return None - OFF per-token 角度从 raw 表按 cu_seqlens_off 重建（每段取 ``[:seg_len]``）， - 再与 ON 用同一 suffix 对齐（cu_seqlens + prefix_lens）后逐元素比 max_diff。 - 角度是 RoPE 的输入，应精确相等（``max_diff == 0``）。 + +def cmp_rope_freqs(dir_on: str, dir_off: str, + layer: int | None = None) -> CheckResult | None: + """对比 pre-RoPE 角度表（angle table，非 cos/sin）— suffix 对齐，应精确相等 max_diff==0。 + + ON ``rope_freqs_on.pt``: per-token 角度 dict {layer: [T_on,1,1,D]}（已 index_select 到 + packed_position_ids，每 token 实际旋转角度） + OFF ``rope_freqs_off.pt``: raw 角度表 dict {layer: [L0,1,1,D]}（freqs[p]=p*inv_freq，未切片） + OFF per-token 从 raw 表按 cu_seqlens_off 重建（每段 [:seg_len]），再 suffix 对齐。 + ``layer`` 给定则只比该层。角度是 RoPE 输入，应精确相等（max_diff==0）。 """ fa = os.path.join(dir_on, "rope_freqs_on.pt") fb = os.path.join(dir_off, "rope_freqs_off.pt") @@ -667,65 +713,99 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None - la, lb = set(on_dict.keys()), set(off_dict.keys()) - if la != lb: - return CheckResult(name="rope_freqs", passed=False, - metrics={"error": "layer set mismatch", - "on_layers": sorted(la), - "off_layers": sorted(lb)}) + layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) + if layer is not None: + layers = [l for l in layers if l == layer] + _name = f"rope_freqs_L{layer}" if layer is not None else "rope_freqs" + if not layers: + return CheckResult(name=_name, passed=False, + metrics={"error": f"layer {layer} 不在双方 rope_freqs 中"}) mb = _load_packed_meta(dir_off) if mb is None: - return CheckResult(name="rope_freqs", passed=False, - metrics={"error": "OFF cu_seqlens missing"}) - cu_off = mb["cu_seqlens"] - T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 - + return CheckResult(name=_name, passed=False, metrics={"error": "OFF cu_seqlens missing"}) ma = _load_packed_meta(dir_on) if ma is None: - return CheckResult(name="rope_freqs", passed=False, - metrics={"error": "ON prefix_lens missing"}) + return CheckResult(name=_name, passed=False, metrics={"error": "ON prefix_lens missing"}) + cu_off = mb["cu_seqlens"] + T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 align_mask = _build_alignment_mask(cu_off, ma["prefix_lens"], T_off) - seqlens = (cu_off[1:] - cu_off[:-1]).tolist() max_diff = 0.0 mismatches: list[dict] = [] - for lyr in sorted(la): - on_freqs = on_dict[lyr] # [T_on, 1, 1, D] - # 从 raw 表重建 OFF per-token：每段 [:seg_len] - off_freqs = torch.cat( - [off_dict[lyr][:s, :, :, :] for s in seqlens], dim=0) # [T_off, 1, 1, D] - - try: - on_aligned, off_aligned = _align_packed( - on_freqs, off_freqs, align_mask) - except ValueError as e: - return CheckResult(name="rope_freqs", passed=False, - metrics={"error": f"align failed L{lyr}: {e}"}) - - diff = (on_aligned - off_aligned).abs() # [N, 1, 1, D] + for lyr in layers: + _aligned = _align_rope_freqs_layer(on_dict, off_dict, lyr, seqlens, align_mask) + if _aligned is None: + continue + on_a, off_a = _aligned + diff = (on_a - off_a).abs() # [N,1,1,D] md = float(diff.max()) max_diff = max(max_diff, md) - if md > 0: - token_diff = diff.squeeze(1).squeeze(1).max(dim=-1) # values [N], indices [N] + token_diff = diff.squeeze(1).squeeze(1).max(dim=-1) # values [N], indices [N] bad_mask = token_diff.values > 0 for t in bad_mask.nonzero(as_tuple=True)[0].tolist(): t = int(t) d = int(token_diff.indices[t]) mismatches.append({ "layer": lyr, "token_idx": t, "dim": d, - "on_val": float(on_aligned[t, 0, 0, d]), - "off_val": float(off_aligned[t, 0, 0, d]), + "on_val": float(on_a[t, 0, 0, d]), + "off_val": float(off_a[t, 0, 0, d]), "diff": float(token_diff.values[t]), }) - metrics: dict = {"max_diff": max_diff, "num_layers": len(la)} + metrics: dict = {"max_diff": max_diff, "num_layers": len(layers)} if mismatches: metrics["mismatches"] = mismatches[:20] metrics["total_mismatches"] = len(mismatches) - return CheckResult(name="rope_freqs", passed=max_diff == 0.0, metrics=metrics) + return CheckResult(name=_name, passed=max_diff == 0.0, metrics=metrics) + + +def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, + layer: int | None = None, + align_mask: torch.Tensor | None = None) -> CheckResult | None: + """rope_freqs 在对齐后 suffix-packed 位置 [pos] 的角度向量对比（应精确相等）。 + + 取 ``layer``（默认最后一层）对齐后第 ``pos`` 个 token 的角度向量 [D]，比 ON/OFF。 + 角度是 RoPE 输入，应逐元素相等 → max_abs 应为 0。 + """ + fa = os.path.join(dir_on, "rope_freqs_on.pt") + fb = os.path.join(dir_off, "rope_freqs_off.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None + on_dict = torch.load(fa, weights_only=True) + off_dict = torch.load(fb, weights_only=True) + if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): + return None + common = set(on_dict.keys()) & set(off_dict.keys()) + rf_layer = layer if layer is not None else (max(common) if common else 0) + _name = f"rope_freqs_L{rf_layer}_pos{pos}" + if rf_layer not in on_dict or rf_layer not in off_dict: + return CheckResult(name=_name, metrics={"error": f"layer {rf_layer} 缺失"}) + + mb = _load_packed_meta(dir_off) + ma = _load_packed_meta(dir_on) + if mb is None or ma is None: + return CheckResult(name=_name, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) + cu_off = mb["cu_seqlens"] + T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 + if align_mask is None: + align_mask = _build_alignment_mask(cu_off, ma["prefix_lens"], T_off) + seqlens = (cu_off[1:] - cu_off[:-1]).tolist() + + _aligned = _align_rope_freqs_layer(on_dict, off_dict, rf_layer, seqlens, align_mask) + if _aligned is None: + return CheckResult(name=_name, metrics={"error": "对齐失败"}) + on_a, off_a = _aligned + n = on_a.shape[0] + if pos < 0 or pos >= n: + return CheckResult(name=_name, + metrics={"error": f"pos {pos} 越界: 对齐后 token 数={n}"}) + on_vec = on_a[pos].reshape(-1) + off_vec = off_a[pos].reshape(-1) + m = _vec_metrics(on_vec, off_vec) + return CheckResult(name=_name, passed=m["max_abs"] == 0.0, metrics=m) # ════════════════════════════════════════════════════════════════ @@ -876,6 +956,7 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, "logits.pt", "attn_outputs.pt", "rope_emb.pt", + "rope_preqk.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] @@ -973,7 +1054,9 @@ def _print_per_layer(r: CheckResult): def _print_rope_emb_per_layer(r: CheckResult): - print(_SEP_SINGLE + "\n [rope_emb] Post-RoPE Q/K Per-Layer Cosine Similarity") + _sec = "rope_preqk" if "preqk" in r.name else "rope_emb" + _stage = "Pre-RoPE" if "preqk" in r.name else "Post-RoPE" + print(_SEP_SINGLE + f"\n [{_sec}] {_stage} Q/K Per-Layer Cosine Similarity") print(_SEP_SINGLE) layers = r.metrics.get("layers") if isinstance(layers, dict): @@ -1201,24 +1284,30 @@ def main(): all_results: list[CheckResult] = [] - # ── packed: rope 角度（suffix 对齐，应精确相等 max_diff==0） ── - r = cmp_rope_freqs(args.dir_on, args.dir_off) + # ── RoPE pipeline per-layer：pre Q/K → rope 角度 → post Q/K ── + # 按计算顺序串联：旋转前 Q/K → 每token旋转角度(freqs) → 旋转后 Q/K， + # 定位分歧出现在 RoPE 哪一步（pre 就偏=上游；freqs 偏=角度表；post 才偏=旋转应用）。 + r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer, stage="pre") if r: all_results.append(r) - _print_rope_freqs(r) + _print_rope_emb_per_layer(r) - # ── packed: attention_output per-layer cos ── - r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) + r = cmp_rope_freqs(args.dir_on, args.dir_off, layer=args.layer) if r: all_results.append(r) - _print_per_layer(r) + _print_rope_freqs(r) - # ── packed: post-RoPE Q/K per-layer cos ── - r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer) + r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer, stage="post") if r: all_results.append(r) _print_rope_emb_per_layer(r) + # ── packed: attention_output per-layer cos（RoPE 下游）── + r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) + if r: + all_results.append(r) + _print_per_layer(r) + # ── packed: packed_token（attn[pos] + logits[pos]，suffix 对齐后） ── # pos 由 --token 指定（默认 0，索引对齐后的 suffix-packed 空间）； # attn 用 --layer 指定的层（默认最后一层）；logits 永远最后一层。 @@ -1253,27 +1342,34 @@ def main(): all_results.append(r) _print_logits_packed(r) - # ── packed: post-RoPE Q/K packed_token ── - rope_pt_results = cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, - align_mask=align_mask) - for r in rope_pt_results: - all_results.append(r) - _print_packed_token(r) - # rope_emb packed_token top-K + # ── RoPE pipeline packed_token：pre Q/K → rope_freqs → post Q/K（指定 pos）── + rope_pt_results: list[CheckResult] = [] + for r in cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, + align_mask=align_mask, stage="pre"): + all_results.append(r); rope_pt_results.append(r); _print_packed_token(r) + _rf = cmp_rope_freqs_token(args.dir_on, args.dir_off, pos, args.layer, + align_mask=align_mask) + if _rf is not None: + all_results.append(_rf); rope_pt_results.append(_rf); _print_packed_token(_rf) + for r in cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, + align_mask=align_mask, stage="post"): + all_results.append(r); rope_pt_results.append(r); _print_packed_token(r) + # rope packed_token top-K（pre Q/K + post Q/K；freqs 角度应精确相等，vec_metrics 已含 max_abs） if args.topk > 0 and rope_pt_results: rope_layer = args.layer if args.layer is not None else ( _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) if rope_layer: - q_on, k_on = _load_rope_emb(args.dir_on, rope_layer) - q_off, k_off = _load_rope_emb(args.dir_off, rope_layer) - vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) - if vecs is not None: - qo, qf, ko, kf = vecs - _print_topk_vec(qo.cpu(), qf.cpu(), args.topk, "val", - f"rope_emb_L{rope_layer}_Q_pos{pos}") - if ko is not None and kf is not None: - _print_topk_vec(ko.cpu(), kf.cpu(), args.topk, "val", - f"rope_emb_L{rope_layer}_K_pos{pos}") + for _stage, _fname, _label in _ROPE_STAGES: + q_on, k_on = _load_rope_emb(args.dir_on, rope_layer, _fname) + q_off, k_off = _load_rope_emb(args.dir_off, rope_layer, _fname) + vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) + if vecs is not None: + qo, qf, ko, kf = vecs + _print_topk_vec(qo.cpu(), qf.cpu(), args.topk, "val", + f"{_label}_L{rope_layer}_Q_pos{pos}") + if ko is not None and kf is not None: + _print_topk_vec(ko.cpu(), kf.cpu(), args.topk, "val", + f"{_label}_L{rope_layer}_K_pos{pos}") # ── 2D: logprobs + entropy ── for fname, cname in [("logprobs", "logp"), ("entropy", "entropy")]: diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index e2a13ac7..65e5151b 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -239,46 +239,82 @@ def dump_rope_emb_verl080(layer_number: int, entry["positions"] = positions.detach().cpu().clone() _ROPE_EMB_BUFFER[layer_number] = entry if layer_number == num_layers: - # 不能用 _save_tensor：它对入参做 .detach().cpu().clone()，dict 没 .detach() → - # AttributeError 被其 except 吞掉，rope_emb.pt 永不写盘（这是 da396612 引入的回归）。 - # entries 在插入时已 detach().cpu().clone()，这里 rank0 直接 torch.save（仿 - # _flush_attn_buffer 写 attn_outputs.pt 的方式）。 - import os as _os - from prefix_sharing.tools.diagnostic_dump import _rank0_only - if _rank0_only(): - try: - torch.save(_ROPE_EMB_BUFFER, _os.path.join(dump_dir, "rope_emb.pt")) - except Exception as _e: - print(f"[PS-diag] rope_emb.pt save failed: {_e}", flush=True) + _flush_dict_buffer("rope_emb.pt", _ROPE_EMB_BUFFER, dump_dir) _ROPE_EMB_BUFFER = None +def _flush_dict_buffer(fname: str, buffer: dict, dump_dir: str) -> None: + """rank0 直接 torch.save 一个 dict buffer。 + + 不能用 _save_tensor：它对入参做 .detach().cpu().clone()，dict 没 .detach() → + AttributeError 被其 except 吞掉，文件永不写盘（rope_emb.pt 曾因此丢失）。 + entries 应在插入时已 detach().cpu().clone()。仿 _flush_attn_buffer。 + """ + import os as _os + from prefix_sharing.tools.diagnostic_dump import _rank0_only + if _rank0_only(): + try: + torch.save(buffer, _os.path.join(dump_dir, fname)) + except Exception as _e: + print(f"[PS-diag] {fname} save failed: {_e}", flush=True) + + +_ROPE_PREQK_BUFFER: dict[int, dict] | None = None + + +def dump_rope_preqk_verl080(layer_number: int, + query: torch.Tensor, + key: torch.Tensor, + num_layers: int) -> None: + """Accumulate one layer's pre-RoPE Q/K. Auto-flush to ``rope_preqk.pt`` on last layer. + + Format: ``{layer_idx: {"query": [T, H, D], "key": [T, H, D]}}``。 + 旋转前的 Q/K（纯 QKV 投影输出，未加位置编码）。ON/OFF 应逐元素相同—— + 同 hidden_states、同 QKV 权重。用来隔离：pre-RoPE 相同但 post-RoPE 不同 → + 问题在 RoPE；pre 就不同 → 问题在上游（hidden_states / 投影）。 + """ + global _ROPE_PREQK_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _ROPE_PREQK_BUFFER is None: + _ROPE_PREQK_BUFFER = {} + _ROPE_PREQK_BUFFER[layer_number] = { + "query": query.detach().cpu().clone(), + "key": key.detach().cpu().clone(), + } + if layer_number == num_layers: + _flush_dict_buffer("rope_preqk.pt", _ROPE_PREQK_BUFFER, dump_dir) + _ROPE_PREQK_BUFFER = None + + @contextlib.contextmanager -def capture_post_rope_qk(): - """Hook mcore 的 ``apply_rotary_pos_emb``，截获 original_forward 内部真实算出的 post-RoPE Q/K。 +def capture_rope_qk(): + """Hook mcore 的 ``apply_rotary_pos_emb``，同时截获旋转前(pre)与旋转后(post)的 Q/K。 mcore ``Attention.forward`` 的 THD prefill 路径每层恰好调用模块级 ``apply_rotary_pos_emb`` 两次——先 Q 后 K（见 megatron/core/transformer/ - attention.py:1097,1110）。这里 monkey-patch 该模块全局函数，把每次调用的 - 返回值（旋转后的张量）追加到 yield 的 list；``finally`` 还原原函数。 + attention.py:1097,1110）。monkey-patch 该模块全局函数，把每次调用的输入 t + (pre-RoPE) 与返回值 out (post-RoPE) 存为 ``{"pre": t, "post": out}`` 追加到 + yield 的 list；``finally`` 还原原函数。 ON 路径不受影响：它在 megatron_runtime.py 里把 ``apply_rotary_pos_emb`` import 进了自有命名空间，不经 attention 模块全局解析，patch 不到它。 用法（OFF 分支，包住 original_forward 调用）:: - with capture_post_rope_qk() as caps: + with capture_rope_qk() as caps: result = original_forward(...) - # caps[0]=rotated_query, caps[1]=rotated_key（单层 Attention.forward 内的调用顺序） + # caps[0]={"pre":Q_pre,"post":Q_post}, caps[1]={"pre":K_pre,"post":K_post} """ import megatron.core.transformer.attention as _attn_mod _orig = _attn_mod.apply_rotary_pos_emb - captures: list = [] + captures: list = [] # 每元素 {"pre": 旋转前 t, "post": 旋转后 out} def _capturing(t, *args, **kwargs): out = _orig(t, *args, **kwargs) - captures.append(out) + captures.append({"pre": t, "post": out}) return out _attn_mod.apply_rotary_pos_emb = _capturing From fad4c9d1191d1b84ecd510873f4e4f1bf34d5314 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 15:01:47 +0800 Subject: [PATCH 23/61] =?UTF-8?q?[refactor]=20rope=20=E5=91=BD=E5=90=8D?= =?UTF-8?q?=E7=BB=9F=E4=B8=80=EF=BC=9Arope=5Femb=E2=86=92rope=5Fpostqk?= =?UTF-8?q?=EF=BC=8Crope=5Ffreqs=20ON/OFF=20=E5=90=88=E5=B9=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 1) rope_emb → rope_postqk 改名（全仓 6 文件） rope_emb.pt 这名字误导（存的是旋转后 Q/K，非 embedding 表）。统一为 rope_postqk.pt， 与 rope_preqk.pt（旋转前）配成一对。函数/变量/buffer/label 一并改： dump_rope_emb_verl080→dump_rope_postqk_verl080, cmp_rope_emb_*→cmp_rope_postqk_*, _ROPE_EMB_BUFFER→_ROPE_POSTQK_BUFFER, 结果名/section 头同步。 判断 pre/post 用 "preqk" in name（postqk 不含 preqk，仍正确分流）。 2) rope_freqs: ON/OFF 统一 per-token，合并 dump 函数 原 rope_freqs_on.pt(per-token) / rope_freqs_off.pt(raw 表) 内容不对称、文件名冗余 （目录已区分 ON/OFF）。统一：两边都存 per-token 角度 rope_freqs.pt [T,1,1,D]。 - dump_rope_freqs_on/off 合并为 dump_rope_freqs（ON/OFF 共用）。 - OFF 调用方 index_select raw 表成 per-token 再 dump。 - cmp 去掉「从 raw 表按 cu_seqlens 重建」逻辑，_align_rope_freqs_layer 简化为直接对齐。 - shapes 表加 rope_freqs.pt；docstring 更新。 注：v070 cmp_diag.py 的 rope_freqs_on/off 引用保留（v070 dump 不写 rope_freqs，dead code）。 --- .../integrations/megatron_runtime.py | 12 +- .../verl080_mcore0161_ms0160/attention.py | 39 ++--- .../prefix_sharing/tools/cmp_diag.py | 14 +- .../prefix_sharing/tools/cmp_diag_verl080.py | 133 ++++++++---------- .../prefix_sharing/tools/diagnostic_dump.py | 75 ++++------ .../tools/diagnostic_dump_verl080.py | 20 +-- 6 files changed, 128 insertions(+), 165 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 6d657325..18031b92 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -317,11 +317,11 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: q_freqs = q_pos_emb.index_select(0, positions) ######### prefix-sharing diag: ON rope_freqs (per-layer) ######### try: - from prefix_sharing.tools.diagnostic_dump import dump_rope_freqs_on - dump_rope_freqs_on(q_freqs, attention_module.layer_number, + from prefix_sharing.tools.diagnostic_dump import dump_rope_freqs + dump_rope_freqs(q_freqs, attention_module.layer_number, attention_module.config.num_layers) except Exception as e: - print(f"rope_freqs_on dump failed: {e}") + print(f"rope_freqs dump failed: {e}") ######### prefix-sharing diag: ON rope_freqs (per-layer) ######### query = apply_rotary_pos_emb( query.unsqueeze(1), @@ -338,15 +338,15 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: ######### prefix-sharing diag: ON post-RoPE Q/K dump (per-layer) ######### try: - from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_emb_verl080 - dump_rope_emb_verl080( + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_postqk_verl080 + dump_rope_postqk_verl080( attention_module.layer_number, query, key, attention_module.config.num_layers, positions=packed_position_ids, ) except Exception as e: - print(f"rope_emb_layer dump failed: {e}") + print(f"rope_postqk_layer dump failed: {e}") ######### prefix-sharing diag: ON post-RoPE Q/K dump end ######### return query, key diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 9017238b..eb091d61 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -63,14 +63,14 @@ def patched_forward( sequence_len_offset=sequence_len_offset, inference_params=inference_params, ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump ##### + # ##### [PS-diag] OFF attn_outputs + rope_freqs + rope_postqk/preqk dump ##### if _diag_on: - import torch # 仅用于构造 debug 用的 positions（cmp 不依赖） + import torch # 仅用于构造 positions（freqs 切 per-token + Q/K debug） from prefix_sharing.tools.diagnostic_dump import ( - dump_attn_off, dump_rope_freqs_off, + dump_attn_off, dump_rope_freqs, ) from prefix_sharing.tools.diagnostic_dump_verl080 import ( - dump_rope_emb_verl080, dump_rope_preqk_verl080, + dump_rope_postqk_verl080, dump_rope_preqk_verl080, ) from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result @@ -84,24 +84,29 @@ def patched_forward( self.layer_number, _bs, self.config.num_layers) if rotary_pos_emb is not None: _q_pos_emb, _k_pos_emb = _unpack_rotary_pos_emb(rotary_pos_emb) - dump_rope_freqs_off(_q_pos_emb, self.layer_number, self.config.num_layers) + # OFF 标准 positions（每 segment 内 0..seg-1）：切 per-token freqs + Q/K debug + _off_positions = None + if (packed_seq_params is not None + and hasattr(packed_seq_params, "cu_seqlens_q_padded")): + _cu = packed_seq_params.cu_seqlens_q_padded + _off_positions = torch.cat([ + torch.arange(_cu[i + 1] - _cu[i]) + for i in range(len(_cu) - 1) + ]).long() + # rope_freqs：存 per-token 角度（与 ON 同款），统一 rope_freqs.pt + if _off_positions is not None: + dump_rope_freqs( + _q_pos_emb.index_select(0, _off_positions), + self.layer_number, self.config.num_layers, + ) # OFF post-RoPE Q/K：直接用 hook 截获的真实张量 # （original_forward 内部 apply_rotary_pos_emb 的返回值）， # captures[0]=Q, captures[1]=K。不再 get_query_key_value_tensors + 重算 RoPE。 if _rope_caps is not None and len(_rope_caps) >= 2: # 每个捕获是 {"pre": 旋转前, "post": 旋转后} _q_cap, _k_cap = _rope_caps[0], _rope_caps[1] - # positions 仅作 debug 记录；cmp 按 cu_seqlens+prefix_lens 对齐，不用它 - _off_positions = None - if (packed_seq_params is not None - and hasattr(packed_seq_params, "cu_seqlens_q_padded")): - _cu = packed_seq_params.cu_seqlens_q_padded - _off_positions = torch.cat([ - torch.arange(_cu[i + 1] - _cu[i]) - for i in range(len(_cu) - 1) - ]).long() # post-RoPE（旋转后） - dump_rope_emb_verl080( + dump_rope_postqk_verl080( self.layer_number, _q_cap["post"], _k_cap["post"], self.config.num_layers, positions=_off_positions, ) @@ -112,12 +117,12 @@ def patched_forward( ) else: print( - f"[PS-diag] OFF rope_emb L{self.layer_number}: " + f"[PS-diag] OFF rope_postqk L{self.layer_number}: " f"expected 2 captures (Q,K), got " f"{len(_rope_caps) if _rope_caps is not None else 'None'}; skip", flush=True, ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs_off + rope_emb dump end ##### + # ##### [PS-diag] OFF attn_outputs + rope_freqs + rope_postqk/preqk dump end ##### return _result # ── prefix-sharing path ── diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag.py b/prefix-sharing/prefix_sharing/tools/cmp_diag.py index e9c04062..77300bbd 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag.py @@ -387,7 +387,7 @@ def cmp_position_ids(dir_a: str, dir_b: str) -> CheckResult | None: # 2. RoPE encoding — absolute equality # ══════════════════════════════════════════════════════════════════ -def cmp_rope_emb(dir_on: str, dir_off: str) -> CheckResult | None: +def cmp_rope_postqk(dir_on: str, dir_off: str) -> CheckResult | None: """Compare post-RoPE Q/K between ON and OFF, aligned by position IDs. ON positions are absolute (preserved from original input). @@ -399,8 +399,8 @@ def cmp_rope_emb(dir_on: str, dir_off: str) -> CheckResult | None: 3. For reuser rows: ON uses absolute positions (prefix_len..), OFF uses relative (0..). Positions differ by design — skip direct comparison. """ - fa = os.path.join(dir_on, "rope_emb.pt") - fb = os.path.join(dir_off, "rope_emb.pt") + fa = os.path.join(dir_on, "rope_postqk.pt") + fb = os.path.join(dir_off, "rope_postqk.pt") if not os.path.exists(fa) or not os.path.exists(fb): return None a = torch.load(fa, weights_only=True) @@ -409,7 +409,7 @@ def cmp_rope_emb(dir_on: str, dir_off: str) -> CheckResult | None: return None la, lb = set(a.keys()), set(b.keys()) if la != lb: - return CheckResult(name="rope_emb", passed=False, + return CheckResult(name="rope_postqk", passed=False, metrics={"error": "layer set mismatch"}) max_diff_q = 0.0 @@ -482,7 +482,7 @@ def cmp_rope_emb(dir_on: str, dir_off: str) -> CheckResult | None: threshold = 1e-7 passed = (first_token_q_diff < threshold and first_token_k_diff < threshold and first_row_q_diff < threshold and first_row_k_diff < threshold) - return CheckResult(name="rope_emb", passed=passed, metrics={ + return CheckResult(name="rope_postqk", passed=passed, metrics={ "num_layers": num_layers, "first_row_len": first_row_len, "first_token_q_maxdiff": first_token_q_diff, @@ -862,7 +862,7 @@ def _print_pos_ids(r: CheckResult): def _print_rope(r: CheckResult): - print(_SEP_SINGLE + f"\n [rope_emb] {_CHECK if r.passed else _CROSS} {'PASS' if r.passed else 'FAIL'}") + print(_SEP_SINGLE + f"\n [rope_postqk] {_CHECK if r.passed else _CROSS} {'PASS' if r.passed else 'FAIL'}") print(_SEP_SINGLE) m = r.metrics if "error" in m: @@ -1111,7 +1111,7 @@ def main(): # ── ②b RoPE encoding (post-apply rotated Q/K) ── if not stop: - r = cmp_rope_emb(args.dir_on, args.dir_off) + r = cmp_rope_postqk(args.dir_on, args.dir_off) if r: all_results.append(r) _print_rope(r) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 32b0d30e..e9ce53cc 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -25,8 +25,9 @@ label_mask_{tag}.pt [B, L_max] bool [prompt-last,L_i-1) PPO loss 范围 logits.pt [N, V//tp] packed logits（ON 裁剪后 / OFF 完整） attn_outputs.pt dict {layer: [N, hidden]} per-layer packed attn output - rope_freqs_on.pt dict {layer: [T_on,1,1,D]} ON per-token RoPE 角度 - rope_freqs_off.pt dict {layer: [L0,1,1,D]} OFF raw 角度表（freqs[p]=p*inv_freq） + rope_freqs.pt dict {layer: [T,1,1,D]} per-token RoPE 角度（ON/OFF 同款） + rope_preqk.pt dict {layer: [T,H,D]} 旋转前 Q/K（pre-RoPE） + rope_postqk.pt dict {layer: [T,H,D]} 旋转后 Q/K（post-RoPE） prefix_lens.pt [B] ON=plan.prefix_lens / OFF=全0 cu_seqlens_q.pt [B+1] NestedTensor offsets（ON 裁剪后 / OFF 完整） cu_seqlens_q_logits.pt [B+1] logits packed 边界（同上） @@ -438,17 +439,17 @@ def cmp_packed_token(dir_on: str, dir_off: str, # Post-RoPE Q/K compare: per-layer + packed_token # ══════════════════════════════════════════════════════════════════ -# RoPE 对比阶段：**先 pre（旋转前，rope_preqk.pt）后 post（旋转后，rope_emb.pt）**。 +# RoPE 对比阶段：**先 pre（旋转前，rope_preqk.pt）后 post（旋转后，rope_postqk.pt）**。 # (stage, fname, label) — label 用作结果名前缀与打印 section 头。 _ROPE_STAGES: list[tuple[str, str, str]] = [ ("pre", "rope_preqk.pt", "rope_preqk"), - ("post", "rope_emb.pt", "rope_emb"), + ("post", "rope_postqk.pt", "rope_postqk"), ] -def _load_rope_emb(dir_path: str, layer: int, fname: str = "rope_emb.pt" +def _load_rope_postqk(dir_path: str, layer: int, fname: str = "rope_postqk.pt" ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - """Load Q/K for a single layer from ``fname`` (rope_emb.pt=post, rope_preqk.pt=pre). + """Load Q/K for a single layer from ``fname`` (rope_postqk.pt=post, rope_preqk.pt=pre). Returns ``(query, key)`` or ``(None, None)``. """ @@ -464,7 +465,7 @@ def _load_rope_emb(dir_path: str, layer: int, fname: str = "rope_emb.pt" return entry.get("query"), entry.get("key") -def _rope_emb_cos_for_layer(qa: torch.Tensor, ka: torch.Tensor, +def _rope_postqk_cos_for_layer(qa: torch.Tensor, ka: torch.Tensor, qb: torch.Tensor, kb: torch.Tensor, align_mask: torch.Tensor | None = None) -> dict: """单层 Q/K suffix 对齐 + per-token cosine（Q 和 K 分别算）。""" @@ -493,13 +494,13 @@ def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, align_mask = _build_attn_align_mask(dir_on, dir_off) if layer is not None: - q_on, k_on = _load_rope_emb(dir_on, layer, fname) - q_off, k_off = _load_rope_emb(dir_off, layer, fname) + q_on, k_on = _load_rope_postqk(dir_on, layer, fname) + q_off, k_off = _load_rope_postqk(dir_off, layer, fname) if q_on is None or q_off is None: return None need = align_mask is not None and q_on.shape[0] != q_off.shape[0] try: - d = _rope_emb_cos_for_layer(q_on, k_on, q_off, k_off, + d = _rope_postqk_cos_for_layer(q_on, k_on, q_off, k_off, align_mask if need else None) except ValueError as e: return CheckResult(name=f"{label}_L{layer}", passed=False, @@ -528,7 +529,7 @@ def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, continue need = align_mask is not None and qa.shape[0] != qb.shape[0] try: - results[lyr] = _rope_emb_cos_for_layer(qa, ka, qb, kb, + results[lyr] = _rope_postqk_cos_for_layer(qa, ka, qb, kb, align_mask if need else None) except ValueError as e: results[lyr] = {"error": str(e)} @@ -536,19 +537,19 @@ def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, metrics={"layers": results}) -def cmp_rope_emb_layer(dir_on: str, dir_off: str, layer: int | None, +def cmp_rope_postqk_layer(dir_on: str, dir_off: str, layer: int | None, stage: str = "post") -> CheckResult | None: """Q/K per-layer cosine（suffix 对齐），单 stage。 - stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_emb.pt（旋转后）。 + stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_postqk.pt（旋转后）。 调用方按 pre → rope_freqs → post 顺序分别调用，便于定位分歧出现在 RoPE 哪一步。 """ if stage == "pre": return _cmp_rope_stage_layer(dir_on, dir_off, layer, "rope_preqk.pt", "rope_preqk") - return _cmp_rope_stage_layer(dir_on, dir_off, layer, "rope_emb.pt", "rope_emb") + return _cmp_rope_stage_layer(dir_on, dir_off, layer, "rope_postqk.pt", "rope_postqk") -def _rope_emb_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, +def _rope_postqk_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, q_off: torch.Tensor | None, k_off: torch.Tensor | None, pos: int, align_mask: torch.Tensor | None @@ -584,9 +585,9 @@ def _rope_emb_vec_at_pos(q_on: torch.Tensor | None, k_on: torch.Tensor | None, def _diag_rope_pos_fail(q_on: torch.Tensor | None, q_off: torch.Tensor | None, pos: int, align_mask: torch.Tensor | None) -> str: - """rope_emb packed_token 取 [pos] 失败时的诊断串：区分 缺失 / 对齐失败 / pos 越界。""" + """rope_postqk packed_token 取 [pos] 失败时的诊断串：区分 缺失 / 对齐失败 / pos 越界。""" if q_on is None or q_off is None: - return f"rope_emb 该层在 {'ON' if q_on is None else 'OFF'} 侧缺失" + return f"rope_postqk 该层在 {'ON' if q_on is None else 'OFF'} 侧缺失" n_on, n_off = q_on.shape[0], q_off.shape[0] if align_mask is not None and n_on != n_off: msum = int(align_mask.sum()) @@ -607,9 +608,9 @@ def _cmp_rope_stage_token(dir_on: str, dir_off: str, pos: int, layer: int | None rope_layer = layer if layer is not None else ( _get_num_layers(dir_on) or _get_num_layers(dir_off)) if rope_layer: - q_on, k_on = _load_rope_emb(dir_on, rope_layer, fname) - q_off, k_off = _load_rope_emb(dir_off, rope_layer, fname) - vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) + q_on, k_on = _load_rope_postqk(dir_on, rope_layer, fname) + q_off, k_off = _load_rope_postqk(dir_off, rope_layer, fname) + vecs = _rope_postqk_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) if vecs is None: results.append(CheckResult( name=f"{label}_L{rope_layer}_pos{pos}", @@ -626,13 +627,13 @@ def _cmp_rope_stage_token(dir_on: str, dir_off: str, pos: int, layer: int | None return results -def cmp_rope_emb_token(dir_on: str, dir_off: str, +def cmp_rope_postqk_token(dir_on: str, dir_off: str, pos: int = 0, layer: int | None = None, align_mask: torch.Tensor | None = None, stage: str = "post") -> list[CheckResult]: """Q/K packed[pos] 对比（**suffix 对齐后**），单 stage。 - stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_emb.pt（旋转后）。 + stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_postqk.pt（旋转后）。 对 Q、K 分别输出 {label}_L{lyr}_Q_pos{pos} / {label}_L{lyr}_K_pos{pos}。 调用方按 pre → rope_freqs → post 顺序分别调用。 """ @@ -640,7 +641,7 @@ def cmp_rope_emb_token(dir_on: str, dir_off: str, return _cmp_rope_stage_token(dir_on, dir_off, pos, layer, align_mask, "rope_preqk.pt", "rope_preqk") return _cmp_rope_stage_token(dir_on, dir_off, pos, layer, align_mask, - "rope_emb.pt", "rope_emb") + "rope_postqk.pt", "rope_postqk") def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: @@ -675,19 +676,15 @@ def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: "cos_avg": cos_avg, "cos_min": cos_min}) -def _align_rope_freqs_layer(on_dict: dict, off_dict: dict, layer: int, - seqlens: list[int], +def _align_rope_freqs_layer(on_freqs: torch.Tensor, off_freqs: torch.Tensor, align_mask: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor] | None: - """单层 rope_freqs：OFF raw 表按 seqlens 重建 per-token + suffix 对齐。 + """单层 rope_freqs（per-token [T,1,1,D]）suffix 对齐。 - 返回 (on_aligned, off_aligned)，shape [N, 1, 1, D]；layer 缺失或对齐失败返回 None。 + 返回 (on_aligned, off_aligned) [N,1,1,D]；对齐失败返回 None。 供 cmp_rope_freqs（per-layer max_diff）与 cmp_rope_freqs_token（[pos] 角度向量）复用。 + ON/OFF 现在都是 per-token，直接对齐即可（不再从 raw 表重建）。 """ - if layer not in on_dict or layer not in off_dict: - return None - on_freqs = on_dict[layer] # [T_on,1,1,D] - off_freqs = torch.cat([off_dict[layer][:s, :, :, :] for s in seqlens], dim=0) # [T_off,1,1,D] try: return _align_packed(on_freqs, off_freqs, align_mask) except ValueError: @@ -696,16 +693,14 @@ def _align_rope_freqs_layer(on_dict: dict, off_dict: dict, layer: int, def cmp_rope_freqs(dir_on: str, dir_off: str, layer: int | None = None) -> CheckResult | None: - """对比 pre-RoPE 角度表（angle table，非 cos/sin）— suffix 对齐，应精确相等 max_diff==0。 + """对比 per-token RoPE 角度 — suffix 对齐，应精确相等 max_diff==0。 - ON ``rope_freqs_on.pt``: per-token 角度 dict {layer: [T_on,1,1,D]}（已 index_select 到 - packed_position_ids，每 token 实际旋转角度） - OFF ``rope_freqs_off.pt``: raw 角度表 dict {layer: [L0,1,1,D]}（freqs[p]=p*inv_freq，未切片） - OFF per-token 从 raw 表按 cu_seqlens_off 重建（每段 [:seg_len]），再 suffix 对齐。 - ``layer`` 给定则只比该层。角度是 RoPE 输入，应精确相等（max_diff==0）。 + ON/OFF 都存 per-token 角度 ``rope_freqs.pt`` {layer: [T,1,1,D]}（cos/sin 之前）， + suffix 对齐后逐元素比。角度是 RoPE 输入，应精确相等（max_diff==0）。 + ``layer`` 给定则只比该层。 """ - fa = os.path.join(dir_on, "rope_freqs_on.pt") - fb = os.path.join(dir_off, "rope_freqs_off.pt") + fa = os.path.join(dir_on, "rope_freqs.pt") + fb = os.path.join(dir_off, "rope_freqs.pt") if not os.path.exists(fa) or not os.path.exists(fb): return None on_dict = torch.load(fa, weights_only=True) @@ -721,21 +716,15 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, return CheckResult(name=_name, passed=False, metrics={"error": f"layer {layer} 不在双方 rope_freqs 中"}) - mb = _load_packed_meta(dir_off) - if mb is None: - return CheckResult(name=_name, passed=False, metrics={"error": "OFF cu_seqlens missing"}) - ma = _load_packed_meta(dir_on) - if ma is None: - return CheckResult(name=_name, passed=False, metrics={"error": "ON prefix_lens missing"}) - cu_off = mb["cu_seqlens"] - T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 - align_mask = _build_alignment_mask(cu_off, ma["prefix_lens"], T_off) - seqlens = (cu_off[1:] - cu_off[:-1]).tolist() + align_mask = _build_attn_align_mask(dir_on, dir_off) + if align_mask is None: + return CheckResult(name=_name, passed=False, + metrics={"error": "cu_seqlens/prefix_lens 缺失"}) max_diff = 0.0 mismatches: list[dict] = [] for lyr in layers: - _aligned = _align_rope_freqs_layer(on_dict, off_dict, lyr, seqlens, align_mask) + _aligned = _align_rope_freqs_layer(on_dict[lyr], off_dict[lyr], align_mask) if _aligned is None: continue on_a, off_a = _aligned @@ -770,8 +759,8 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, 取 ``layer``（默认最后一层）对齐后第 ``pos`` 个 token 的角度向量 [D]，比 ON/OFF。 角度是 RoPE 输入，应逐元素相等 → max_abs 应为 0。 """ - fa = os.path.join(dir_on, "rope_freqs_on.pt") - fb = os.path.join(dir_off, "rope_freqs_off.pt") + fa = os.path.join(dir_on, "rope_freqs.pt") + fb = os.path.join(dir_off, "rope_freqs.pt") if not os.path.exists(fa) or not os.path.exists(fb): return None on_dict = torch.load(fa, weights_only=True) @@ -784,17 +773,12 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, if rf_layer not in on_dict or rf_layer not in off_dict: return CheckResult(name=_name, metrics={"error": f"layer {rf_layer} 缺失"}) - mb = _load_packed_meta(dir_off) - ma = _load_packed_meta(dir_on) - if mb is None or ma is None: - return CheckResult(name=_name, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) - cu_off = mb["cu_seqlens"] - T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 if align_mask is None: - align_mask = _build_alignment_mask(cu_off, ma["prefix_lens"], T_off) - seqlens = (cu_off[1:] - cu_off[:-1]).tolist() + align_mask = _build_attn_align_mask(dir_on, dir_off) + if align_mask is None: + return CheckResult(name=_name, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) - _aligned = _align_rope_freqs_layer(on_dict, off_dict, rf_layer, seqlens, align_mask) + _aligned = _align_rope_freqs_layer(on_dict[rf_layer], off_dict[rf_layer], align_mask) if _aligned is None: return CheckResult(name=_name, metrics={"error": "对齐失败"}) on_a, off_a = _aligned @@ -894,7 +878,7 @@ def _shape_of(dir_path: str, filename: str) -> str: if isinstance(obj, dict): # per-layer dict（attn_outputs / rope_freqs_*）：显示层数 + 首层 shape sample = next(iter(obj.values())) if obj else None - # rope_emb.pt：每层值是 {"query","key"[,"positions"]} dict，取 query 的 shape 代表 + # rope_postqk.pt：每层值是 {"query","key"[,"positions"]} dict，取 query 的 shape 代表 if isinstance(sample, dict): _q = sample.get("query") sample_shape = f",Q{tuple(_q.shape)}" if _q is not None else "" @@ -955,8 +939,9 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, f"attention_mask_{tag}.pt", "logits.pt", "attn_outputs.pt", - "rope_emb.pt", + "rope_postqk.pt", "rope_preqk.pt", + "rope_freqs.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] @@ -1053,8 +1038,8 @@ def _print_per_layer(r: CheckResult): print() -def _print_rope_emb_per_layer(r: CheckResult): - _sec = "rope_preqk" if "preqk" in r.name else "rope_emb" +def _print_rope_postqk_per_layer(r: CheckResult): + _sec = "rope_preqk" if "preqk" in r.name else "rope_postqk" _stage = "Pre-RoPE" if "preqk" in r.name else "Post-RoPE" print(_SEP_SINGLE + f"\n [{_sec}] {_stage} Q/K Per-Layer Cosine Similarity") print(_SEP_SINGLE) @@ -1287,20 +1272,20 @@ def main(): # ── RoPE pipeline per-layer：pre Q/K → rope 角度 → post Q/K ── # 按计算顺序串联：旋转前 Q/K → 每token旋转角度(freqs) → 旋转后 Q/K， # 定位分歧出现在 RoPE 哪一步（pre 就偏=上游；freqs 偏=角度表；post 才偏=旋转应用）。 - r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer, stage="pre") + r = cmp_rope_postqk_layer(args.dir_on, args.dir_off, args.layer, stage="pre") if r: all_results.append(r) - _print_rope_emb_per_layer(r) + _print_rope_postqk_per_layer(r) r = cmp_rope_freqs(args.dir_on, args.dir_off, layer=args.layer) if r: all_results.append(r) _print_rope_freqs(r) - r = cmp_rope_emb_layer(args.dir_on, args.dir_off, args.layer, stage="post") + r = cmp_rope_postqk_layer(args.dir_on, args.dir_off, args.layer, stage="post") if r: all_results.append(r) - _print_rope_emb_per_layer(r) + _print_rope_postqk_per_layer(r) # ── packed: attention_output per-layer cos（RoPE 下游）── r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) @@ -1344,14 +1329,14 @@ def main(): # ── RoPE pipeline packed_token：pre Q/K → rope_freqs → post Q/K（指定 pos）── rope_pt_results: list[CheckResult] = [] - for r in cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, + for r in cmp_rope_postqk_token(args.dir_on, args.dir_off, pos, args.layer, align_mask=align_mask, stage="pre"): all_results.append(r); rope_pt_results.append(r); _print_packed_token(r) _rf = cmp_rope_freqs_token(args.dir_on, args.dir_off, pos, args.layer, align_mask=align_mask) if _rf is not None: all_results.append(_rf); rope_pt_results.append(_rf); _print_packed_token(_rf) - for r in cmp_rope_emb_token(args.dir_on, args.dir_off, pos, args.layer, + for r in cmp_rope_postqk_token(args.dir_on, args.dir_off, pos, args.layer, align_mask=align_mask, stage="post"): all_results.append(r); rope_pt_results.append(r); _print_packed_token(r) # rope packed_token top-K（pre Q/K + post Q/K；freqs 角度应精确相等，vec_metrics 已含 max_abs） @@ -1360,9 +1345,9 @@ def main(): _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) if rope_layer: for _stage, _fname, _label in _ROPE_STAGES: - q_on, k_on = _load_rope_emb(args.dir_on, rope_layer, _fname) - q_off, k_off = _load_rope_emb(args.dir_off, rope_layer, _fname) - vecs = _rope_emb_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) + q_on, k_on = _load_rope_postqk(args.dir_on, rope_layer, _fname) + q_off, k_off = _load_rope_postqk(args.dir_off, rope_layer, _fname) + vecs = _rope_postqk_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) if vecs is not None: qo, qf, ko, kf = vecs _print_topk_vec(qo.cpu(), qf.cpu(), args.topk, "val", diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py index 308a2401..3af88d19 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py @@ -10,7 +10,7 @@ 2. Positional encoding: - ``position_ids.pt`` — packed position ids [N] - - ``rope_emb.pt`` — per-layer RoPE encoding dict + - ``rope_postqk.pt`` — per-layer RoPE encoding dict {layer_idx: {"query": rotated_q, "key": rotated_k}} 3. 2D format: @@ -39,8 +39,7 @@ _META_SAVED: set[str] = set() # saved metadata keys (dedup per key) _ATTN_BUFFER: dict[int, torch.Tensor] | None = None # {layer_idx: tensor} _ROPE_BUFFER: dict[int, dict] | None = None # {layer_idx: {"query": q, "key": k}} -_ROPE_FREQS_ON_BUFFER: dict[int, torch.Tensor] | None = None # {layer_idx: q_freqs per-token} -_ROPE_FREQS_OFF_BUFFER: dict[int, torch.Tensor] | None = None # {layer_idx: q_pos_emb table} +_ROPE_FREQS_BUFFER: dict[int, torch.Tensor] | None = None # {layer_idx: per-token RoPE 角度} def _get_dump_dir() -> str | None: @@ -264,7 +263,7 @@ def _add_to_rope_buffer(layer_number: int, rotated_query: torch.Tensor, def _flush_rope_buffer(dump_dir: str) -> None: - """Write accumulated rope_emb dict to disk and clear buffer.""" + """Write accumulated rope_postqk dict to disk and clear buffer.""" global _ROPE_BUFFER if _ROPE_BUFFER is None: return @@ -272,67 +271,41 @@ def _flush_rope_buffer(dump_dir: str) -> None: _ROPE_BUFFER = None return try: - torch.save(_ROPE_BUFFER, os.path.join(dump_dir, "rope_emb.pt")) - _log.warning("rope_emb.pt saved (%d layers)", len(_ROPE_BUFFER)) + torch.save(_ROPE_BUFFER, os.path.join(dump_dir, "rope_postqk.pt")) + _log.warning("rope_postqk.pt saved (%d layers)", len(_ROPE_BUFFER)) _ROPE_BUFFER = None except Exception as e: - _log.warning("rope_emb.pt save failed: %s", e) + _log.warning("rope_postqk.pt save failed: %s", e) # ── RoPE angle dump (pre-apply, per-layer) ────────────────────── -def dump_rope_freqs_on(q_freqs: torch.Tensor, layer_number: int, - num_layers: int) -> None: - """Accumulate ON-mode per-token RoPE angles. Auto-flush on last layer. +def dump_rope_freqs(q_freqs: torch.Tensor, layer_number: int, + num_layers: int) -> None: + """Accumulate per-token RoPE angles (ON/OFF 共用). Auto-flush on last layer. - ``q_freqs`` is the result of ``q_pos_emb.index_select(0, packed_position_ids)`` - — shape [T_on, 1, 1, D], each token's actual rotation angles **before - cos/sin**. Stored as ``rope_freqs_on.pt`` (dict {layer_idx: tensor}). + ``q_freqs`` = 每 token 实际旋转角度（cos/sin 之前），shape [T, 1, 1, D]。 + ON 由 ``q_pos_emb.index_select(0, packed_position_ids)`` 得到；OFF 由 raw 表 + 按 cu_seqlens 切 per-token（每段 0..seg-1）得到。两边语义统一，写到各自 dump + 目录的 ``rope_freqs.pt``，cmp 侧 suffix 对齐后比 max_diff（角度应精确相等）。 """ - global _ROPE_FREQS_ON_BUFFER + global _ROPE_FREQS_BUFFER dump_dir = _get_dump_dir() if dump_dir is None: return - if _ROPE_FREQS_ON_BUFFER is None: - _ROPE_FREQS_ON_BUFFER = {} - _ROPE_FREQS_ON_BUFFER[layer_number] = q_freqs.detach().cpu().clone() + if _ROPE_FREQS_BUFFER is None: + _ROPE_FREQS_BUFFER = {} + _ROPE_FREQS_BUFFER[layer_number] = q_freqs.detach().cpu().clone() if layer_number == num_layers: if _rank0_only(): try: - torch.save(_ROPE_FREQS_ON_BUFFER, - os.path.join(dump_dir, "rope_freqs_on.pt")) - _log.warning("rope_freqs_on.pt saved (%d layers)", - len(_ROPE_FREQS_ON_BUFFER)) + torch.save(_ROPE_FREQS_BUFFER, + os.path.join(dump_dir, "rope_freqs.pt")) + _log.warning("rope_freqs.pt saved (%d layers)", + len(_ROPE_FREQS_BUFFER)) except Exception as e: - _log.warning("rope_freqs_on.pt save failed: %s", e) - _ROPE_FREQS_ON_BUFFER = None - - -def dump_rope_freqs_off(q_pos_emb: torch.Tensor, layer_number: int, - num_layers: int) -> None: - """Accumulate OFF-mode raw RoPE angle table. Auto-flush on last layer. - - ``q_pos_emb`` is the raw angle table (before per-token slicing) — - shape [L0, 1, 1, D], where ``freqs[p] = p * inv_freq``. - Stored as ``rope_freqs_off.pt`` (dict {layer_idx: tensor}). - """ - global _ROPE_FREQS_OFF_BUFFER - dump_dir = _get_dump_dir() - if dump_dir is None: - return - if _ROPE_FREQS_OFF_BUFFER is None: - _ROPE_FREQS_OFF_BUFFER = {} - _ROPE_FREQS_OFF_BUFFER[layer_number] = q_pos_emb.detach().cpu().clone() - if layer_number == num_layers: - if _rank0_only(): - try: - torch.save(_ROPE_FREQS_OFF_BUFFER, - os.path.join(dump_dir, "rope_freqs_off.pt")) - _log.warning("rope_freqs_off.pt saved (%d layers)", - len(_ROPE_FREQS_OFF_BUFFER)) - except Exception as e: - _log.warning("rope_freqs_off.pt save failed: %s", e) - _ROPE_FREQS_OFF_BUFFER = None + _log.warning("rope_freqs.pt save failed: %s", e) + _ROPE_FREQS_BUFFER = None # ── Position IDs dump ─────────────────────────────────────────── @@ -351,7 +324,7 @@ def dump_position_ids(position_ids: torch.Tensor) -> None: # ── RoPE encoding dump (per layer) ────────────────────────────── -def dump_rope_emb_layer(layer_number: int, rotated_query: torch.Tensor, +def dump_rope_postqk_layer(layer_number: int, rotated_query: torch.Tensor, rotated_key: torch.Tensor, num_layers: int, positions: torch.Tensor | None = None) -> None: """Accumulate one layer's post-RoPE query/key. Auto-flush on last layer. diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 65e5151b..d7121546 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -210,44 +210,44 @@ def dump_entropy_2d_verl080(ent_2d: torch.Tensor | None, tag: str) -> None: # Post-RoPE Q/K dump (per layer) # ════════════════════════════════════════════════════════════════ -_ROPE_EMB_BUFFER: dict[int, dict] | None = None +_ROPE_POSTQK_BUFFER: dict[int, dict] | None = None -def dump_rope_emb_verl080(layer_number: int, +def dump_rope_postqk_verl080(layer_number: int, rotated_query: torch.Tensor, rotated_key: torch.Tensor, num_layers: int, positions: torch.Tensor | None = None) -> None: - """Accumulate one layer's post-RoPE Q/K. Auto-flush to ``rope_emb.pt`` on last layer. + """Accumulate one layer's post-RoPE Q/K. Auto-flush to ``rope_postqk.pt`` on last layer. Format: ``{layer_idx: {"query": [T, H, D], "key": [T, H, D], "positions": [T] or None}}`` ON packed 只含 suffix（裁剪后），OFF packed 含完整序列。 cmp 侧用 prefix_lens + cu_seqlens 做 suffix 对齐后对比（同 attn_output 模式）。 positions 可选，用于手动排查时的位置回溯。 """ - global _ROPE_EMB_BUFFER + global _ROPE_POSTQK_BUFFER dump_dir = _get_dump_dir() if dump_dir is None: return - if _ROPE_EMB_BUFFER is None: - _ROPE_EMB_BUFFER = {} + if _ROPE_POSTQK_BUFFER is None: + _ROPE_POSTQK_BUFFER = {} entry = { "query": rotated_query.detach().cpu().clone(), "key": rotated_key.detach().cpu().clone(), } if positions is not None: entry["positions"] = positions.detach().cpu().clone() - _ROPE_EMB_BUFFER[layer_number] = entry + _ROPE_POSTQK_BUFFER[layer_number] = entry if layer_number == num_layers: - _flush_dict_buffer("rope_emb.pt", _ROPE_EMB_BUFFER, dump_dir) - _ROPE_EMB_BUFFER = None + _flush_dict_buffer("rope_postqk.pt", _ROPE_POSTQK_BUFFER, dump_dir) + _ROPE_POSTQK_BUFFER = None def _flush_dict_buffer(fname: str, buffer: dict, dump_dir: str) -> None: """rank0 直接 torch.save 一个 dict buffer。 不能用 _save_tensor：它对入参做 .detach().cpu().clone()，dict 没 .detach() → - AttributeError 被其 except 吞掉，文件永不写盘（rope_emb.pt 曾因此丢失）。 + AttributeError 被其 except 吞掉，文件永不写盘（rope_postqk.pt 曾因此丢失）。 entries 应在插入时已 detach().cpu().clone()。仿 _flush_attn_buffer。 """ import os as _os From b61e353d10ddbd8191a60b4e99b0f6e9896b4eaf Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 15:19:36 +0800 Subject: [PATCH 24/61] =?UTF-8?q?[diag]=20cmp:=20rope=20top-K=20dim=20?= =?UTF-8?q?=E8=B7=A8=20stage=20=E5=AF=B9=E9=BD=90=EF=BC=88postqk=20?= =?UTF-8?q?=E5=9F=BA=E5=87=86=20=E2=86=92=20preqk/freqs=20=E5=90=8C=20dim?= =?UTF-8?q?=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --topk 时 rope 三个 stage 的 dim 不再各自独立排序，而是以 rope_postqk 的 sort-err top-K dim 为基准，rope_preqk 显示同样 dim，rope_freqs 显示 dim%D（角度按 head_dim 共享），便于逐 dim 追溯误差来源： - preqk[dim] 已偏 → 上游（投影/hidden_states） - freqs[dim%D] 偏 → 角度/位置 ID - 两者都 ~0 但 postqk[dim] 偏 → 旋转应用本身 实现： - _print_topk_vec 返回选中的 dim 列表。 - 新增 _print_vec_at_dims（按指定 dim 打印，不排序）。 - 新增 _load_rope_freqs_vec_at_pos（对齐后取 [pos] 角度向量 [D]）。 - main rope topk 段：postqk Q/K sort-err top-K → preqk 同 dim、freqs dim%D。 - postqk 排序用 --sort-err（默认 abs），原先硬编码 val。 --- .../prefix_sharing/tools/cmp_diag_verl080.py | 99 ++++++++++++++++--- 1 file changed, 87 insertions(+), 12 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index e9ce53cc..bde53670 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -691,6 +691,36 @@ def _align_rope_freqs_layer(on_freqs: torch.Tensor, off_freqs: torch.Tensor, return None +def _load_rope_freqs_vec_at_pos(dir_on: str, dir_off: str, layer: int, pos: int, + align_mask: torch.Tensor | None = None + ) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """加载 rope_freqs 对齐后 [pos] 的角度向量 [D]，返回 (on_vec, off_vec) 或 (None, None)。 + + 供 top-K 跨 stage 对齐用（freqs dim = Q/K dim % D，角度按 head_dim 共享）。 + """ + fa = os.path.join(dir_on, "rope_freqs.pt") + fb = os.path.join(dir_off, "rope_freqs.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None, None + on_dict = torch.load(fa, weights_only=True) + off_dict = torch.load(fb, weights_only=True) + if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): + return None, None + if layer not in on_dict or layer not in off_dict: + return None, None + if align_mask is None: + align_mask = _build_attn_align_mask(dir_on, dir_off) + if align_mask is None: + return None, None + _aligned = _align_rope_freqs_layer(on_dict[layer], off_dict[layer], align_mask) + if _aligned is None: + return None, None + on_a, off_a = _aligned + if pos < 0 or pos >= on_a.shape[0]: + return None, None + return on_a[pos].reshape(-1), off_a[pos].reshape(-1) + + def cmp_rope_freqs(dir_on: str, dir_off: str, layer: int | None = None) -> CheckResult | None: """对比 per-token RoPE 角度 — suffix 对齐，应精确相等 max_diff==0。 @@ -1153,6 +1183,29 @@ def _print_topk_vec(on_vec: torch.Tensor, off_vec: torch.Tensor, for i in idx.tolist(): print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" f" {float(abs_err[i]):>12.6e}") + return idx.tolist() + + +def _print_vec_at_dims(on_vec: torch.Tensor, off_vec: torch.Tensor, + dims, label: str, show_rel: bool = True): + """在指定 dims 上打印 ON/OFF/ABS_ERR（不排序），跨 stage 对齐同一批 dim。 + + 供 rope 流水线 top-K 对齐：dims 取自 rope_postqk 的 sort-err top-K， + 在 rope_preqk / rope_freqs 上显示同样的 dim，逐 dim 追溯误差来源。 + """ + abs_err = (on_vec - off_vec).abs() + rel_err = abs_err / torch.maximum(on_vec.abs(), off_vec.abs()).clamp(min=1e-8) + print(f"\n [{label}] at {len(dims)} dims") + if show_rel: + print(f" {'DIM':>6s} {'ON':>14s} {'OFF':>14s} {'ABS_ERR':>12s} {'REL_ERR':>12s}") + for i in dims: + print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" + f" {float(abs_err[i]):>12.6e} {float(rel_err[i]):>12.6e}") + else: + print(f" {'DIM':>6s} {'ON':>14s} {'OFF':>14s} {'ABS_ERR':>12s}") + for i in dims: + print(f" {i:>6d} {float(on_vec[i]):>14.6e} {float(off_vec[i]):>14.6e}" + f" {float(abs_err[i]):>12.6e}") def _print_topk_2d(on_t: torch.Tensor, off_t: torch.Tensor, @@ -1339,22 +1392,44 @@ def main(): for r in cmp_rope_postqk_token(args.dir_on, args.dir_off, pos, args.layer, align_mask=align_mask, stage="post"): all_results.append(r); rope_pt_results.append(r); _print_packed_token(r) - # rope packed_token top-K（pre Q/K + post Q/K；freqs 角度应精确相等，vec_metrics 已含 max_abs） + # rope packed_token top-K —— dim 跨 stage 对齐：以 rope_postqk 的 sort-err top-K dim 为基准， + # rope_preqk 显示同样 dim，rope_freqs 显示 dim%D（角度按 head_dim 共享），逐 dim 追溯误差。 if args.topk > 0 and rope_pt_results: rope_layer = args.layer if args.layer is not None else ( _get_num_layers(args.dir_on) or _get_num_layers(args.dir_off)) if rope_layer: - for _stage, _fname, _label in _ROPE_STAGES: - q_on, k_on = _load_rope_postqk(args.dir_on, rope_layer, _fname) - q_off, k_off = _load_rope_postqk(args.dir_off, rope_layer, _fname) - vecs = _rope_postqk_vec_at_pos(q_on, k_on, q_off, k_off, pos, align_mask) - if vecs is not None: - qo, qf, ko, kf = vecs - _print_topk_vec(qo.cpu(), qf.cpu(), args.topk, "val", - f"{_label}_L{rope_layer}_Q_pos{pos}") - if ko is not None and kf is not None: - _print_topk_vec(ko.cpu(), kf.cpu(), args.topk, "val", - f"{_label}_L{rope_layer}_K_pos{pos}") + pre_q_on, pre_k_on = _load_rope_postqk(args.dir_on, rope_layer, "rope_preqk.pt") + pre_q_off, pre_k_off = _load_rope_postqk(args.dir_off, rope_layer, "rope_preqk.pt") + post_q_on, post_k_on = _load_rope_postqk(args.dir_on, rope_layer, "rope_postqk.pt") + post_q_off, post_k_off = _load_rope_postqk(args.dir_off, rope_layer, "rope_postqk.pt") + pre_vecs = _rope_postqk_vec_at_pos(pre_q_on, pre_k_on, pre_q_off, pre_k_off, pos, align_mask) + post_vecs = _rope_postqk_vec_at_pos(post_q_on, post_k_on, post_q_off, post_k_off, pos, align_mask) + freq_on, freq_off = _load_rope_freqs_vec_at_pos( + args.dir_on, args.dir_off, rope_layer, pos, align_mask) + if post_vecs is not None: + pqo, pqf, pko, pkf = post_vecs + # Q: postqk sort-err top-K → preqk / freqs 同 dim + q_dims = _print_topk_vec(pqo.cpu(), pqf.cpu(), args.topk, args.sort_err, + f"rope_postqk_L{rope_layer}_Q_pos{pos}") + if pre_vecs is not None: + _print_vec_at_dims(pre_vecs[0].cpu(), pre_vecs[1].cpu(), q_dims, + f"rope_preqk_L{rope_layer}_Q_pos{pos} (same dims)") + if freq_on is not None and freq_off is not None: + _D = freq_on.numel() + _print_vec_at_dims(freq_on.cpu(), freq_off.cpu(), + [d % _D for d in q_dims], + f"rope_freqs_L{rope_layer}_Q_pos{pos} (dim%D)") + # K: 同样 + if pko is not None and pkf is not None: + k_dims = _print_topk_vec(pko.cpu(), pkf.cpu(), args.topk, args.sort_err, + f"rope_postqk_L{rope_layer}_K_pos{pos}") + if pre_vecs is not None and pre_vecs[2] is not None and pre_vecs[3] is not None: + _print_vec_at_dims(pre_vecs[2].cpu(), pre_vecs[3].cpu(), k_dims, + f"rope_preqk_L{rope_layer}_K_pos{pos} (same dims)") + if freq_on is not None and freq_off is not None: + _print_vec_at_dims(freq_on.cpu(), freq_off.cpu(), + [d % _D for d in k_dims], + f"rope_freqs_L{rope_layer}_K_pos{pos} (dim%D)") # ── 2D: logprobs + entropy ── for fname, cname in [("logprobs", "logp"), ("entropy", "entropy")]: From dc70f97e2f940bbc88dfc88775bb7fdce896dd60 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 15:45:44 +0800 Subject: [PATCH 25/61] =?UTF-8?q?[fix]=20OFF=20rope=5Ffreqs=20dump=20?= =?UTF-8?q?=E5=B4=A9=20forward=EF=BC=9A=5Foff=5Fpositions=20=E7=BC=BA=20de?= =?UTF-8?q?vice=20=E5=8F=82=E6=95=B0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 重构 OFF diag 块把 _off_positions 上提时漏抄了 device=（原 device=_off_q.device）， arange 默认建在 CPU，而 _q_pos_emb 在 GPU → _q_pos_emb.index_select(0, _off_positions) 抛 RuntimeError(device 不匹配)。该 diag 块无 try/except，异常直接崩 OFF forward， 导致 rope_freqs/rope_postqk/rope_preqk 三个文件都没写盘。 修：arange 显式 device=_cu.device（cu_seqlens 在 GPU），并用 int() 把 GPU 标量转 Python int 再给 arange。 --- .../setup/patches/verl080_mcore0161_ms0160/attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index eb091d61..34dac1c8 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -89,8 +89,10 @@ def patched_forward( if (packed_seq_params is not None and hasattr(packed_seq_params, "cu_seqlens_q_padded")): _cu = packed_seq_params.cu_seqlens_q_padded + # device 必须显式到 _cu.device(GPU)：arange 默认 CPU，否则后面 + # _q_pos_emb.index_select(0, _off_positions) 会 device 不匹配崩 forward。 _off_positions = torch.cat([ - torch.arange(_cu[i + 1] - _cu[i]) + torch.arange(int(_cu[i + 1] - _cu[i]), device=_cu.device) for i in range(len(_cu) - 1) ]).long() # rope_freqs：存 per-token 角度（与 ON 同款），统一 rope_freqs.pt From cc40ce6c09965fbc92a2912d136aef753115aa94 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 16:29:30 +0800 Subject: [PATCH 26/61] =?UTF-8?q?[diag]=20attn=5Fkv=20=E5=AF=B9=E6=AF=94?= =?UTF-8?q?=EF=BC=9AON=20expanded=5Fkv=20vs=20OFF=20full=5Fkv=EF=BC=88pref?= =?UTF-8?q?ix=20=E5=A4=8D=E7=94=A8=E6=A0=A1=E9=AA=8C=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 定位「postqk 对但 attention_output 偏」的根因：dump attention 实际用的完整 KV， 逐元素比 ON(expanded) vs OFF(full)。 dump 侧： - ON(megatron_runtime): build_kv 后 dump expanded_key/value → expanded_kv.pt （prefix 复用 + suffix，attention 实际用的完整 KV）。 - OFF(attention): dump 完整 K(hook 截的 post-RoPE K) + 完整 V(get_qkv，V 不旋转 故忠实) → full_kv.pt。 cmp 侧(cmp_diag_verl080)： - cmp_attn_kv：ON expanded_kv.pt vs OFF full_kv.pt，逐层 K/V 分别 max_diff + cos， 两者 reshape [T,-1]，全量不需 align_mask。阈值极严（expanded 应精确==full）。 - _print_attn_kv 表插在 postqk 与 attn_output 之间。 - K/V 都 OK → attention 输入一致，偏差必来自 attention 计算/mask； K 或 V DIFF → bug 在 build_kv 的 prefix 复用。 - shapes 表加 expanded_kv.pt / full_kv.pt（per-side 文件，对侧显示 missing 正常）。 --- .../integrations/megatron_runtime.py | 9 ++ .../verl080_mcore0161_ms0160/attention.py | 13 +- .../prefix_sharing/tools/cmp_diag_verl080.py | 114 ++++++++++++++++++ .../tools/diagnostic_dump_verl080.py | 51 ++++++++ 4 files changed, 186 insertions(+), 1 deletion(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 18031b92..feb095b0 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -106,6 +106,15 @@ def prefix_attention( f"built expanded kv: expanded_key_shape={tuple(expanded_key.shape)}, expanded_value_shape={tuple(expanded_value.shape)}" ) + ######### prefix-sharing diag: ON expanded K/V dump（build_kv 输出，attention 实际用的完整 KV）######### + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_expanded_kv_on + dump_expanded_kv_on(layer_id, expanded_key, expanded_value, + attention_module.config.num_layers) + except Exception as _e: + print(f"expanded_kv dump failed: {_e}", flush=True) + ######### prefix-sharing diag: ON expanded K/V dump end ######### + # 注意力计算 core_attn_out = attention_backend.attention( query, diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 34dac1c8..e2865b1f 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -70,7 +70,7 @@ def patched_forward( dump_attn_off, dump_rope_freqs, ) from prefix_sharing.tools.diagnostic_dump_verl080 import ( - dump_rope_postqk_verl080, dump_rope_preqk_verl080, + dump_rope_postqk_verl080, dump_rope_preqk_verl080, dump_full_kv_off, ) from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result @@ -117,6 +117,17 @@ def patched_forward( self.layer_number, _q_cap["pre"], _k_cap["pre"], self.config.num_layers, ) + # full KV（post-RoPE K from hook + V from get_qkv），供与 ON expanded_kv 对比 + _off_v = self.get_query_key_value_tensors( + hidden_states, key_value_states, + split_qkv=True, output_gate=False, + )[2] + if _off_v.dim() > 2: + _off_v = _off_v.squeeze(1) + dump_full_kv_off( + self.layer_number, _k_cap["post"], _off_v, + self.config.num_layers, + ) else: print( f"[PS-diag] OFF rope_postqk L{self.layer_number}: " diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index bde53670..fe360f6e 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -822,6 +822,112 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, return CheckResult(name=_name, passed=m["max_abs"] == 0.0, metrics=m) +# ════════════════════════════════════════════════════════════════ +# Attention KV: ON expanded_kv vs OFF full_kv（prefix 复用校验） +# ════════════════════════════════════════════════════════════════ + +def _load_attn_kv(dir_path: str, layer: int, + fname: str) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """load {key, value} for a layer from fname. Returns (key, value) or (None, None).""" + fp = os.path.join(dir_path, fname) + if not os.path.exists(fp): + return None, None + d = torch.load(fp, weights_only=True) + if not isinstance(d, dict): + return None, None + entry = d.get(layer) + if entry is None: + return None, None + return entry.get("key"), entry.get("value") + + +def cmp_attn_kv(dir_on: str, dir_off: str, + layer: int | None = None) -> CheckResult | None: + """对比 ON expanded_kv vs OFF full_kv（K/V 分别），逐元素 max_diff + cos。 + + 两者都应是 full（prefix+suffix）且**逐元素相同**（prefix-sharing 的 KV 展开应精确还原 + 完整 KV）。相同 → attention 输入一致，attention_output 差异必来自 attention 计算/mask； + 不同 → bug 在 build_kv 的 prefix 复用（store/expand）。 + """ + fa = os.path.join(dir_on, "expanded_kv.pt") + fb = os.path.join(dir_off, "full_kv.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None + on_dict = torch.load(fa, weights_only=True) + off_dict = torch.load(fb, weights_only=True) + if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): + return None + layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) + if layer is not None: + layers = [l for l in layers if l == layer] + _name = f"attn_kv_L{layer}" if layer is not None else "attn_kv" + if not layers: + return CheckResult(name=_name, passed=False, + metrics={"error": f"layer {layer} 不在双方 attn_kv 中"}) + + per_layer: dict = {} + worst = {"max_diff": 0.0, "cos_min": 1.0} + for lyr in layers: + ek, ev = on_dict[lyr].get("key"), on_dict[lyr].get("value") + fk, fv = off_dict[lyr].get("key"), off_dict[lyr].get("value") + d: dict = {} + for _tag, (_a, _b) in [("K", (ek, fk)), ("V", (ev, fv))]: + if _a is None or _b is None: + d[_tag] = {"error": "缺失"} + continue + if _a.shape != _b.shape: + d[_tag] = {"error": f"shape mismatch ON{tuple(_a.shape)} vs OFF{tuple(_b.shape)}"} + continue + _af = _a.reshape(_a.shape[0], -1).float() + _bf = _b.reshape(_b.shape[0], -1).float() + _diff = (_af - _bf).abs() + _cos = _cosine_sim(_af, _bf, dim=-1) + _md = float(_diff.max()) + d[_tag] = {"max_diff": _md, + "cos_avg": float(_cos.mean()), "cos_min": float(_cos.min()), + "n_tokens": _af.shape[0]} + worst["max_diff"] = max(worst["max_diff"], _md) + worst["cos_min"] = min(worst["cos_min"], float(_cos.min())) + per_layer[lyr] = d + # expanded 应精确等于 full → 阈值极严 + passed = worst["max_diff"] < 1e-5 and worst["cos_min"] > 0.9999 + return CheckResult(name=_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst["max_diff"], + "cos_min": worst["cos_min"], "num_layers": len(layers)}) + + +def _print_attn_kv(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] ON expanded_kv vs OFF full_kv（K/V 逐元素）") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + print(f" {'LAYER':>6s} {'K_MAXDIFF':>12s} {'K_COS':>10s} " + f"{'V_MAXDIFF':>12s} {'V_COS':>10s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 12} {'─' * 10} {'─' * 8}") + bad = [] + for lyr in sorted(layers): + d = layers[lyr] + kd, vd = d.get("K", {}), d.get("V", {}) + if "error" in kd or "error" in vd: + print(f" {lyr:>6d} K:{kd.get('error','')} V:{vd.get('error','')}") + bad.append(lyr); continue + kmd, kcos = kd["max_diff"], kd["cos_avg"] + vmd, vcos = vd["max_diff"], vd["cos_avg"] + ok = kmd < 1e-5 and vmd < 1e-5 + if not ok: + bad.append(lyr) + print(f" {lyr:>6d} {kmd:>12.3e} {kcos:>10.6f} " + f"{vmd:>12.3e} {vcos:>10.6f} {'OK' if ok else 'DIFF':>8s}") + print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " + f"{_CHECK if r.passed else _CROSS} " + f"{'PASS（KV 一致）' if r.passed else 'FAIL（KV 不一致 → build_kv prefix 复用）'}") + if bad: + print(f" ⚠ 首个 KV 不一致层: {bad[0]}") + print() + + # ════════════════════════════════════════════════════════════════ # 2D mask loading # ════════════════════════════════════════════════════════════════ @@ -972,6 +1078,8 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, "rope_postqk.pt", "rope_preqk.pt", "rope_freqs.pt", + "expanded_kv.pt", + "full_kv.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] @@ -1340,6 +1448,12 @@ def main(): all_results.append(r) _print_rope_postqk_per_layer(r) + # ── packed: attention KV（ON expanded vs OFF full，prefix 复用校验）── + r = cmp_attn_kv(args.dir_on, args.dir_off, args.layer) + if r: + all_results.append(r) + _print_attn_kv(r) + # ── packed: attention_output per-layer cos（RoPE 下游）── r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) if r: diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index d7121546..8b1b8bf5 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -288,6 +288,57 @@ def dump_rope_preqk_verl080(layer_number: int, _ROPE_PREQK_BUFFER = None +_EXPANDED_KV_BUFFER: dict[int, dict] | None = None + + +def dump_expanded_kv_on(layer_number: int, expanded_key: torch.Tensor, + expanded_value: torch.Tensor, num_layers: int) -> None: + """ON: 累加 build_kv 输出（expanded K/V = prefix 复用 + suffix，全量），满层 flush ``expanded_kv.pt``。 + + Format: ``{layer_idx: {"key": [T,H,D], "value": [T,H,D]}}``。这是 ON attention 实际用的 + 完整 K/V，应与 OFF ``full_kv.pt`` 逐元素相同——验证 prefix-sharing 的 KV 展开/复用 + 是否正确还原了完整 KV。 + """ + global _EXPANDED_KV_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _EXPANDED_KV_BUFFER is None: + _EXPANDED_KV_BUFFER = {} + _EXPANDED_KV_BUFFER[layer_number] = { + "key": expanded_key.detach().cpu().clone(), + "value": expanded_value.detach().cpu().clone(), + } + if layer_number == num_layers: + _flush_dict_buffer("expanded_kv.pt", _EXPANDED_KV_BUFFER, dump_dir) + _EXPANDED_KV_BUFFER = None + + +_FULL_KV_BUFFER: dict[int, dict] | None = None + + +def dump_full_kv_off(layer_number: int, key: torch.Tensor, value: torch.Tensor, + num_layers: int) -> None: + """OFF: 累加完整 K（post-RoPE）/ V，满层 flush ``full_kv.pt``。 + + Format: ``{layer_idx: {"key": [T,H,D], "value": [T,H,D]}}``。key 应为 post-RoPE 完整 K + （与 ON expanded_key 同语义），value 为完整 V。供与 ON expanded_kv 逐元素对比。 + """ + global _FULL_KV_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _FULL_KV_BUFFER is None: + _FULL_KV_BUFFER = {} + _FULL_KV_BUFFER[layer_number] = { + "key": key.detach().cpu().clone(), + "value": value.detach().cpu().clone(), + } + if layer_number == num_layers: + _flush_dict_buffer("full_kv.pt", _FULL_KV_BUFFER, dump_dir) + _FULL_KV_BUFFER = None + + @contextlib.contextmanager def capture_rope_qk(): """Hook mcore 的 ``apply_rotary_pos_emb``，同时截获旋转前(pre)与旋转后(post)的 Q/K。 From 04d06c625578b6715b64045df1692de289fa7292 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 17:20:45 +0800 Subject: [PATCH 27/61] =?UTF-8?q?[diag]=20OFF=20V=20=E6=94=B9=20in-context?= =?UTF-8?q?=20=E6=88=AA=E8=8E=B7=EF=BC=88hook=20get=5Fquery=5Fkey=5Fvalue?= =?UTF-8?q?=5Ftensors=EF=BC=89=EF=BC=8C=E4=B8=8D=E5=86=8D=20re-call?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 之前 OFF 的 V 用 get_query_key_value_tensors 事后重算（re-call），而 K 走 apply_rotary hook 是 in-context。数据显示 L1「K 精确(0)、V 偏 1.2e-4」的不对称——K 和 V 同出 get_qkv， 理应同对/同偏，怀疑是 re-call 时 hidden_states 已被 forward 的 in-place 操作改写导致 V 失真。 改：capture_rope_qk(attention_module) 额外 hook 实例方法 get_query_key_value_tensors， 把返回的 V（第 3 元素）in-context 截下。yield (qk_captures, v_captures)。 attention.py OFF 块传 self、解包、用 _v_caps[-1] 作为 full_kv 的 V（删掉 re-call）。 重跑后看 attn_kv 的 V_MAXDIFF： - 若 V 也归零（和 K 一致）→ 之前是 dump artifact，expanded_kv 其实没问题； - 若 V 仍偏 → ON 侧 V 真有问题，深挖 build_kv/store。 --- .../verl080_mcore0161_ms0160/attention.py | 34 +++++++----- .../tools/diagnostic_dump_verl080.py | 55 ++++++++++++------- 2 files changed, 55 insertions(+), 34 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index e2865b1f..9bf7c18e 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -44,7 +44,7 @@ def patched_forward( _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None if _diag_on and rotary_pos_emb is not None: from prefix_sharing.tools.diagnostic_dump_verl080 import capture_rope_qk - _rope_cm = capture_rope_qk() + _rope_cm = capture_rope_qk(self) else: _rope_cm = nullcontext() with _rope_cm as _rope_caps: @@ -104,9 +104,14 @@ def patched_forward( # OFF post-RoPE Q/K：直接用 hook 截获的真实张量 # （original_forward 内部 apply_rotary_pos_emb 的返回值）， # captures[0]=Q, captures[1]=K。不再 get_query_key_value_tensors + 重算 RoPE。 - if _rope_caps is not None and len(_rope_caps) >= 2: + # _rope_caps = (qk_captures, v_captures)，均 in-context 截获 + if _rope_caps is not None: + _qk_caps, _v_caps = _rope_caps + else: + _qk_caps, _v_caps = None, None + if _qk_caps is not None and len(_qk_caps) >= 2: # 每个捕获是 {"pre": 旋转前, "post": 旋转后} - _q_cap, _k_cap = _rope_caps[0], _rope_caps[1] + _q_cap, _k_cap = _qk_caps[0], _qk_caps[1] # post-RoPE（旋转后） dump_rope_postqk_verl080( self.layer_number, _q_cap["post"], _k_cap["post"], @@ -117,22 +122,23 @@ def patched_forward( self.layer_number, _q_cap["pre"], _k_cap["pre"], self.config.num_layers, ) - # full KV（post-RoPE K from hook + V from get_qkv），供与 ON expanded_kv 对比 - _off_v = self.get_query_key_value_tensors( - hidden_states, key_value_states, - split_qkv=True, output_gate=False, - )[2] - if _off_v.dim() > 2: + # full KV（post-RoPE K + V，均 in-context 截获，不再事后 re-call） + _off_v = _v_caps[-1] if _v_caps else None + if _off_v is not None and _off_v.dim() > 2: _off_v = _off_v.squeeze(1) - dump_full_kv_off( - self.layer_number, _k_cap["post"], _off_v, - self.config.num_layers, - ) + if _off_v is not None: + dump_full_kv_off( + self.layer_number, _k_cap["post"], _off_v, + self.config.num_layers, + ) + else: + print(f"[PS-diag] OFF full_kv L{self.layer_number}: V 未截获; skip", + flush=True) else: print( f"[PS-diag] OFF rope_postqk L{self.layer_number}: " f"expected 2 captures (Q,K), got " - f"{len(_rope_caps) if _rope_caps is not None else 'None'}; skip", + f"{len(_qk_caps) if _qk_caps is not None else 'None'}; skip", flush=True, ) # ##### [PS-diag] OFF attn_outputs + rope_freqs + rope_postqk/preqk dump end ##### diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 8b1b8bf5..28005ba8 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -340,36 +340,51 @@ def dump_full_kv_off(layer_number: int, key: torch.Tensor, value: torch.Tensor, @contextlib.contextmanager -def capture_rope_qk(): - """Hook mcore 的 ``apply_rotary_pos_emb``，同时截获旋转前(pre)与旋转后(post)的 Q/K。 +def capture_rope_qk(attention_module): + """Hook apply_rotary_pos_emb（Q/K pre+post）+ get_query_key_value_tensors（V），全 in-context。 - mcore ``Attention.forward`` 的 THD prefill 路径每层恰好调用模块级 - ``apply_rotary_pos_emb`` 两次——先 Q 后 K（见 megatron/core/transformer/ - attention.py:1097,1110）。monkey-patch 该模块全局函数，把每次调用的输入 t - (pre-RoPE) 与返回值 out (post-RoPE) 存为 ``{"pre": t, "post": out}`` 追加到 - yield 的 list；``finally`` 还原原函数。 + mcore ``Attention.forward`` THD prefill 每层调模块级 ``apply_rotary_pos_emb`` 两次（先 Q 后 + K），把每次输入/返回存为 ``{"pre": t, "post": out}``。V 不走 rotary，所以单独 hook + ``attention_module.get_query_key_value_tensors``（实例级），截其返回的 V（第 3 个元素）。 - ON 路径不受影响：它在 megatron_runtime.py 里把 ``apply_rotary_pos_emb`` - import 进了自有命名空间，不经 attention 模块全局解析，patch 不到它。 + 全部 in-context 截获（在 original_forward 内部），避免事后 re-call get_query_key_value_tensors + 因 hidden_states 被 in-place 改写而失真——K 走 rotary hook 一向准确，V 之前用 re-call 才显得偏。 - 用法（OFF 分支，包住 original_forward 调用）:: + 用法:: - with capture_rope_qk() as caps: + with capture_rope_qk(self) as (qk_caps, v_caps): result = original_forward(...) - # caps[0]={"pre":Q_pre,"post":Q_post}, caps[1]={"pre":K_pre,"post":K_post} + # qk_caps[0]={"pre":Q_pre,"post":Q_post}, qk_caps[1]={...K...}; v_caps[-1]=V """ import megatron.core.transformer.attention as _attn_mod + import types as _types - _orig = _attn_mod.apply_rotary_pos_emb - captures: list = [] # 每元素 {"pre": 旋转前 t, "post": 旋转后 out} + _orig_arpe = _attn_mod.apply_rotary_pos_emb + _orig_get_qkv = attention_module.get_query_key_value_tensors # 已绑定方法 + qk_captures: list = [] + v_captures: list = [] - def _capturing(t, *args, **kwargs): - out = _orig(t, *args, **kwargs) - captures.append({"pre": t, "post": out}) + def _capturing_arpe(t, *args, **kwargs): + out = _orig_arpe(t, *args, **kwargs) + qk_captures.append({"pre": t, "post": out}) return out - _attn_mod.apply_rotary_pos_emb = _capturing + def _capturing_get_qkv(*args, **kwargs): + out = _orig_get_qkv(*args, **kwargs) + try: + v_captures.append(out[2]) # (Q, K, V[, gate]) → V + except Exception: + pass + return out + + _attn_mod.apply_rotary_pos_emb = _capturing_arpe + attention_module.get_query_key_value_tensors = _types.MethodType( + _capturing_get_qkv, attention_module) try: - yield captures + yield qk_captures, v_captures finally: - _attn_mod.apply_rotary_pos_emb = _orig + _attn_mod.apply_rotary_pos_emb = _orig_arpe + try: + del attention_module.get_query_key_value_tensors # 还原：删实例属性，回落到类方法 + except AttributeError: + pass From 68b0c44c52aa36d560f6c79448b0af0fe57a4c3e Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 17:32:32 +0800 Subject: [PATCH 28/61] =?UTF-8?q?[fix]=20capture=5Frope=5Fqk:=20get=5Fqkv?= =?UTF-8?q?=20=E5=8F=8C=E9=87=8D=E7=BB=91=E5=AE=9A=E5=AF=BC=E8=87=B4=20out?= =?UTF-8?q?put=5Fgate=20multiple=20values?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _orig_get_qkv 之前取成绑定方法（attention_module.get_query_key_value_tensors，self 已绑）， _capturing_get_qkv 又经 MethodType 绑定把 attention_module 作为 self_ 传入，原方法收到 两次 self，位置参数错位 → 'got multiple values for argument output_gate'。 修：_orig_get_qkv 改取未绑定的类方法函数 type(attention_module).get_query_key_value_tensors， _capturing_get_qkv 显式接收 self_ 并传给未绑定原函数（调用方经 MethodType 自动注入 self_）。 --- .../prefix_sharing/tools/diagnostic_dump_verl080.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 28005ba8..c278ff9f 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -360,7 +360,7 @@ def capture_rope_qk(attention_module): import types as _types _orig_arpe = _attn_mod.apply_rotary_pos_emb - _orig_get_qkv = attention_module.get_query_key_value_tensors # 已绑定方法 + _orig_get_qkv = type(attention_module).get_query_key_value_tensors # 未绑定的类方法函数 qk_captures: list = [] v_captures: list = [] @@ -369,8 +369,9 @@ def _capturing_arpe(t, *args, **kwargs): qk_captures.append({"pre": t, "post": out}) return out - def _capturing_get_qkv(*args, **kwargs): - out = _orig_get_qkv(*args, **kwargs) + def _capturing_get_qkv(self_, *args, **kwargs): + # MethodType 绑定后 self_=attention_module；_orig_get_qkv 是未绑定函数，需显式传 self_ + out = _orig_get_qkv(self_, *args, **kwargs) try: v_captures.append(out[2]) # (Q, K, V[, gate]) → V except Exception: From c7dcc29910c2460bcb00048a0e81dac66839b184 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 19:16:01 +0800 Subject: [PATCH 29/61] =?UTF-8?q?[diag]=20dump=20build=5Fkv=20=E8=BE=93?= =?UTF-8?q?=E5=85=A5=20V=EF=BC=8C=E5=AE=9A=E4=BD=8D=20V=20=E5=81=8F?= =?UTF-8?q?=E5=9C=A8=20build=5Fkv=20=E4=B9=8B=E5=89=8D=E8=BF=98=E6=98=AF?= =?UTF-8?q?=E4=B9=8B=E5=90=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit postqk 对、build_kv 出来 V 偏，需区分 V 是 build_kv 之前就偏（get_qkv/hidden_states） 还是 build_kv 引入。在 megatron_runtime build_kv 调用前 dump 进去的 value（raw V， get_qkv 出来、build_kv 之前），存 build_kv_input_v.pt。 cmp_build_kv_input_v：ON build_kv_input_v vs OFF full_kv V（suffix 对齐）。 输出 ON_T / OFF_T —— 若 ON_T < OFF_T 说明 ON 把 hidden_states 裁成 suffix-only。 判读：PASS（build_kv 前 V 一致）→ 偏由 build_kv 引入；FAIL（已偏）→ 根因在 get_qkv/hidden_states。 --- .../integrations/megatron_runtime.py | 7 ++ .../prefix_sharing/tools/cmp_diag_verl080.py | 88 +++++++++++++++++++ .../tools/diagnostic_dump_verl080.py | 23 +++++ 3 files changed, 118 insertions(+) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index feb095b0..e8034ad4 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -90,6 +90,13 @@ def prefix_attention( # 前缀共享：provider 存储激活值，reuser 拼接激活值 attention_backend = prefix_sharing_context.attention_backend or TorchReferenceBackend() + ######### prefix-sharing diag: ON build_kv 输入 V dump（build_kv 之前）######### + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_build_kv_input_v_on + dump_build_kv_input_v_on(layer_id, value, attention_module.config.num_layers) + except Exception as _e: + print(f"build_kv_input_v dump failed: {_e}", flush=True) + ######### prefix-sharing diag: ON build_kv 输入 V dump end ######### expanded_key, expanded_value = attention_backend.build_kv( key, value, diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index fe360f6e..5f05cfe5 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -928,6 +928,87 @@ def _print_attn_kv(r: CheckResult): print() +def cmp_build_kv_input_v(dir_on: str, dir_off: str, + layer: int | None = None) -> CheckResult | None: + """对比 ON build_kv 输入 V（build_kv 之前）vs OFF full_kv V（suffix 对齐）。 + + 定位 V 偏是在 build_kv 之前（get_qkv/hidden_states）还是 build_kv 引入。 + ON_T vs OFF_T 还能看出 ON 有没有把 hidden_states 裁成 suffix-only。 + """ + fa = os.path.join(dir_on, "build_kv_input_v.pt") + fb = os.path.join(dir_off, "full_kv.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None + on_dict = torch.load(fa, weights_only=True) + off_dict = torch.load(fb, weights_only=True) + if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): + return None + layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) + if layer is not None: + layers = [l for l in layers if l == layer] + _name = f"build_kv_input_v_L{layer}" if layer is not None else "build_kv_input_v" + if not layers: + return CheckResult(name=_name, passed=False, metrics={"error": "no layers"}) + + align_mask = _build_attn_align_mask(dir_on, dir_off) + per_layer: dict = {} + worst_md = 0.0 + worst_cos = 1.0 + for lyr in layers: + on_v = on_dict[lyr] + off_entry = off_dict[lyr] + off_v = off_entry.get("value") if isinstance(off_entry, dict) else None + if on_v is None or off_v is None: + per_layer[lyr] = {"error": "缺失"}; continue + on_f = on_v.reshape(on_v.shape[0], -1).float() + off_f = off_v.reshape(off_v.shape[0], -1).float() + on_T, off_T = int(on_v.shape[0]), int(off_v.shape[0]) + if align_mask is not None and on_f.shape[0] != off_f.shape[0]: + try: + on_f, off_f = _align_packed(on_f, off_f, align_mask) + except ValueError as e: + per_layer[lyr] = {"error": str(e), "on_T": on_T, "off_T": off_T} + continue + diff = (on_f - off_f).abs() + cos = _cosine_sim(on_f, off_f, dim=-1) + md = float(diff.max()) + per_layer[lyr] = {"max_diff": md, "cos_avg": float(cos.mean()), + "cos_min": float(cos.min()), "n_tokens": on_f.shape[0], + "on_T": on_T, "off_T": off_T} + worst_md = max(worst_md, md) + worst_cos = min(worst_cos, float(cos.min())) + passed = worst_md < 1e-5 + return CheckResult(name=_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_md, "cos_min": worst_cos}) + + +def _print_build_kv_input_v(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] ON build_kv 输入 V vs OFF full_kv V（suffix 对齐）") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS':>10s} " + f"{'ON_T':>8s} {'OFF_T':>8s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") + for lyr in sorted(layers): + d = layers[lyr] + if "max_diff" not in d: + print(f" {lyr:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") + continue + md, cos = d["max_diff"], d["cos_avg"] + ok = md < 1e-5 + _crop = " (cropped)" if d.get("on_T") != d.get("off_T") else "" + print(f" {lyr:>6d} {md:>12.3e} {cos:>10.6f} " + f"{d.get('on_T', '—'):>8} {d.get('off_T', '—'):>8} " + f"{'OK' if ok else 'DIFF':>8s}{_crop}") + print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " + f"{_CHECK if r.passed else _CROSS} " + f"{'PASS（build_kv 前 V 一致 → 偏由 build_kv 引入）' if r.passed else 'FAIL（build_kv 前 V 已偏 → 根因在 get_qkv/hidden_states）'}") + print() + + # ════════════════════════════════════════════════════════════════ # 2D mask loading # ════════════════════════════════════════════════════════════════ @@ -1080,6 +1161,7 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, "rope_freqs.pt", "expanded_kv.pt", "full_kv.pt", + "build_kv_input_v.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] @@ -1454,6 +1536,12 @@ def main(): all_results.append(r) _print_attn_kv(r) + # ── packed: build_kv 输入 V（build_kv 前）vs OFF full_kv V —— 定位 V 偏在 build_kv 之前还是之后 ── + r = cmp_build_kv_input_v(args.dir_on, args.dir_off, args.layer) + if r: + all_results.append(r) + _print_build_kv_input_v(r) + # ── packed: attention_output per-layer cos（RoPE 下游）── r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) if r: diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index c278ff9f..46b10325 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -339,6 +339,29 @@ def dump_full_kv_off(layer_number: int, key: torch.Tensor, value: torch.Tensor, _FULL_KV_BUFFER = None +_BUILD_KV_INPUT_V_BUFFER: dict[int, torch.Tensor] | None = None + + +def dump_build_kv_input_v_on(layer_number: int, value: torch.Tensor, + num_layers: int) -> None: + """ON: build_kv 输入的 V（get_qkv 出来、build_kv 之前的 raw V）。满层 flush ``build_kv_input_v.pt``。 + + Format: ``{layer_idx: [T_on, ...]}``。供与 OFF ``full_kv.pt`` 的 V 做 suffix 对比—— + 定位 V 是在 build_kv 之前（get_qkv/hidden_states）就偏，还是 build_kv 引入。 + T_on vs T_off 还能看出 ON 有没有把 hidden_states 裁剪成 suffix-only。 + """ + global _BUILD_KV_INPUT_V_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _BUILD_KV_INPUT_V_BUFFER is None: + _BUILD_KV_INPUT_V_BUFFER = {} + _BUILD_KV_INPUT_V_BUFFER[layer_number] = value.detach().cpu().clone() + if layer_number == num_layers: + _flush_dict_buffer("build_kv_input_v.pt", _BUILD_KV_INPUT_V_BUFFER, dump_dir) + _BUILD_KV_INPUT_V_BUFFER = None + + @contextlib.contextmanager def capture_rope_qk(attention_module): """Hook apply_rotary_pos_emb（Q/K pre+post）+ get_query_key_value_tensors（V），全 in-context。 From 2235b1c4cad51075d28ce1e61a28dad8247649b0 Mon Sep 17 00:00:00 2001 From: Boundless Date: Sat, 27 Jun 2026 20:43:24 +0800 Subject: [PATCH 30/61] =?UTF-8?q?[diag]=20preqk/postqk=20per-layer=20?= =?UTF-8?q?=E5=8A=A0=20max=5Fdiff=EF=BC=8C=E4=B8=8E=20build=5Fkv=5Finput?= =?UTF-8?q?=5Fv=20=E7=9A=84=20V=20=E5=90=8C=E5=8F=A3=E5=BE=84=E5=AF=B9?= =?UTF-8?q?=E6=AF=94?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _rope_postqk_cos_for_layer 额外算 Q_max_diff / K_max_diff（对齐后整体 abs max）。 printer 表格改显 max_diff（替换 cos_min 列；cos_min 仍在 metrics 用于 WARN 判定）。 目的：验证 Q/K 的 max_diff 是否和 V 的 ~6e-5 一致。若一致 → Q/K/V 同源同偏（裁剪 matmul tiling）；若 Q/K≈0 而 V~6e-5 → V 走了不同路径，需深查模型 attention 实现。 --- .../prefix_sharing/tools/cmp_diag_verl080.py | 22 +++++++++++-------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 5f05cfe5..d19edba7 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -485,6 +485,8 @@ def _rope_postqk_cos_for_layer(qa: torch.Tensor, ka: torch.Tensor, "n_tokens": qa_flat.shape[0], "Q_cos_avg": float(q_cos.mean()), "Q_cos_min": float(q_cos.min()), "K_cos_avg": float(k_cos.mean()), "K_cos_min": float(k_cos.min()), + "Q_max_diff": float((qa_flat - qb_flat).abs().max()), + "K_max_diff": float((ka_flat - kb_flat).abs().max()), } @@ -1265,10 +1267,10 @@ def _print_rope_postqk_per_layer(r: CheckResult): print(_SEP_SINGLE) layers = r.metrics.get("layers") if isinstance(layers, dict): - print(f" {'LAYER':>6s} {'Q_COS_AVG':>14s} {'Q_COS_MIN':>14s} " - f"{'K_COS_AVG':>14s} {'K_COS_MIN':>14s} " + print(f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} " + f"{'K_MAXDIFF':>12s} {'K_COS_AVG':>12s} " f"{'TOKENS':>8s} {'STATUS':>8s}") - print(f" {'─' * 6} {'─' * 14} {'─' * 14} {'─' * 14} {'─' * 14} " + print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 12} " f"{'─' * 8} {'─' * 8}") bad = [] for lyr in sorted(layers.keys()): @@ -1279,21 +1281,23 @@ def _print_rope_postqk_per_layer(r: CheckResult): continue ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) - print(f" {lyr:>6d} {d['Q_cos_avg']:>14.6e} {d['Q_cos_min']:>14.6e} " - f"{d['K_cos_avg']:>14.6e} {d['K_cos_min']:>14.6e} " + print(f" {lyr:>6d} {d.get('Q_max_diff', 0.0):>12.3e} " + f"{d['Q_cos_avg']:>12.6e} " + f"{d.get('K_max_diff', 0.0):>12.3e} {d['K_cos_avg']:>12.6e} " f"{d['n_tokens']:>8d} {'PASS' if ok else 'WARN':>8s}") if not ok: bad.append(lyr) if bad: print(f"\n ⚠ First deviating layer: {bad[0]}") + print(f" （Q/K max_diff 与 build_kv_input_v 的 V max_diff 同口径，可直接对比）") elif "Q_cos_avg" in r.metrics: d = r.metrics ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) - print(f" L{d['layer']} Q_cos_avg={d['Q_cos_avg']:.6e} " - f"Q_cos_min={d['Q_cos_min']:.6e} " - f"K_cos_avg={d['K_cos_avg']:.6e} K_cos_min={d['K_cos_min']:.6e} " - f"{'PASS' if ok else 'WARN'}") + print(f" L{d['layer']} Q_maxdiff={d.get('Q_max_diff', 0.0):.3e} " + f"Q_cos_avg={d['Q_cos_avg']:.6e} " + f"K_maxdiff={d.get('K_max_diff', 0.0):.3e} " + f"K_cos_avg={d['K_cos_avg']:.6e} {'PASS' if ok else 'WARN'}") elif "error" in r.metrics: print(f" {_CROSS} {r.metrics['error']}") print() From bfc06faff0c26d52deadd094d64d995baf09e387 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 11:36:09 +0800 Subject: [PATCH 31/61] =?UTF-8?q?[refactor]=20ON=20pre-RoPE=20Q/K/V=20dump?= =?UTF-8?q?=20=E7=BB=9F=E4=B8=80=E5=88=B0=20attention.py=20squeeze=20?= =?UTF-8?q?=E4=B9=8B=E5=90=8E?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 之前 ON 的 pre-RoPE Q/K dump 在 _apply_positioned_rope（megatron_runtime）， V dump 在 build_kv 前（megatron_runtime），分散两处。现统一到 attention.py ON 分支 get_qkv + squeeze(1) 之后、_apply_positioned_rope/build_kv 之前——数据等价 （squeeze 无损），但集中一处、口径清晰，与 OFF baseline（hook 在 get_qkv 输出） 同点。 删掉 megatron_runtime 里分散的 dump_rope_preqk_verl080 / dump_build_kv_input_v_on 两处调用。 --- .../integrations/megatron_runtime.py | 16 ---------------- .../verl080_mcore0161_ms0160/attention.py | 19 +++++++++++++++++++ 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index e8034ad4..90f88019 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -90,13 +90,6 @@ def prefix_attention( # 前缀共享：provider 存储激活值，reuser 拼接激活值 attention_backend = prefix_sharing_context.attention_backend or TorchReferenceBackend() - ######### prefix-sharing diag: ON build_kv 输入 V dump（build_kv 之前）######### - try: - from prefix_sharing.tools.diagnostic_dump_verl080 import dump_build_kv_input_v_on - dump_build_kv_input_v_on(layer_id, value, attention_module.config.num_layers) - except Exception as _e: - print(f"build_kv_input_v dump failed: {_e}", flush=True) - ######### prefix-sharing diag: ON build_kv 输入 V dump end ######### expanded_key, expanded_value = attention_backend.build_kv( key, value, @@ -177,15 +170,6 @@ def _apply_positioned_rope( positions = packed_position_ids.to(device=query.device, dtype=torch.long) max_needed = positions.max().item() + 1 - ######### prefix-sharing diag: pre-RoPE Q/K dump（旋转前，尚未加位置编码）######### - try: - from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_preqk_verl080 - dump_rope_preqk_verl080(attention_module.layer_number, query, key, - attention_module.config.num_layers) - except Exception as _e: - print(f"rope_preqk (pre-RoPE) dump failed: {_e}", flush=True) - ######### prefix-sharing diag: pre-RoPE Q/K dump end ######### - # 当 packed_position_ids 所需要的最大 position id 超过了 q_pos_emb / k_pos_emb 的当前长度时， # 就需要对 q_pos_emb / k_pos_emb 进行扩展。 # THD 模式下生成的 pos_emb 仅覆盖 positions 0 .. max_seqlen_q-1 这段范围， diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 9bf7c18e..4458a0f7 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -159,6 +159,25 @@ def patched_forward( key = key.squeeze(1) value = value.squeeze(1) + # ##### [PS-diag] ON pre-RoPE Q/K/V 统一 dump（get_qkv 之后、RoPE 之前）##### + # 全部在此点 dump（squeeze 后、_apply_positioned_rope / build_kv 之前）， + # 与 OFF baseline（hook 在 get_qkv 输出处截）同口径，集中对比，避免分散。 + import os as _os + if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + from prefix_sharing.tools.diagnostic_dump_verl080 import ( + dump_rope_preqk_verl080, dump_build_kv_input_v_on, + ) + try: + dump_rope_preqk_verl080(self.layer_number, query, key, + self.config.num_layers) + except Exception as _e: + print(f"rope_preqk (pre-RoPE Q/K) dump failed: {_e}", flush=True) + try: + dump_build_kv_input_v_on(self.layer_number, value, + self.config.num_layers) + except Exception as _e: + print(f"build_kv_input_v (pre-RoPE V) dump failed: {_e}", flush=True) + # delegate to verified integrations code from prefix_sharing.integrations.megatron_runtime import ( prefix_attention, From 5e68962b91620eb79d3bc9fa50ac3715466ca70c Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 11:50:28 +0800 Subject: [PATCH 32/61] =?UTF-8?q?[refactor]=20OFF=20Q/K/V=20=E7=BB=9F?= =?UTF-8?q?=E4=B8=80=E4=BB=8E=20get=5Fqkv=20=E8=BF=94=E5=9B=9E=E6=88=AA?= =?UTF-8?q?=EF=BC=88=E4=B8=8E=20ON=20=E5=AF=B9=E7=A7=B0=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit capture_rope_qk 改为从 get_query_key_value_tensors 返回统一截 (Q,K,V) pre-squeeze， 不再 Q/K 走 apply_rotary 输入、V 走 get_qkv 返回两个不对称来源。apply_rotary hook 只保留 返回（post-RoPE Q/K）。 OFF diag 块用 get_qkv 截到的 Q/K/V（手动 squeeze，与 ON 侧 get_qkv+squeeze 后 dump 同口径） dump rope_preqk + full_kv。两侧完全对称，避免 Q/K/V 来源不同引入的干扰。 yield 改为 (qk_caps, qkv_caps)，qkv_caps[-1]=(Q_pre,K_pre,V_pre)。 --- .../verl080_mcore0161_ms0160/attention.py | 42 +++++++++---------- .../tools/diagnostic_dump_verl080.py | 32 +++++++------- 2 files changed, 35 insertions(+), 39 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 4458a0f7..b0cb7e4e 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -101,38 +101,36 @@ def patched_forward( _q_pos_emb.index_select(0, _off_positions), self.layer_number, self.config.num_layers, ) - # OFF post-RoPE Q/K：直接用 hook 截获的真实张量 - # （original_forward 内部 apply_rotary_pos_emb 的返回值）， - # captures[0]=Q, captures[1]=K。不再 get_query_key_value_tensors + 重算 RoPE。 - # _rope_caps = (qk_captures, v_captures)，均 in-context 截获 + # _rope_caps = (qk_caps[post-RoPE], qkv_caps[pre-RoPE Q/K/V]) + # 与 ON 侧对称：ON 在 get_qkv+squeeze 之后 dump Q/K/V；OFF 这里用 hook + # 截的 get_qkv 返回（pre-squeeze），手动 squeeze 后与 ON 同口径。 if _rope_caps is not None: - _qk_caps, _v_caps = _rope_caps + _qk_caps, _qkv_caps = _rope_caps else: - _qk_caps, _v_caps = None, None + _qk_caps, _qkv_caps = None, None + # post-RoPE Q/K（apply_rotary 返回） if _qk_caps is not None and len(_qk_caps) >= 2: - # 每个捕获是 {"pre": 旋转前, "post": 旋转后} - _q_cap, _k_cap = _qk_caps[0], _qk_caps[1] - # post-RoPE（旋转后） + _q_post, _k_post = _qk_caps[0]["post"], _qk_caps[1]["post"] dump_rope_postqk_verl080( - self.layer_number, _q_cap["post"], _k_cap["post"], + self.layer_number, _q_post, _k_post, self.config.num_layers, positions=_off_positions, ) - # pre-RoPE（旋转前，纯 QKV 投影） - dump_rope_preqk_verl080( - self.layer_number, _q_cap["pre"], _k_cap["pre"], - self.config.num_layers, - ) - # full KV（post-RoPE K + V，均 in-context 截获，不再事后 re-call） - _off_v = _v_caps[-1] if _v_caps else None - if _off_v is not None and _off_v.dim() > 2: - _off_v = _off_v.squeeze(1) - if _off_v is not None: + # pre-RoPE Q/K + V：从 get_qkv 返回统一取（与 ON 同源） + if _qkv_caps: + _pre_q, _pre_k, _pre_v = _qkv_caps[-1][:3] + # squeeze 到 [T, H, D]（与 ON 侧 squeeze 后 dump 一致） + _sq = lambda _t: _t.squeeze(1) if _t.dim() > 2 else _t + _pre_q, _pre_k, _pre_v = _sq(_pre_q), _sq(_pre_k), _sq(_pre_v) + dump_rope_preqk_verl080( + self.layer_number, _pre_q, _pre_k, + self.config.num_layers, + ) dump_full_kv_off( - self.layer_number, _k_cap["post"], _off_v, + self.layer_number, _k_post, _pre_v, self.config.num_layers, ) else: - print(f"[PS-diag] OFF full_kv L{self.layer_number}: V 未截获; skip", + print(f"[PS-diag] OFF preqkv L{self.layer_number}: get_qkv 未截获; skip", flush=True) else: print( diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 46b10325..2b1da1d1 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -364,39 +364,37 @@ def dump_build_kv_input_v_on(layer_number: int, value: torch.Tensor, @contextlib.contextmanager def capture_rope_qk(attention_module): - """Hook apply_rotary_pos_emb（Q/K pre+post）+ get_query_key_value_tensors（V），全 in-context。 + """Hook apply_rotary_pos_emb（post-RoPE Q/K）+ get_query_key_value_tensors（pre-RoPE Q/K/V）。 - mcore ``Attention.forward`` THD prefill 每层调模块级 ``apply_rotary_pos_emb`` 两次（先 Q 后 - K），把每次输入/返回存为 ``{"pre": t, "post": out}``。V 不走 rotary，所以单独 hook - ``attention_module.get_query_key_value_tensors``（实例级），截其返回的 V（第 3 个元素）。 - - 全部 in-context 截获（在 original_forward 内部），避免事后 re-call get_query_key_value_tensors - 因 hidden_states 被 in-place 改写而失真——K 走 rotary hook 一向准确，V 之前用 re-call 才显得偏。 + 两个 hook，全 in-context： + - get_qkv 返回值统一截 Q/K/V（pre-squeeze，与 ON 侧 get_qkv+squeeze 之后 dump 同源）。 + - apply_rotary_pos_emb 截 post-RoPE Q/K（返回值）。 + Q/K/V 都从 get_qkv 一次返回取（pre-RoPE），避免 ON 侧 Q/K/V 也走两个不对称来源。 用法:: - with capture_rope_qk(self) as (qk_caps, v_caps): + with capture_rope_qk(self) as (qk_caps, qkv_caps): result = original_forward(...) - # qk_caps[0]={"pre":Q_pre,"post":Q_post}, qk_caps[1]={...K...}; v_caps[-1]=V + # qk_caps[0]={"post":Q_post}, qk_caps[1]={"post":K_post}（apply_rotary 返回） + # qkv_caps[-1] = (Q_pre, K_pre, V_pre)（get_qkv 返回，pre-squeeze） """ import megatron.core.transformer.attention as _attn_mod import types as _types _orig_arpe = _attn_mod.apply_rotary_pos_emb - _orig_get_qkv = type(attention_module).get_query_key_value_tensors # 未绑定的类方法函数 - qk_captures: list = [] - v_captures: list = [] + _orig_get_qkv = type(attention_module).get_query_key_value_tensors + qk_captures: list = [] # apply_rotary 返回（post-RoPE） + qkv_captures: list = [] # get_qkv 返回 (Q, K, V) pre-squeeze def _capturing_arpe(t, *args, **kwargs): out = _orig_arpe(t, *args, **kwargs) - qk_captures.append({"pre": t, "post": out}) + qk_captures.append({"post": out}) return out def _capturing_get_qkv(self_, *args, **kwargs): - # MethodType 绑定后 self_=attention_module；_orig_get_qkv 是未绑定函数，需显式传 self_ out = _orig_get_qkv(self_, *args, **kwargs) try: - v_captures.append(out[2]) # (Q, K, V[, gate]) → V + qkv_captures.append(tuple(out[:3])) # (Q, K, V) except Exception: pass return out @@ -405,10 +403,10 @@ def _capturing_get_qkv(self_, *args, **kwargs): attention_module.get_query_key_value_tensors = _types.MethodType( _capturing_get_qkv, attention_module) try: - yield qk_captures, v_captures + yield qk_captures, qkv_captures finally: _attn_mod.apply_rotary_pos_emb = _orig_arpe try: - del attention_module.get_query_key_value_tensors # 还原：删实例属性，回落到类方法 + del attention_module.get_query_key_value_tensors except AttributeError: pass From 4cd87dceb04be33dbb88c53d3a55e5be9d1045f2 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 12:00:47 +0800 Subject: [PATCH 33/61] =?UTF-8?q?[diag]=20OFF=20pre-RoPE=20Q/K/V=20?= =?UTF-8?q?=E4=BE=B5=E5=85=A5=E5=BC=8F=20dump=EF=BC=88megatron=20attention?= =?UTF-8?q?.py=20squeeze=20=E4=B9=8B=E5=90=8E=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 在 megatron attention.py 的 squeeze if-block 之后（line 1070）直接插入 dump， post-squeeze、pre-RoPE，与 ON 侧（attention.py ON 分支 get_qkv+squeeze 之后 dump） 完全对称——同一代码结构、同一处理点，避免 hook 式截取的来源不对称。 OFF diag 块移除 hook 式 rope_preqk dump（megatron 已侵入式 dump，避免双写）。 保留 apply_rotary hook（post-RoPE Q/K）+ get_qkv hook（full_kv 的 V）。 --- .../megatron/core/transformer/attention.py | 15 +++++++++++++ .../verl080_mcore0161_ms0160/attention.py | 21 +++++++------------ 2 files changed, 22 insertions(+), 14 deletions(-) diff --git a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py index 2200b558..31a4f85c 100644 --- a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py +++ b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py @@ -1069,6 +1069,21 @@ def forward( value = value.squeeze(1) nvtx_range_pop(suffix="adjust_key_value") + # [PS-diag] OFF pre-RoPE Q/K/V dump — 侵入式，post-squeeze、pre-RoPE， + # 与 ON 侧（attention.py ON 分支 get_qkv+squeeze 之后 dump）完全对称。 + import os as _ps_os + if _ps_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import ( + dump_rope_preqk_verl080, dump_build_kv_input_v_on, + ) + dump_rope_preqk_verl080(self.layer_number, query, key, + self.config.num_layers) + dump_build_kv_input_v_on(self.layer_number, value, + self.config.num_layers) + except Exception as _ps_e: + print(f"[PS-diag] OFF preqkv dump failed: {_ps_e}", flush=True) + # ================================================ # relative positional embedding (rotary embedding) # ================================================ diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index b0cb7e4e..778eb8d9 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -102,8 +102,8 @@ def patched_forward( self.layer_number, self.config.num_layers, ) # _rope_caps = (qk_caps[post-RoPE], qkv_caps[pre-RoPE Q/K/V]) - # 与 ON 侧对称：ON 在 get_qkv+squeeze 之后 dump Q/K/V；OFF 这里用 hook - # 截的 get_qkv 返回（pre-squeeze），手动 squeeze 后与 ON 同口径。 + # pre-RoPE Q/K/V 已由 megatron attention.py 侵入式 dump（line 1070 之后）， + # 这里只处理 post-RoPE Q/K（apply_rotary 返回）+ full_kv。 if _rope_caps is not None: _qk_caps, _qkv_caps = _rope_caps else: @@ -115,23 +115,16 @@ def patched_forward( self.layer_number, _q_post, _k_post, self.config.num_layers, positions=_off_positions, ) - # pre-RoPE Q/K + V：从 get_qkv 返回统一取（与 ON 同源） + # full KV（post-RoPE K + V）：V 从 get_qkv 截（侵入式已在 megatron dump build_kv_input_v， + # 这里 full_kv 另存一份 post-K + V 供 attn_kv 对比） if _qkv_caps: - _pre_q, _pre_k, _pre_v = _qkv_caps[-1][:3] - # squeeze 到 [T, H, D]（与 ON 侧 squeeze 后 dump 一致） - _sq = lambda _t: _t.squeeze(1) if _t.dim() > 2 else _t - _pre_q, _pre_k, _pre_v = _sq(_pre_q), _sq(_pre_k), _sq(_pre_v) - dump_rope_preqk_verl080( - self.layer_number, _pre_q, _pre_k, - self.config.num_layers, - ) + _pre_v = _qkv_caps[-1][2] + if _pre_v.dim() > 2: + _pre_v = _pre_v.squeeze(1) dump_full_kv_off( self.layer_number, _k_post, _pre_v, self.config.num_layers, ) - else: - print(f"[PS-diag] OFF preqkv L{self.layer_number}: get_qkv 未截获; skip", - flush=True) else: print( f"[PS-diag] OFF rope_postqk L{self.layer_number}: " From 7aa39da40fb6a7a2d321c901edf0f8693de44726 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 12:05:41 +0800 Subject: [PATCH 34/61] =?UTF-8?q?[diag]=20OFF=20post-RoPE=20Q/K=20+=20full?= =?UTF-8?q?=5Fkv=20=E4=B9=9F=E4=BE=B5=E5=85=A5=E5=BC=8F=20dump=EF=BC=88rot?= =?UTF-8?q?ary=20block=20=E4=B9=8B=E5=90=8E=EF=BC=89?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 在 megatron attention.py rotary block 结束后（nvtx_range_pop 之后、core attention 之前） 侵入式 dump rope_postqk（Q/K post-RoPE）+ full_kv（post-K + V）。 此处 query/key 是 post-RoPE、value 是 raw，都在 scope——和 ON 侧（_apply_positioned_rope 返回后 dump postqk；build_kv 后 dump expanded_kv）对称。 OFF 侧全部诊断（preqk/build_kv_input_v/postqk/full_kv）现在都在 megatron 源码里 侵入式 dump，与 ON 同代码结构、同处理点。apply_rotary hook 保留但冗余（双写同值）。 --- .../megatron/core/transformer/attention.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py index 31a4f85c..70725a8a 100644 --- a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py +++ b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py @@ -1141,6 +1141,21 @@ def forward( # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb) nvtx_range_pop(suffix="rotary_pos_emb") + # [PS-diag] OFF post-RoPE Q/K + full_kv dump — 侵入式，rotary block 之后、core attention 之前， + # 与 ON 侧（_apply_positioned_rope 返回后 dump rope_postqk；build_kv 后 dump expanded_kv）对称。 + # 此处 query/key 是 post-RoPE，value 是 raw（未旋转），都在 scope。 + if _ps_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import ( + dump_rope_postqk_verl080, dump_full_kv_off, + ) + dump_rope_postqk_verl080(self.layer_number, query, key, + self.config.num_layers) + dump_full_kv_off(self.layer_number, key, value, + self.config.num_layers) + except Exception as _ps_e2: + print(f"[PS-diag] OFF postqk/full_kv dump failed: {_ps_e2}", flush=True) + # ================================== # core attention computation # ================================== From 5b95d704110449f7eb3346a6065450e9da724b8d Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 15:15:50 +0800 Subject: [PATCH 35/61] [fix] OFF dump: remove hook-based rope_postqk/full_kv writes (double-flush overwrite); cmp: build_kv_input_v ON vs OFF same-source - attention.py (patch): remove capture_rope_qk hook + associated dump_rope_postqk_verl080 / dump_full_kv_off calls from OFF path. These were duplicating the invasive dumps inside Megatron attention.py and overwriting the complete 24-layer files with 1-layer partial files on the last layer flush. - cmp_diag_verl080.py: cmp_build_kv_input_v now compares ON build_kv_input_v.pt vs OFF build_kv_input_v.pt (same-source, both raw V from get_qkv) instead of OFF full_kv.pt V entry. Co-Authored-By: Claude Fable 5 --- .../verl080_mcore0161_ms0160/attention.py | 84 +++++-------------- .../prefix_sharing/tools/cmp_diag_verl080.py | 12 +-- 2 files changed, 29 insertions(+), 67 deletions(-) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 778eb8d9..518c73c7 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -36,42 +36,32 @@ def patched_forward( ctx = current_prefix_sharing_context() if ctx is None: # ── normal path: 调用原始 forward ── - # diag: hook 截获 original_forward 内部 apply_rotary_pos_emb 的真实 post-RoPE - # Q/K（mcore Attention.forward THD prefill 每层调两次：先 Q 后 K）。post-RoPE Q/K - # 是 forward 内部中间变量，唯一能拿到真实张量的办法就是 hook 那个模块级 rotary 函数。 + # post-RoPE Q/K / full_kv / preqk 由 Megatron attention.py 侵入式 dump 写入； + # patch 层只负责 attn_outputs + rope_freqs（侵入式未覆盖的）。 import os as _os - from contextlib import nullcontext _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None - if _diag_on and rotary_pos_emb is not None: - from prefix_sharing.tools.diagnostic_dump_verl080 import capture_rope_qk - _rope_cm = capture_rope_qk(self) - else: - _rope_cm = nullcontext() - with _rope_cm as _rope_caps: - _result = original_forward( - self, - hidden_states, - attention_mask, - key_value_states=key_value_states, - inference_context=inference_context, - rotary_pos_emb=rotary_pos_emb, - rotary_pos_cos=rotary_pos_cos, - rotary_pos_sin=rotary_pos_sin, - rotary_pos_cos_sin=rotary_pos_cos_sin, - attention_bias=attention_bias, - packed_seq_params=packed_seq_params, - sequence_len_offset=sequence_len_offset, - inference_params=inference_params, - ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs + rope_postqk/preqk dump ##### + _result = original_forward( + self, + hidden_states, + attention_mask, + key_value_states=key_value_states, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + inference_params=inference_params, + ) + # ##### [PS-diag] OFF attn_outputs + rope_freqs dump ##### if _diag_on: import torch # 仅用于构造 positions（freqs 切 per-token + Q/K debug） from prefix_sharing.tools.diagnostic_dump import ( dump_attn_off, dump_rope_freqs, ) - from prefix_sharing.tools.diagnostic_dump_verl080 import ( - dump_rope_postqk_verl080, dump_rope_preqk_verl080, dump_full_kv_off, - ) + # rope_postqk / preqk / full_kv 已由 Megatron 侵入式 dump 覆盖 from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb _attn_out = _result[0] if isinstance(_result, tuple) else _result _bs = ( @@ -101,38 +91,10 @@ def patched_forward( _q_pos_emb.index_select(0, _off_positions), self.layer_number, self.config.num_layers, ) - # _rope_caps = (qk_caps[post-RoPE], qkv_caps[pre-RoPE Q/K/V]) - # pre-RoPE Q/K/V 已由 megatron attention.py 侵入式 dump（line 1070 之后）， - # 这里只处理 post-RoPE Q/K（apply_rotary 返回）+ full_kv。 - if _rope_caps is not None: - _qk_caps, _qkv_caps = _rope_caps - else: - _qk_caps, _qkv_caps = None, None - # post-RoPE Q/K（apply_rotary 返回） - if _qk_caps is not None and len(_qk_caps) >= 2: - _q_post, _k_post = _qk_caps[0]["post"], _qk_caps[1]["post"] - dump_rope_postqk_verl080( - self.layer_number, _q_post, _k_post, - self.config.num_layers, positions=_off_positions, - ) - # full KV（post-RoPE K + V）：V 从 get_qkv 截（侵入式已在 megatron dump build_kv_input_v， - # 这里 full_kv 另存一份 post-K + V 供 attn_kv 对比） - if _qkv_caps: - _pre_v = _qkv_caps[-1][2] - if _pre_v.dim() > 2: - _pre_v = _pre_v.squeeze(1) - dump_full_kv_off( - self.layer_number, _k_post, _pre_v, - self.config.num_layers, - ) - else: - print( - f"[PS-diag] OFF rope_postqk L{self.layer_number}: " - f"expected 2 captures (Q,K), got " - f"{len(_qk_caps) if _qk_caps is not None else 'None'}; skip", - flush=True, - ) - # ##### [PS-diag] OFF attn_outputs + rope_freqs + rope_postqk/preqk dump end ##### + # rope_postqk + full_kv 已由 megatron attention.py 侵入式 dump + # （rotary block 之后），不在 patch 层重复——避免 hook 二次 flush + # 覆盖侵入式已写好的完整 24 层文件。 + # ##### [PS-diag] OFF attn_outputs + rope_freqs dump end ##### return _result # ── prefix-sharing path ── diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index d19edba7..bb5e5123 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -932,13 +932,14 @@ def _print_attn_kv(r: CheckResult): def cmp_build_kv_input_v(dir_on: str, dir_off: str, layer: int | None = None) -> CheckResult | None: - """对比 ON build_kv 输入 V（build_kv 之前）vs OFF full_kv V（suffix 对齐）。 + """对比 ON build_kv_input_v vs OFF build_kv_input_v（suffix 对齐）。 - 定位 V 偏是在 build_kv 之前（get_qkv/hidden_states）还是 build_kv 引入。 + 两边都存 get_qkv 后、build_kv/RoPE 前的 raw V（``{layer: tensor}``）。同源对比， + 应逐元素相同——若不同则问题在 QKV 投影阶段（hidden_states / QKV 权重）。 ON_T vs OFF_T 还能看出 ON 有没有把 hidden_states 裁成 suffix-only。 """ fa = os.path.join(dir_on, "build_kv_input_v.pt") - fb = os.path.join(dir_off, "full_kv.pt") + fb = os.path.join(dir_off, "build_kv_input_v.pt") if not os.path.exists(fa) or not os.path.exists(fb): return None on_dict = torch.load(fa, weights_only=True) @@ -958,8 +959,7 @@ def cmp_build_kv_input_v(dir_on: str, dir_off: str, worst_cos = 1.0 for lyr in layers: on_v = on_dict[lyr] - off_entry = off_dict[lyr] - off_v = off_entry.get("value") if isinstance(off_entry, dict) else None + off_v = off_dict[lyr] if on_v is None or off_v is None: per_layer[lyr] = {"error": "缺失"}; continue on_f = on_v.reshape(on_v.shape[0], -1).float() @@ -985,7 +985,7 @@ def cmp_build_kv_input_v(dir_on: str, dir_off: str, def _print_build_kv_input_v(r: CheckResult): - print(_SEP_SINGLE + f"\n [{r.name}] ON build_kv 输入 V vs OFF full_kv V（suffix 对齐）") + print(_SEP_SINGLE + f"\n [{r.name}] ON build_kv_input_v vs OFF build_kv_input_v（suffix 对齐）") print(_SEP_SINGLE) m = r.metrics if "error" in m: From 91282bdc7ba940246dec6692a56d5e17548a48c3 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 16:27:46 +0800 Subject: [PATCH 36/61] =?UTF-8?q?[diag]=20add=20hidden=5Fstates=20dump=20+?= =?UTF-8?q?=20cmp=20at=20attention=20entrance=20(ON=20=E4=BE=B5=E5=85=A5?= =?UTF-8?q?=E5=BC=8F=20+=20OFF=20=E4=BE=B5=E5=85=A5=E5=BC=8F)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - diagnostic_dump_verl080.py: new dump_hidden_states_on + buffer - attention.py (ON): dump hidden_states alongside Q/K/V - Megatron attention.py (OFF 侵入式): same dump at same point - cmp_diag_verl080.py: cmp_hidden_states + print + shapes + main — per-layer suffix-aligned max_diff/cos; pinpoints whether V discrepancy originates in GEMM or upstream hidden_states Co-Authored-By: Claude Fable 5 --- .../megatron/core/transformer/attention.py | 4 +- .../verl080_mcore0161_ms0160/attention.py | 7 ++ .../prefix_sharing/tools/cmp_diag_verl080.py | 88 ++++++++++++++++++- .../tools/diagnostic_dump_verl080.py | 22 +++++ 4 files changed, 119 insertions(+), 2 deletions(-) diff --git a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py index 70725a8a..100cbfba 100644 --- a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py +++ b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py @@ -1075,12 +1075,14 @@ def forward( if _ps_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: try: from prefix_sharing.tools.diagnostic_dump_verl080 import ( - dump_rope_preqk_verl080, dump_build_kv_input_v_on, + dump_rope_preqk_verl080, dump_build_kv_input_v_on, dump_hidden_states_on, ) dump_rope_preqk_verl080(self.layer_number, query, key, self.config.num_layers) dump_build_kv_input_v_on(self.layer_number, value, self.config.num_layers) + dump_hidden_states_on(self.layer_number, hidden_states, + self.config.num_layers) except Exception as _ps_e: print(f"[PS-diag] OFF preqkv dump failed: {_ps_e}", flush=True) diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 518c73c7..d0d06865 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -130,6 +130,13 @@ def patched_forward( self.config.num_layers) except Exception as _e: print(f"build_kv_input_v (pre-RoPE V) dump failed: {_e}", flush=True) + # [PS-diag] dump hidden_states for input-level comparison + try: + from prefix_sharing.tools.diagnostic_dump_verl080 import dump_hidden_states_on + dump_hidden_states_on(self.layer_number, hidden_states, + self.config.num_layers) + except Exception as _e: + print(f"hidden_states dump failed: {_e}", flush=True) # delegate to verified integrations code from prefix_sharing.integrations.megatron_runtime import ( diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index bb5e5123..acbff5fa 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -1011,6 +1011,85 @@ def _print_build_kv_input_v(r: CheckResult): print() +def cmp_hidden_states(dir_on: str, dir_off: str, + layer: int | None = None) -> CheckResult | None: + """对比 ON vs OFF hidden_states（suffix 对齐，注意力层入口）。 + + 这是 QKV 投影的 INPUT。如果 hidden_states 一致但 V 不一致 → GEMM 精度差异； + 如果 hidden_states 就不一致 → 根因在上游（embedding / input_layernorm）。 + """ + fa = os.path.join(dir_on, "hidden_states.pt") + fb = os.path.join(dir_off, "hidden_states.pt") + if not os.path.exists(fa) or not os.path.exists(fb): + return None + on_dict = torch.load(fa, weights_only=True) + off_dict = torch.load(fb, weights_only=True) + if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): + return None + layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) + if layer is not None: + layers = [l for l in layers if l == layer] + _name = f"hidden_states_L{layer}" if layer is not None else "hidden_states" + if not layers: + return CheckResult(name=_name, passed=False, metrics={"error": "no layers"}) + + align_mask = _build_attn_align_mask(dir_on, dir_off) + per_layer: dict = {} + worst_md = 0.0 + worst_cos = 1.0 + for lyr in layers: + on_hs = on_dict[lyr] + off_hs = off_dict[lyr] + if on_hs is None or off_hs is None: + per_layer[lyr] = {"error": "缺失"}; continue + on_f = on_hs.reshape(on_hs.shape[0], -1).float() + off_f = off_hs.reshape(off_hs.shape[0], -1).float() + on_T, off_T = int(on_hs.shape[0]), int(off_hs.shape[0]) + if align_mask is not None and on_f.shape[0] != off_f.shape[0]: + try: + on_f, off_f = _align_packed(on_f, off_f, align_mask) + except ValueError as e: + per_layer[lyr] = {"error": str(e), "on_T": on_T, "off_T": off_T} + continue + diff = (on_f - off_f).abs() + cos = _cosine_sim(on_f, off_f, dim=-1) + md = float(diff.max()) + per_layer[lyr] = {"max_diff": md, "cos_avg": float(cos.mean()), + "cos_min": float(cos.min()), "n_tokens": on_f.shape[0], + "on_T": on_T, "off_T": off_T} + worst_md = max(worst_md, md) + worst_cos = min(worst_cos, float(cos.min())) + passed = worst_md < 1e-5 + return CheckResult(name=_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_md, "cos_min": worst_cos}) + + +def _print_hidden_states(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] ON vs OFF hidden_states（suffix 对齐，注意力入口）") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS':>10s} " + f"{'ON_T':>8s} {'OFF_T':>8s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") + for lyr in sorted(layers): + d = layers[lyr] + if "max_diff" not in d: + print(f" {lyr:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") + continue + md, cos = d["max_diff"], d["cos_avg"] + ok = md < 1e-5 + print(f" {lyr:>6d} {md:>12.3e} {cos:>10.6f} " + f"{d.get('on_T', '—'):>8} {d.get('off_T', '—'):>8} " + f"{'OK' if ok else 'DIFF':>8s}") + print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " + f"{_CHECK if r.passed else _CROSS} " + f"{'PASS（hidden_states 一致 → V 差异在 GEMM）' if r.passed else 'FAIL（hidden_states 不一致 → 根因上游）'}") + print() + + # ════════════════════════════════════════════════════════════════ # 2D mask loading # ════════════════════════════════════════════════════════════════ @@ -1164,6 +1243,7 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, "expanded_kv.pt", "full_kv.pt", "build_kv_input_v.pt", + "hidden_states.pt", "prefix_lens.pt", "cu_seqlens_q.pt", ] @@ -1540,12 +1620,18 @@ def main(): all_results.append(r) _print_attn_kv(r) - # ── packed: build_kv 输入 V（build_kv 前）vs OFF full_kv V —— 定位 V 偏在 build_kv 之前还是之后 ── + # ── packed: build_kv 输入 V — ON vs OFF 同源对比 ── r = cmp_build_kv_input_v(args.dir_on, args.dir_off, args.layer) if r: all_results.append(r) _print_build_kv_input_v(r) + # ── packed: hidden_states（注意力入口）— 隔离 QKV 投影 vs 上游 ── + r = cmp_hidden_states(args.dir_on, args.dir_off, args.layer) + if r: + all_results.append(r) + _print_hidden_states(r) + # ── packed: attention_output per-layer cos（RoPE 下游）── r = cmp_attn_layer(args.dir_on, args.dir_off, args.layer) if r: diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 2b1da1d1..d2eff422 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -362,6 +362,28 @@ def dump_build_kv_input_v_on(layer_number: int, value: torch.Tensor, _BUILD_KV_INPUT_V_BUFFER = None +_HIDDEN_STATES_BUFFER: dict[int, torch.Tensor] | None = None + + +def dump_hidden_states_on(layer_number: int, hidden_states: torch.Tensor, + num_layers: int) -> None: + """Accumulate hidden_states at attention entrance. Auto-flush ``hidden_states.pt`` on last layer. + + Used to verify whether ON/OFF hidden_states are bit-identical for suffix tokens. + Format: ``{layer_idx: [T, H]}``. + """ + global _HIDDEN_STATES_BUFFER + dump_dir = _get_dump_dir() + if dump_dir is None: + return + if _HIDDEN_STATES_BUFFER is None: + _HIDDEN_STATES_BUFFER = {} + _HIDDEN_STATES_BUFFER[layer_number] = hidden_states.detach().cpu().clone() + if layer_number == num_layers: + _flush_dict_buffer("hidden_states.pt", _HIDDEN_STATES_BUFFER, dump_dir) + _HIDDEN_STATES_BUFFER = None + + @contextlib.contextmanager def capture_rope_qk(attention_module): """Hook apply_rotary_pos_emb（post-RoPE Q/K）+ get_query_key_value_tensors（pre-RoPE Q/K/V）。 From 2d5270ec39bbdcb7059a5174f5cbb33528d85ef3 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 20:21:58 +0800 Subject: [PATCH 37/61] [cmp] add baseline scripts: cross_batch + within_batch GEMM precision - cmp_baseline_cross_batch.py: compare single-copy dump vs N-copy stacked dump (cross-batch-size baseline). Measures GEMM kernel precision difference from batch size change. - cmp_baseline_within_batch.py: pairwise compare N copies within a single dump (within-batch baseline). Verifies same-kernel bit-identical reproduction. - Both reuse CheckResult, _cosine_sim, _error_abs_rel, _pearson_r, _dump_json from cmp_diag_verl080. - Compare pre-RoPE quantities: hidden_states, rope_preqk Q/K, build_kv_input_v via direct cu_seqlens slicing. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 346 ++++++++++++++++++ .../tools/cmp_baseline_within_batch.py | 328 +++++++++++++++++ 2 files changed, 674 insertions(+) create mode 100644 prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py create mode 100644 prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py new file mode 100644 index 00000000..ebd01ffd --- /dev/null +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -0,0 +1,346 @@ +"""GEMM Precision Baseline — Cross-Batch-Size Comparison. + +Compares the SAME data processed at DIFFERENT batch sizes (1 copy vs N copies) +to quantify GEMM floating-point noise independent of prefix-sharing logic. + +Usage:: + + # Run 1: single copy + export PREFIX_SHARING_DIAG_DUMP=/dump_single + # forward with batch=[A] + + # Run 2: stacked copies + export PREFIX_SHARING_DIAG_DUMP=/dump_stacked + # forward with batch=[A x N] + + python cmp_baseline_cross_batch.py \ + --dir-single /dump_single --dir-stacked /dump_stacked --num-copies 4 +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +from dataclasses import dataclass, field + +import torch + +from prefix_sharing.tools.cmp_diag_verl080 import ( + CheckResult, + _SEP_DOUBLE, + _SEP_SINGLE, + _CHECK, + _CROSS, + _cosine_sim, + _error_abs_rel, + _pearson_r, + _dump_json, +) + + +# ══════════════════════════════════════════════════════════════════ +# Local helpers +# ══════════════════════════════════════════════════════════════════ + +def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: + """Load a per-layer ``{layer_idx: ...}`` dict from ``dir_path/filename``.""" + fp = os.path.join(dir_path, filename) + if not os.path.exists(fp): + return None + d = torch.load(fp, weights_only=True) + return d if isinstance(d, dict) else None + + +def _load_cu_seqlens(dir_path: str) -> torch.Tensor | None: + fp = os.path.join(dir_path, "cu_seqlens_q.pt") + if not os.path.exists(fp): + return None + return torch.load(fp, weights_only=True) + + +def _extract_copy(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: + """Slice copy ``copy_idx`` from packed: ``packed[cu[i] : cu[i+1]]``.""" + return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] + + +def _tensor_metrics(a: torch.Tensor, b: torch.Tensor) -> dict: + """Compute per-token cosine + global error for ``[T, ...]`` tensors. + + Unlike :func:`_vec_metrics` (which expects 1-D vectors), this handles + multi-token tensors by flattening all non-token dims. + """ + a_f = a.reshape(a.shape[0], -1).float() + b_f = b.reshape(b.shape[0], -1).float() + err = _error_abs_rel(a_f, b_f) + cos = _cosine_sim(a_f, b_f, dim=-1) + pr = _pearson_r(a_f, b_f) + return { + "max_abs": err["abs_max"], + "mean_abs": err["abs_mean"], + "rel_max": err["rel_max"], + "rel_mean": err["rel_mean"], + "cos_avg": float(cos.mean()), + "cos_min": float(cos.min()), + "pearson": pr, + "n_tokens": a_f.shape[0], + } + + +def _get_layers(data: dict) -> list[int]: + return sorted(int(k) for k in data.keys()) + + +# ══════════════════════════════════════════════════════════════════ +# Comparison drivers +# ══════════════════════════════════════════════════════════════════ + +def _compare_plain_file(dir_single: str, dir_stacked: str, filename: str, + cu_single: torch.Tensor, cu_stacked: torch.Tensor, + n_copies: int, layer: int | None, + label: str) -> CheckResult: + """Compare ``filename`` (``{layer: [T, ...]}``) across copies.""" + sd = _load_per_layer_dict(dir_single, filename) + md = _load_per_layer_dict(dir_stacked, filename) + if sd is None or md is None: + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing in one or both dirs"}) + + layers = _get_layers(sd) + if layer is not None: + layers = [l for l in layers if l == layer] + if not layers: + return CheckResult(name=label, passed=False, + metrics={"error": "no common layers"}) + + per_layer: dict = {} + worst_md = 0.0 + for lyr in layers: + st = sd[lyr].float() + mt = md[lyr].float() + T = st.shape[0] + copies: list[dict] = [] + copy_max = 0.0 + for i in range(n_copies): + ct = _extract_copy(mt, cu_stacked, i) + if ct.shape[0] != T: + copies.append({"copy": i, "error": + f"length mismatch: single={T} copy={ct.shape[0]}"}) + continue + m = _tensor_metrics(st, ct) + m["copy"] = i + copies.append(m) + copy_max = max(copy_max, m["max_abs"]) + per_layer[lyr] = {"copies": copies, "max_across_copies": copy_max} + worst_md = max(worst_md, copy_max) + + return CheckResult(name=label, passed=worst_md < 1e-5, + metrics={"layers": per_layer, "worst_max_abs": worst_md}) + + +def _compare_rope_preqk(dir_single: str, dir_stacked: str, + cu_single: torch.Tensor, cu_stacked: torch.Tensor, + n_copies: int, layer: int | None, + label: str) -> CheckResult: + """Compare ``rope_preqk.pt`` (``{layer: {"query", "key"}}``) across copies.""" + sd = _load_per_layer_dict(dir_single, "rope_preqk.pt") + md = _load_per_layer_dict(dir_stacked, "rope_preqk.pt") + if sd is None or md is None: + return CheckResult(name=label, passed=False, + metrics={"error": "rope_preqk.pt missing"}) + + layers = _get_layers(sd) + if layer is not None: + layers = [l for l in layers if l == layer] + if not layers: + return CheckResult(name=label, passed=False, metrics={"error": "no common layers"}) + + per_layer: dict = {} + worst_md = 0.0 + for lyr in layers: + sq = sd[lyr]["query"].float() + sk = sd[lyr]["key"].float() + mq = md[lyr]["query"].float() + mk = md[lyr]["key"].float() + Tq, Tk = sq.shape[0], sk.shape[0] + layer_worst = 0.0 + layer_copies: list[dict] = [] + for i in range(n_copies): + cq = _extract_copy(mq, cu_stacked, i) + ck = _extract_copy(mk, cu_stacked, i) + if cq.shape[0] != Tq or ck.shape[0] != Tk: + layer_copies.append({"copy": i, "error": + f"length mismatch Q: single={Tq} copy={cq.shape[0]}" + f" K: single={Tk} copy={ck.shape[0]}"}) + continue + qm = _tensor_metrics(sq, cq) + km = _tensor_metrics(sk, ck) + layer_copies.append({ + "copy": i, + "Q": qm, + "K": km, + }) + layer_worst = max(layer_worst, qm["max_abs"], km["max_abs"]) + per_layer[lyr] = {"copies": layer_copies, "max_across_copies": layer_worst} + worst_md = max(worst_md, layer_worst) + + return CheckResult(name=label, passed=worst_md < 1e-5, + metrics={"layers": per_layer, "worst_max_abs": worst_md}) + + +# ══════════════════════════════════════════════════════════════════ +# Output +# ══════════════════════════════════════════════════════════════════ + +def _print_header(dir_single: str, dir_stacked: str, n_copies: int, + cu_single: torch.Tensor, cu_stacked: torch.Tensor): + T = int(cu_single[-1]) + ok_single = cu_single.numel() == 2 # B=1 + ok_stacked = (cu_stacked.numel() == n_copies + 1 + and all(int(cu_stacked[i + 1]) - int(cu_stacked[i]) == T + for i in range(n_copies))) + print(_SEP_DOUBLE) + print(" GEMM Precision Baseline — Cross-Batch-Size Comparison") + print(f" Single : {dir_single} (1 sequence, {T} tokens)") + print(f" Stacked: {dir_stacked} ({n_copies} copies, {n_copies * T} tokens)") + print(f" Copies : {n_copies}") + print(_SEP_DOUBLE) + print(f" cu_seqlens single: {cu_single.tolist()}" + f" {' ' + _CHECK if ok_single else ' ' + _CROSS + ' expected B=1'}") + print(f" cu_seqlens stacked: {cu_stacked.tolist()}" + f" {' ' + _CHECK if ok_stacked else ' ' + _CROSS + ' expected B=' + str(n_copies)}") + print() + + +def _print_plain_table(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] Per-layer max_abs across copies") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + n_copies = max(len(v.get("copies", [])) for v in layers.values()) if layers else 0 + hdr = (f" {'LAYER':>6s} " + + " ".join(f"{'COPY_'+str(i):>11s}" for i in range(n_copies)) + + f" {'MAX':>11s} {'MEAN':>11s}") + print(hdr) + print(f" {'─' * 6} " + " ".join("─" * 11 for _ in range(n_copies + 2))) + for lyr in sorted(layers): + d = layers[lyr] + copies = d.get("copies", []) + vals = [c.get("max_abs", float("nan")) for c in copies] + mx = max(v for v in vals if not math.isnan(v)) if vals else float("nan") + mn = sum(v for v in vals if not math.isnan(v)) / max(1, sum(1 for v in vals if not math.isnan(v))) + row = f" {lyr:>6d} " + " ".join(f"{v:>11.3e}" for v in vals) + \ + f" {mx:>11.3e} {mn:>11.3e}" + print(row) + print() + + +def _print_rope_preqk_table(r: CheckResult, component: str): + print(_SEP_SINGLE + f"\n [{r.name}] {component} Per-layer max_abs across copies") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + n_copies = 0 + for v in layers.values(): + n_copies = max(n_copies, len(v.get("copies", []))) + if n_copies == 0: + print(" no data\n"); return + hdr = (f" {'LAYER':>6s} " + + " ".join(f"{'COPY_'+str(i):>11s}" for i in range(n_copies)) + + f" {'MAX':>11s} {'MEAN':>11s}") + print(hdr) + print(f" {'─' * 6} " + " ".join("─" * 11 for _ in range(n_copies + 2))) + for lyr in sorted(layers): + d = layers[lyr] + copies = d.get("copies", []) + vals = [c.get(component, {}).get("max_abs", float("nan")) for c in copies] + mx = max(v for v in vals if not math.isnan(v)) if vals else float("nan") + mn = sum(v for v in vals if not math.isnan(v)) / max(1, sum(1 for v in vals if not math.isnan(v))) + row = f" {lyr:>6d} " + " ".join(f"{v:>11.3e}" for v in vals) + \ + f" {mx:>11.3e} {mn:>11.3e}" + print(row) + print() + + +def _print_summary(all_results: list[CheckResult]): + print(_SEP_DOUBLE + "\n AGGREGATE SUMMARY (worst max_abs across all layers & copies)") + print(_SEP_DOUBLE) + hdr = f" {'FILE':<24s} {'WORST_MAX_ABS':>14s} {'PASS?':>8s}" + print(hdr + "\n " + "─" * (len(hdr) - 2)) + for r in all_results: + wm = r.metrics.get("worst_max_abs", "—") + wm_s = f"{wm:>14.3e}" if isinstance(wm, float) else f"{wm:>14s}" + s = f" {_CHECK} PASS" if r.passed else f" {_CROSS} FAIL" + print(f" {r.name:<24s} {wm_s} {s}") + print(_SEP_DOUBLE) + print() + + +# ══════════════════════════════════════════════════════════════════ +# Main +# ══════════════════════════════════════════════════════════════════ + +def main(): + ap = argparse.ArgumentParser( + description="GEMM precision baseline — cross-batch-size (single vs N copies)", + epilog=__doc__, + ) + ap.add_argument("--dir-single", required=True, + help="Single-copy dump directory (batch=[A])") + ap.add_argument("--dir-stacked", required=True, + help="Stacked-copies dump directory (batch=[A x N])") + ap.add_argument("--num-copies", type=int, required=True, + help="Number of stacked copies N") + ap.add_argument("--layer", type=int, default=None, + help="Compare specific layer 1-indexed (default: all)") + ap.add_argument("--output", "-o", default=None, + help="Write JSON report to this path") + args = ap.parse_args() + + # ── Load cu_seqlens & validate ── + cu_single = _load_cu_seqlens(args.dir_single) + cu_stacked = _load_cu_seqlens(args.dir_stacked) + if cu_single is None or cu_stacked is None: + print(f"{_CROSS} cu_seqlens_q.pt missing in one or both dump dirs") + return 1 + _print_header(args.dir_single, args.dir_stacked, args.num_copies, + cu_single, cu_stacked) + + all_results: list[CheckResult] = [] + + # ── hidden_states ── + r = _compare_plain_file(args.dir_single, args.dir_stacked, + "hidden_states.pt", cu_single, cu_stacked, + args.num_copies, args.layer, "hidden_states") + all_results.append(r) + _print_plain_table(r) + + # ── build_kv_input_v ── + r = _compare_plain_file(args.dir_single, args.dir_stacked, + "build_kv_input_v.pt", cu_single, cu_stacked, + args.num_copies, args.layer, "build_kv_input_v") + all_results.append(r) + _print_plain_table(r) + + # ── rope_preqk ── + r = _compare_rope_preqk(args.dir_single, args.dir_stacked, + cu_single, cu_stacked, + args.num_copies, args.layer, "rope_preqk") + all_results.append(r) + _print_rope_preqk_table(r, "Q") + _print_rope_preqk_table(r, "K") + + _print_summary(all_results) + + if args.output: + _dump_json(all_results, args.output, args.dir_single, args.dir_stacked, + tag=f"cross_batch_N{args.num_copies}", dir_off2=None) + + +if __name__ == "__main__": + main() diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py new file mode 100644 index 00000000..d8bac901 --- /dev/null +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -0,0 +1,328 @@ +"""GEMM Precision Baseline — Within-Batch Pairwise Comparison. + +Compares N identical copies WITHIN a single forward pass to verify that +the same GEMM kernel produces bit-identical results for the same data. + +Expected result: max_abs == 0.0 for all pairs (same batch size → same kernel). +Non-zero results indicate non-determinism beyond batch-size effects. + +Usage:: + + export PREFIX_SHARING_DIAG_DUMP=/dump_multi + # forward with batch=[A x N] + + python cmp_baseline_within_batch.py --dir-multi /dump_multi --num-copies 4 +""" + +from __future__ import annotations + +import argparse +import math +import os +from dataclasses import dataclass, field + +import torch + +from prefix_sharing.tools.cmp_diag_verl080 import ( + CheckResult, + _SEP_DOUBLE, + _SEP_SINGLE, + _CHECK, + _CROSS, + _cosine_sim, + _error_abs_rel, + _pearson_r, + _dump_json, +) + + +# ══════════════════════════════════════════════════════════════════ +# Local helpers (same as cross_batch) +# ══════════════════════════════════════════════════════════════════ + +def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: + fp = os.path.join(dir_path, filename) + if not os.path.exists(fp): + return None + d = torch.load(fp, weights_only=True) + return d if isinstance(d, dict) else None + + +def _load_cu_seqlens(dir_path: str) -> torch.Tensor | None: + fp = os.path.join(dir_path, "cu_seqlens_q.pt") + if not os.path.exists(fp): + return None + return torch.load(fp, weights_only=True) + + +def _extract_copy(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: + return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] + + +def _tensor_metrics(a: torch.Tensor, b: torch.Tensor) -> dict: + a_f = a.reshape(a.shape[0], -1).float() + b_f = b.reshape(b.shape[0], -1).float() + err = _error_abs_rel(a_f, b_f) + cos = _cosine_sim(a_f, b_f, dim=-1) + pr = _pearson_r(a_f, b_f) + return { + "max_abs": err["abs_max"], + "mean_abs": err["abs_mean"], + "rel_max": err["rel_max"], + "rel_mean": err["rel_mean"], + "cos_avg": float(cos.mean()), + "cos_min": float(cos.min()), + "pearson": pr, + "n_tokens": a_f.shape[0], + } + + +def _get_layers(data: dict) -> list[int]: + return sorted(int(k) for k in data.keys()) + + +# ══════════════════════════════════════════════════════════════════ +# Comparison drivers +# ══════════════════════════════════════════════════════════════════ + +def _compare_plain_file_within(dir_multi: str, filename: str, + cu: torch.Tensor, n_copies: int, + layer: int | None, + label: str) -> CheckResult: + """Pairwise comparison of copies within a single dump for ``filename``.""" + d = _load_per_layer_dict(dir_multi, filename) + if d is None: + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing"}) + + layers = _get_layers(d) + if layer is not None: + layers = [l for l in layers if l == layer] + if not layers: + return CheckResult(name=label, passed=False, + metrics={"error": "no layers"}) + + n_pairs = n_copies * (n_copies - 1) // 2 + per_layer: dict = {} + worst_md = 0.0 + worst_pair: tuple | None = None + worst_layer: int | None = None + + for lyr in layers: + mt = d[lyr].float() + copies = [_extract_copy(mt, cu, i) for i in range(n_copies)] + T = copies[0].shape[0] + layer_worst = 0.0 + layer_pair: tuple | None = None + for i in range(n_copies): + for j in range(i + 1, n_copies): + m = _tensor_metrics(copies[i], copies[j]) + if m["max_abs"] > layer_worst: + layer_worst = m["max_abs"] + layer_pair = (i, j) + per_layer[lyr] = { + "max_abs": layer_worst, + "worst_pair": list(layer_pair) if layer_pair else None, + "n_pairs": n_pairs, + "n_tokens": T, + } + if layer_worst > worst_md: + worst_md = layer_worst + worst_pair = layer_pair + worst_layer = lyr + + return CheckResult( + name=label, + passed=worst_md == 0.0, + metrics={ + "layers": per_layer, + "worst_max_abs": worst_md, + "worst_layer": worst_layer, + "worst_pair": worst_pair, + "n_pairs": n_pairs, + }, + ) + + +def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, + n_copies: int, layer: int | None, + label: str) -> CheckResult: + """Pairwise comparison of rope_preqk Q and K within a single dump.""" + d = _load_per_layer_dict(dir_multi, "rope_preqk.pt") + if d is None: + return CheckResult(name=label, passed=False, + metrics={"error": "rope_preqk.pt missing"}) + + layers = _get_layers(d) + if layer is not None: + layers = [l for l in layers if l == layer] + if not layers: + return CheckResult(name=label, passed=False, metrics={"error": "no layers"}) + + n_pairs = n_copies * (n_copies - 1) // 2 + per_layer: dict = {} + worst_md = 0.0 + worst_pair: tuple | None = None + worst_layer: int | None = None + + for lyr in layers: + mq = d[lyr]["query"].float() + mk = d[lyr]["key"].float() + q_copies = [_extract_copy(mq, cu, i) for i in range(n_copies)] + k_copies = [_extract_copy(mk, cu, i) for i in range(n_copies)] + layer_worst = 0.0 + layer_pair: tuple | None = None + for i in range(n_copies): + for j in range(i + 1, n_copies): + qm = _tensor_metrics(q_copies[i], q_copies[j]) + km = _tensor_metrics(k_copies[i], k_copies[j]) + w = max(qm["max_abs"], km["max_abs"]) + if w > layer_worst: + layer_worst = w + layer_pair = (i, j) + per_layer[lyr] = { + "max_abs": layer_worst, + "worst_pair": list(layer_pair) if layer_pair else None, + "n_pairs": n_pairs, + } + if layer_worst > worst_md: + worst_md = layer_worst + worst_pair = layer_pair + worst_layer = lyr + + return CheckResult( + name=label, + passed=worst_md == 0.0, + metrics={ + "layers": per_layer, + "worst_max_abs": worst_md, + "worst_layer": worst_layer, + "worst_pair": worst_pair, + "n_pairs": n_pairs, + }, + ) + + +# ══════════════════════════════════════════════════════════════════ +# Output +# ══════════════════════════════════════════════════════════════════ + +def _print_header(dir_multi: str, cu: torch.Tensor, n_copies: int): + n_pair = n_copies * (n_copies - 1) // 2 + T = int(cu[1]) - int(cu[0]) + ok = (cu.numel() == n_copies + 1 + and all(int(cu[i + 1]) - int(cu[i]) == T for i in range(n_copies))) + print(_SEP_DOUBLE) + print(" GEMM Precision Baseline — Within-Batch Pairwise Comparison") + print(f" Directory: {dir_multi}") + print(f" Copies: {n_copies} → {n_pair} pair(s)") + print(f" Tokens per copy: {T}") + print(_SEP_DOUBLE) + print(f" cu_seqlens: {cu.tolist()}" + f" {' ' + _CHECK if ok else ' ' + _CROSS + ' copies not uniform'}") + print() + + +def _print_within_plain_table(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] Max abs error across all C(N,2) pairs") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + n_pairs = m.get("n_pairs", "—") + print(f" {'LAYER':>6s} {'PAIRS':>6s} {'MAX_ABS':>12s} {'WORST_PAIR':>12s} " + f"{'COS_MIN':>10s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 10} {'─' * 8}") + for lyr in sorted(layers): + d = layers[lyr] + md = d["max_abs"] + wp = d.get("worst_pair", "—") + wp_s = str(wp) if wp else "—" + ok = md == 0.0 + # cos_min not tracked per-layer in simple mode; use "—" + print(f" {lyr:>6d} {n_pairs:>6} {md:>12.3e} {wp_s:>12} " + f"{'—':>10} {'PASS' if ok else 'DIFF':>8s}") + print() + + +def _print_within_verdict(r: CheckResult): + print(_SEP_DOUBLE) + m = r.metrics + if m.get("worst_max_abs", 1.0) == 0.0: + print(f" RESULT: ALL max_abs == 0.0 {_CHECK}") + print(" → No within-batch non-determinism detected.") + print(" → Any non-zero diff in cross-batch baseline is purely from") + print(" batch-size-induced GEMM kernel selection.") + else: + print(f" RESULT: max_abs = {m.get('worst_max_abs', '?'):.3e} {_CROSS}") + print(f" → Non-determinism detected!") + print(f" Layer: {m.get('worst_layer', '?')}") + print(f" Pair: {m.get('worst_pair', '?')}") + print(" → Check: dropout disabled? model.eval()? non-deterministic CUDA?") + print(_SEP_DOUBLE) + print() + + +# ══════════════════════════════════════════════════════════════════ +# Main +# ══════════════════════════════════════════════════════════════════ + +def main(): + ap = argparse.ArgumentParser( + description="GEMM precision baseline — within-batch pairwise comparison", + epilog=__doc__, + ) + ap.add_argument("--dir-multi", required=True, + help="Multi-copy dump directory (batch=[A x N])") + ap.add_argument("--num-copies", type=int, required=True, + help="Number of copies N") + ap.add_argument("--layer", type=int, default=None, + help="Compare specific layer 1-indexed (default: all)") + ap.add_argument("--output", "-o", default=None, + help="Write JSON report to this path") + args = ap.parse_args() + + # ── Load cu_seqlens & validate ── + cu = _load_cu_seqlens(args.dir_multi) + if cu is None: + print(f"{_CROSS} cu_seqlens_q.pt missing") + return 1 + _print_header(args.dir_multi, cu, args.num_copies) + + all_results: list[CheckResult] = [] + + # ── hidden_states ── + r = _compare_plain_file_within(args.dir_multi, "hidden_states.pt", + cu, args.num_copies, args.layer, + "hidden_states") + all_results.append(r) + _print_within_plain_table(r) + + # ── build_kv_input_v ── + r = _compare_plain_file_within(args.dir_multi, "build_kv_input_v.pt", + cu, args.num_copies, args.layer, + "build_kv_input_v") + all_results.append(r) + _print_within_plain_table(r) + + # ── rope_preqk ── + r = _compare_rope_preqk_within(args.dir_multi, cu, args.num_copies, + args.layer, "rope_preqk") + all_results.append(r) + _print_within_plain_table(r) + + # ── verdict ── + all_passed = all(r.passed for r in all_results) + worst_md = max(r.metrics.get("worst_max_abs", 0) for r in all_results) + combined = CheckResult(name="WITHIN_BATCH_OVERALL", passed=all_passed, + metrics={"worst_max_abs": worst_md}) + _print_within_verdict(combined) + + if args.output: + _dump_json(all_results, args.output, "—", args.dir_multi, + tag=f"within_batch_N{args.num_copies}", dir_off2=None) + + +if __name__ == "__main__": + main() From e261dcaf4fd081071466d1632932db31933ccaa3 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 20:26:22 +0800 Subject: [PATCH 38/61] [cmp] baseline: add rope_freqs, rope_postqk, attn_outputs to both scripts - cross_batch: now compares 7 quantities per layer (hidden_states, build_kv_input_v, rope_preqk Q/K, rope_freqs, rope_postqk Q/K, attn_outputs) - within_batch: same 7 quantities, pairwise - _compare_rope_preqk accepts fname param for rope_preqk / rope_postqk Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 33 ++++++++++++++++--- .../tools/cmp_baseline_within_batch.py | 29 +++++++++++++--- 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index ebd01ffd..290d372c 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -142,13 +142,13 @@ def _compare_plain_file(dir_single: str, dir_stacked: str, filename: str, def _compare_rope_preqk(dir_single: str, dir_stacked: str, cu_single: torch.Tensor, cu_stacked: torch.Tensor, n_copies: int, layer: int | None, - label: str) -> CheckResult: - """Compare ``rope_preqk.pt`` (``{layer: {"query", "key"}}``) across copies.""" - sd = _load_per_layer_dict(dir_single, "rope_preqk.pt") - md = _load_per_layer_dict(dir_stacked, "rope_preqk.pt") + label: str, fname: str = "rope_preqk.pt") -> CheckResult: + """Compare ``{layer: {"query", "key"}}`` dict file across copies.""" + sd = _load_per_layer_dict(dir_single, fname) + md = _load_per_layer_dict(dir_stacked, fname) if sd is None or md is None: return CheckResult(name=label, passed=False, - metrics={"error": "rope_preqk.pt missing"}) + metrics={"error": f"{fname} missing"}) layers = _get_layers(sd) if layer is not None: @@ -335,6 +335,29 @@ def main(): _print_rope_preqk_table(r, "Q") _print_rope_preqk_table(r, "K") + # ── rope_freqs ── + r = _compare_plain_file(args.dir_single, args.dir_stacked, + "rope_freqs.pt", cu_single, cu_stacked, + args.num_copies, args.layer, "rope_freqs") + all_results.append(r) + _print_plain_table(r) + + # ── rope_postqk ── + r = _compare_rope_preqk(args.dir_single, args.dir_stacked, + cu_single, cu_stacked, + args.num_copies, args.layer, "rope_postqk", + fname="rope_postqk.pt") + all_results.append(r) + _print_rope_preqk_table(r, "Q") + _print_rope_preqk_table(r, "K") + + # ── attn_outputs ── + r = _compare_plain_file(args.dir_single, args.dir_stacked, + "attn_outputs.pt", cu_single, cu_stacked, + args.num_copies, args.layer, "attn_outputs") + all_results.append(r) + _print_plain_table(r) + _print_summary(all_results) if args.output: diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index d8bac901..8db49de5 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -146,12 +146,12 @@ def _compare_plain_file_within(dir_multi: str, filename: str, def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, n_copies: int, layer: int | None, - label: str) -> CheckResult: - """Pairwise comparison of rope_preqk Q and K within a single dump.""" - d = _load_per_layer_dict(dir_multi, "rope_preqk.pt") + label: str, fname: str = "rope_preqk.pt") -> CheckResult: + """Pairwise comparison of ``{layer: {"query","key"}}`` dict within a single dump.""" + d = _load_per_layer_dict(dir_multi, fname) if d is None: return CheckResult(name=label, passed=False, - metrics={"error": "rope_preqk.pt missing"}) + metrics={"error": f"{fname} missing"}) layers = _get_layers(d) if layer is not None: @@ -312,6 +312,27 @@ def main(): all_results.append(r) _print_within_plain_table(r) + # ── rope_freqs ── + r = _compare_plain_file_within(args.dir_multi, "rope_freqs.pt", + cu, args.num_copies, args.layer, + "rope_freqs") + all_results.append(r) + _print_within_plain_table(r) + + # ── rope_postqk ── + r = _compare_rope_preqk_within(args.dir_multi, cu, args.num_copies, + args.layer, "rope_postqk", + fname="rope_postqk.pt") + all_results.append(r) + _print_within_plain_table(r) + + # ── attn_outputs ── + r = _compare_plain_file_within(args.dir_multi, "attn_outputs.pt", + cu, args.num_copies, args.layer, + "attn_outputs") + all_results.append(r) + _print_within_plain_table(r) + # ── verdict ── all_passed = all(r.passed for r in all_results) worst_md = max(r.metrics.get("worst_max_abs", 0) for r in all_results) From a429b5f5882d709d36051521f753164ab2c1d13c Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 20:31:03 +0800 Subject: [PATCH 39/61] [cmp] baseline: add full_kv, logits, logprobs, entropy to both scripts Cross-batch: now compares 11 quantities: per-layer: hidden_states, build_kv_input_v, rope_preqk Q/K, rope_freqs, rope_postqk Q/K, attn_outputs, full_kv K/V single: logits 2D: logprobs_{tag}, entropy_{tag} Within-batch: same 11 quantities, pairwise within one dump. Added --tag param (default 'old') for logprobs/entropy filenames. Generic _compare_per_layer_kv handles any {layer: {a, b}} dict. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 153 +++++++++++++++--- .../tools/cmp_baseline_within_batch.py | 117 +++++++++++++- 2 files changed, 243 insertions(+), 27 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 290d372c..8972d407 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -143,7 +143,16 @@ def _compare_rope_preqk(dir_single: str, dir_stacked: str, cu_single: torch.Tensor, cu_stacked: torch.Tensor, n_copies: int, layer: int | None, label: str, fname: str = "rope_preqk.pt") -> CheckResult: - """Compare ``{layer: {"query", "key"}}`` dict file across copies.""" + return _compare_per_layer_kv(dir_single, dir_stacked, cu_single, cu_stacked, + n_copies, layer, label, fname, "query", "key") + + +def _compare_per_layer_kv(dir_single: str, dir_stacked: str, + cu_single: torch.Tensor, cu_stacked: torch.Tensor, + n_copies: int, layer: int | None, + label: str, fname: str, + field_a: str, field_b: str) -> CheckResult: + """Compare ``{layer: {field_a, field_b}}`` dict file across copies.""" sd = _load_per_layer_dict(dir_single, fname) md = _load_per_layer_dict(dir_stacked, fname) if sd is None or md is None: @@ -159,29 +168,25 @@ def _compare_rope_preqk(dir_single: str, dir_stacked: str, per_layer: dict = {} worst_md = 0.0 for lyr in layers: - sq = sd[lyr]["query"].float() - sk = sd[lyr]["key"].float() - mq = md[lyr]["query"].float() - mk = md[lyr]["key"].float() - Tq, Tk = sq.shape[0], sk.shape[0] + sa = sd[lyr][field_a].float() + sb = sd[lyr][field_b].float() + ma = md[lyr][field_a].float() + mb = md[lyr][field_b].float() + Ta, Tb = sa.shape[0], sb.shape[0] layer_worst = 0.0 layer_copies: list[dict] = [] for i in range(n_copies): - cq = _extract_copy(mq, cu_stacked, i) - ck = _extract_copy(mk, cu_stacked, i) - if cq.shape[0] != Tq or ck.shape[0] != Tk: + ca = _extract_copy(ma, cu_stacked, i) + cb = _extract_copy(mb, cu_stacked, i) + if ca.shape[0] != Ta or cb.shape[0] != Tb: layer_copies.append({"copy": i, "error": - f"length mismatch Q: single={Tq} copy={cq.shape[0]}" - f" K: single={Tk} copy={ck.shape[0]}"}) + f"length mismatch {field_a}: single={Ta} copy={ca.shape[0]}" + f" {field_b}: single={Tb} copy={cb.shape[0]}"}) continue - qm = _tensor_metrics(sq, cq) - km = _tensor_metrics(sk, ck) - layer_copies.append({ - "copy": i, - "Q": qm, - "K": km, - }) - layer_worst = max(layer_worst, qm["max_abs"], km["max_abs"]) + am = _tensor_metrics(sa, ca) + bm = _tensor_metrics(sb, cb) + layer_copies.append({"copy": i, field_a: am, field_b: bm}) + layer_worst = max(layer_worst, am["max_abs"], bm["max_abs"]) per_layer[lyr] = {"copies": layer_copies, "max_across_copies": layer_worst} worst_md = max(worst_md, layer_worst) @@ -189,6 +194,69 @@ def _compare_rope_preqk(dir_single: str, dir_stacked: str, metrics={"layers": per_layer, "worst_max_abs": worst_md}) +def _compare_single_tensor(dir_single: str, dir_stacked: str, + filename: str, cu_single: torch.Tensor, + cu_stacked: torch.Tensor, + n_copies: int, label: str) -> CheckResult: + """Compare a single packed tensor (e.g. logits.pt) across copies.""" + st = _load_per_layer_dict(dir_single, filename) # not actually a dict + # Actually single-tensor files are not dicts; use _load_tensor pattern + fp_s = os.path.join(dir_single, filename) + fp_m = os.path.join(dir_stacked, filename) + if not os.path.exists(fp_s) or not os.path.exists(fp_m): + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing in one or both dirs"}) + single = torch.load(fp_s, weights_only=True).float() + multi = torch.load(fp_m, weights_only=True).float() + T = single.shape[0] + copies: list[dict] = [] + worst_md = 0.0 + for i in range(n_copies): + ct = _extract_copy(multi, cu_stacked, i) + if ct.shape[0] != T: + copies.append({"copy": i, "error": + f"length mismatch: single={T} copy={ct.shape[0]}"}) + continue + m = _tensor_metrics(single, ct) + m["copy"] = i + copies.append(m) + worst_md = max(worst_md, m["max_abs"]) + return CheckResult(name=label, passed=worst_md < 1e-5, + metrics={"copies": copies, "worst_max_abs": worst_md}) + + +def _compare_2d_tensor(dir_single: str, dir_stacked: str, + filename: str, label: str) -> CheckResult: + """Compare 2D [B, L_max] tensor (logprobs/entropy) row by row.""" + fp_s = os.path.join(dir_single, filename) + fp_m = os.path.join(dir_stacked, filename) + if not os.path.exists(fp_s) or not os.path.exists(fp_m): + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing"}) + single = torch.load(fp_s, weights_only=True).float() # [1, L_max] + multi = torch.load(fp_m, weights_only=True).float() # [N, L_max] + if single.dim() < 2 or multi.dim() < 2: + return CheckResult(name=label, passed=False, + metrics={"error": "not 2D"}) + B = multi.shape[0] + copies: list[dict] = [] + worst_md = 0.0 + for i in range(B): + # single is [1, L_max], copy i is multi[i, :] → [L_max] + m = _tensor_metrics(single.reshape(-1), multi[i].reshape(-1)) + m["copy"] = i + copies.append(m) + worst_md = max(worst_md, m["max_abs"]) + if single.shape[0] > 1: + # single also has multiple rows → compare all pairs + for i in range(single.shape[0]): + m = _tensor_metrics(single[i].reshape(-1), multi[i].reshape(-1)) + copies[i] = {**copies[i], **m} + worst_md = max(worst_md, m["max_abs"]) + return CheckResult(name=label, passed=worst_md < 1e-5, + metrics={"copies": copies, "worst_max_abs": worst_md}) + + # ══════════════════════════════════════════════════════════════════ # Output # ══════════════════════════════════════════════════════════════════ @@ -298,6 +366,8 @@ def main(): help="Number of stacked copies N") ap.add_argument("--layer", type=int, default=None, help="Compare specific layer 1-indexed (default: all)") + ap.add_argument("--tag", default="old", + help="2D file tag for logprobs/entropy (default: old)") ap.add_argument("--output", "-o", default=None, help="Write JSON report to this path") args = ap.parse_args() @@ -358,6 +428,51 @@ def main(): all_results.append(r) _print_plain_table(r) + # ── full_kv (key + value) ── + r = _compare_per_layer_kv(args.dir_single, args.dir_stacked, + cu_single, cu_stacked, + args.num_copies, args.layer, "full_kv", + fname="full_kv.pt", field_a="key", field_b="value") + all_results.append(r) + _print_rope_preqk_table(r, "key") + _print_rope_preqk_table(r, "value") + + # ── logits (single packed tensor) ── + r = _compare_single_tensor(args.dir_single, args.dir_stacked, + "logits.pt", cu_single, cu_stacked, + args.num_copies, "logits") + if r.metrics.get("error") is None: + all_results.append(r) + _print_plain_table(r) + + # ── logprobs (2D) ── + r = _compare_2d_tensor(args.dir_single, args.dir_stacked, + f"logprobs_{args.tag}.pt", f"logprobs_{args.tag}") + if r.metrics.get("error") is None: + all_results.append(r) + print(_SEP_SINGLE + f"\n [{r.name}] 2D per-row comparison") + print(_SEP_SINGLE) + m = r.metrics + copies = m.get("copies", []) + vals = [c.get("max_abs", float("nan")) for c in copies] + print(f" copies: {len(copies)}, max_abs: {max(v for v in vals if not math.isnan(v)):.3e}" + if vals else " no data") + print() + + # ── entropy (2D) ── + r = _compare_2d_tensor(args.dir_single, args.dir_stacked, + f"entropy_{args.tag}.pt", f"entropy_{args.tag}") + if r.metrics.get("error") is None: + all_results.append(r) + print(_SEP_SINGLE + f"\n [{r.name}] 2D per-row comparison") + print(_SEP_SINGLE) + m = r.metrics + copies = m.get("copies", []) + vals = [c.get("max_abs", float("nan")) for c in copies] + print(f" copies: {len(copies)}, max_abs: {max(v for v in vals if not math.isnan(v)):.3e}" + if vals else " no data") + print() + _print_summary(all_results) if args.output: diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 8db49de5..651a69ad 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -147,7 +147,15 @@ def _compare_plain_file_within(dir_multi: str, filename: str, def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, n_copies: int, layer: int | None, label: str, fname: str = "rope_preqk.pt") -> CheckResult: - """Pairwise comparison of ``{layer: {"query","key"}}`` dict within a single dump.""" + return _compare_per_layer_kv_within(dir_multi, cu, n_copies, layer, label, + fname, "query", "key") + + +def _compare_per_layer_kv_within(dir_multi: str, cu: torch.Tensor, + n_copies: int, layer: int | None, + label: str, fname: str, + field_a: str, field_b: str) -> CheckResult: + """Pairwise comparison of ``{layer: {field_a, field_b}}`` dict within a single dump.""" d = _load_per_layer_dict(dir_multi, fname) if d is None: return CheckResult(name=label, passed=False, @@ -166,17 +174,17 @@ def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, worst_layer: int | None = None for lyr in layers: - mq = d[lyr]["query"].float() - mk = d[lyr]["key"].float() - q_copies = [_extract_copy(mq, cu, i) for i in range(n_copies)] - k_copies = [_extract_copy(mk, cu, i) for i in range(n_copies)] + ma = d[lyr][field_a].float() + mb = d[lyr][field_b].float() + a_copies = [_extract_copy(ma, cu, i) for i in range(n_copies)] + b_copies = [_extract_copy(mb, cu, i) for i in range(n_copies)] layer_worst = 0.0 layer_pair: tuple | None = None for i in range(n_copies): for j in range(i + 1, n_copies): - qm = _tensor_metrics(q_copies[i], q_copies[j]) - km = _tensor_metrics(k_copies[i], k_copies[j]) - w = max(qm["max_abs"], km["max_abs"]) + am = _tensor_metrics(a_copies[i], a_copies[j]) + bm = _tensor_metrics(b_copies[i], b_copies[j]) + w = max(am["max_abs"], bm["max_abs"]) if w > layer_worst: layer_worst = w layer_pair = (i, j) @@ -203,6 +211,56 @@ def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, ) +def _compare_single_tensor_within(dir_multi: str, filename: str, + cu: torch.Tensor, n_copies: int, + label: str) -> CheckResult: + """Pairwise compare a single packed tensor (e.g. logits.pt) within a dump.""" + fp = os.path.join(dir_multi, filename) + if not os.path.exists(fp): + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing"}) + mt = torch.load(fp, weights_only=True).float() + copies = [_extract_copy(mt, cu, i) for i in range(n_copies)] + n_pairs = n_copies * (n_copies - 1) // 2 + worst_md = 0.0 + worst_pair = None + for i in range(n_copies): + for j in range(i + 1, n_copies): + m = _tensor_metrics(copies[i], copies[j]) + if m["max_abs"] > worst_md: + worst_md = m["max_abs"] + worst_pair = (i, j) + return CheckResult(name=label, passed=worst_md == 0.0, + metrics={"worst_max_abs": worst_md, + "worst_pair": worst_pair, "n_pairs": n_pairs}) + + +def _compare_2d_tensor_within(dir_multi: str, filename: str, + label: str) -> CheckResult: + """Pairwise compare 2D [B, L_max] rows within a dump.""" + fp = os.path.join(dir_multi, filename) + if not os.path.exists(fp): + return CheckResult(name=label, passed=False, + metrics={"error": f"{filename} missing"}) + mt = torch.load(fp, weights_only=True).float() # [B, L_max] + if mt.dim() < 2: + return CheckResult(name=label, passed=False, + metrics={"error": "not 2D"}) + B = mt.shape[0] + n_pairs = B * (B - 1) // 2 + worst_md = 0.0 + worst_pair = None + for i in range(B): + for j in range(i + 1, B): + m = _tensor_metrics(mt[i].reshape(-1), mt[j].reshape(-1)) + if m["max_abs"] > worst_md: + worst_md = m["max_abs"] + worst_pair = (i, j) + return CheckResult(name=label, passed=worst_md == 0.0, + metrics={"worst_max_abs": worst_md, + "worst_pair": worst_pair, "n_pairs": n_pairs}) + + # ══════════════════════════════════════════════════════════════════ # Output # ══════════════════════════════════════════════════════════════════ @@ -279,6 +337,8 @@ def main(): help="Number of copies N") ap.add_argument("--layer", type=int, default=None, help="Compare specific layer 1-indexed (default: all)") + ap.add_argument("--tag", default="old", + help="2D file tag for logprobs/entropy (default: old)") ap.add_argument("--output", "-o", default=None, help="Write JSON report to this path") args = ap.parse_args() @@ -333,6 +393,47 @@ def main(): all_results.append(r) _print_within_plain_table(r) + # ── full_kv (key + value) ── + r = _compare_per_layer_kv_within(args.dir_multi, cu, args.num_copies, + args.layer, "full_kv", + fname="full_kv.pt", + field_a="key", field_b="value") + all_results.append(r) + _print_within_plain_table(r) + + # ── logits (single packed tensor) ── + r = _compare_single_tensor_within(args.dir_multi, "logits.pt", + cu, args.num_copies, "logits") + if r.metrics.get("error") is None: + all_results.append(r) + _print_within_plain_table(r) + + # ── logprobs (2D) ── + r = _compare_2d_tensor_within(args.dir_multi, + f"logprobs_{args.tag}.pt", + f"logprobs_{args.tag}") + if r.metrics.get("error") is None: + all_results.append(r) + print(_SEP_SINGLE + f"\n [{r.name}] 2D pairwise comparison") + print(_SEP_SINGLE) + print(f" max_abs: {r.metrics.get('worst_max_abs', '—'):.3e}" + if isinstance(r.metrics.get("worst_max_abs"), float) + else f" {r.metrics.get('error', '—')}") + print() + + # ── entropy (2D) ── + r = _compare_2d_tensor_within(args.dir_multi, + f"entropy_{args.tag}.pt", + f"entropy_{args.tag}") + if r.metrics.get("error") is None: + all_results.append(r) + print(_SEP_SINGLE + f"\n [{r.name}] 2D pairwise comparison") + print(_SEP_SINGLE) + print(f" max_abs: {r.metrics.get('worst_max_abs', '—'):.3e}" + if isinstance(r.metrics.get("worst_max_abs"), float) + else f" {r.metrics.get('error', '—')}") + print() + # ── verdict ── all_passed = all(r.passed for r in all_results) worst_md = max(r.metrics.get("worst_max_abs", 0) for r in all_results) From 18d99c140bb8ef2106768b53e9221d06a2949975 Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 20:53:56 +0800 Subject: [PATCH 40/61] [inject] add baseline synthetic data injection with shuffle + stack MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - inject_baseline_synthetic.py: creates independent sequences from JSON, deterministic shuffle per sequence (seed + seq_index), stacks each sequence 'stack' times for batch-size-controlled baseline. - Env vars: PREFIX_SHARING_BASELINE_SYNTHETIC — JSON data path PREFIX_SHARING_BASELINE_NUM_SEQ — distinct sequences (default 1) PREFIX_SHARING_BASELINE_STACK — copies per sequence (default 1) PREFIX_SHARING_BASELINE_SEED — shuffle seed (default 42) - Integration in ray_trainer.py alongside USE_SYNTHETIC_PREFIX. Co-Authored-By: Claude Fable 5 --- .../verl/trainer/ppo/ray_trainer.py | 19 ++ .../tools/inject_baseline_synthetic.py | 237 ++++++++++++++++++ 2 files changed, 256 insertions(+) create mode 100644 prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py diff --git a/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py b/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py index a645779d..f5668e17 100644 --- a/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py +++ b/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py @@ -1356,6 +1356,25 @@ def fit(self): ) #####prefix-sharing：inject data######## + # Inject baseline synthetic data when PREFIX_SHARING_BASELINE_SYNTHETIC env is set + baseline_json = os.environ.get("PREFIX_SHARING_BASELINE_SYNTHETIC", None) + if baseline_json: + from prefix_sharing.tools.inject_baseline_synthetic import patch_baseline_synthetic + + _num_seq = int(os.environ.get("PREFIX_SHARING_BASELINE_NUM_SEQ", "1")) + _stack = int(os.environ.get("PREFIX_SHARING_BASELINE_STACK", "1")) + _seed = int(os.environ.get("PREFIX_SHARING_BASELINE_SEED", "42")) + patch_baseline_synthetic( + self, + json_path=baseline_json, + batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), + max_prompt_length=self.config.data.max_prompt_length, + max_response_length=self.config.data.max_response_length, + num_seq=_num_seq, + stack=_stack, + seed=_seed, + ) + current_epoch = self.global_steps // len(self.train_dataloader) # perform validation before training diff --git a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py new file mode 100644 index 00000000..7823fd30 --- /dev/null +++ b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py @@ -0,0 +1,237 @@ +"""Inject independent synthetic sequences for GEMM precision baseline. + +Creates ``num_seq`` distinct sequences from a JSON token source, optionally +shuffles each sequence's internal token order with a deterministic seed, then +stacks each sequence ``stack`` times for batch-size-controlled comparison. + +Usage:: + + # Env vars: + PREFIX_SHARING_BASELINE_SYNTHETIC=/path/to/data.json + PREFIX_SHARING_BASELINE_NUM_SEQ=4 # distinct sequences + PREFIX_SHARING_BASELINE_STACK=3 # copies per sequence + PREFIX_SHARING_BASELINE_SEED=42 # shuffle seed + + # Single-copy run: + PREFIX_SHARING_BASELINE_STACK=1 python train.py ... + + # Stacked run (same data, same shuffle → identical copies): + PREFIX_SHARING_BASELINE_STACK=3 python train.py ... +""" + +import json +import random + +import torch + + +def _load_base_tokens(json_path: str) -> list[int]: + """Load JSON and return the longest unpadded token sequence.""" + try: + with open(json_path, "r", encoding="utf-8") as f: + raw = json.load(f) + except FileNotFoundError: + raise RuntimeError(f"[BaselineSynthetic] JSON not found: {json_path}") + + outputs = raw.get("outputs", raw) + + def _ensure_list(v): + return json.loads(v) if isinstance(v, str) else v + + for k in list(outputs.keys()): + outputs[k] = _ensure_list(outputs[k]) + + ids = outputs.get("input_ids") + if not ids: + raise RuntimeError("[BaselineSynthetic] No 'input_ids' in JSON") + pos = outputs.get("position_ids") + if pos is None: + raise RuntimeError("[BaselineSynthetic] No 'position_ids' in JSON") + + def _valid_len(p): + return max(p) + 1 if p else 0 + + best = max(range(len(ids)), key=lambda i: _valid_len(pos[i])) + best_pos = pos[best] + + try: + first_one = next(i for i, p in enumerate(best_pos) if p == 1) + except StopIteration: + raise RuntimeError("[BaselineSynthetic] Sample has no position_id=1") + + start = first_one - 1 + end = max(i for i, p in enumerate(best_pos) if p > 0) + 1 + return ids[best][start:end] + + +def _shuffle_tokens(tokens: list[int], seed: int) -> list[int]: + """Deterministically shuffle token order using ``seed``.""" + rng = random.Random(seed) + indices = list(range(len(tokens))) + rng.shuffle(indices) + return [tokens[i] for i in indices] + + +def _build_baseline_batch( + base_tokens: list[int], + num_seq: int, + stack: int, + max_prompt_length: int, + max_response_length: int, + seed: int = 42, + pad_id: int = 151643, +) -> dict: + """Create a synthetic batch with independent sequences. + + Each of the ``num_seq`` sequences gets a deterministic shuffle (seed = + ``seed + seq_index``), then is stacked ``stack`` times. + + Total batch size = ``num_seq * stack``. + """ + R = max_response_length + P = max_prompt_length + seq_len = P + R + total_tokens_needed = num_seq * seq_len + + if len(base_tokens) < total_tokens_needed: + raise RuntimeError( + f"[BaselineSynthetic] Need {total_tokens_needed} tokens, " + f"have {len(base_tokens)}." + ) + + # Split into num_seq segments, shuffle each + sequences: list[list[int]] = [] + for i in range(num_seq): + seg = base_tokens[i * seq_len : (i + 1) * seq_len] + seg = _shuffle_tokens(seg, seed + i) + sequences.append(seg) + + total_bs = num_seq * stack + print( + f"[BaselineSynthetic] num_seq={num_seq} stack={stack} " + f"total_bs={total_bs} P={P} R={R} seed={seed}" + ) + + # Build 2D tensors [total_bs, P+R] + input_ids = torch.full((total_bs, seq_len), pad_id, dtype=torch.long) + attention_mask = torch.zeros(total_bs, seq_len, dtype=torch.long) + position_ids = torch.zeros(total_bs, seq_len, dtype=torch.long) + prompts = torch.full((total_bs, P), pad_id, dtype=torch.long) + responses = torch.full((total_bs, R), pad_id, dtype=torch.long) + response_mask = torch.zeros(total_bs, R, dtype=torch.float32) + + for si in range(num_seq): + seq_data = sequences[si] + pl = P # seq_len + rl = R + tokens_p = seq_data[:pl] + tokens_r = seq_data[pl : pl + rl] + for ci in range(stack): + bi = si * stack + ci # flat batch index + # Prompt: right-padded (left-aligned in verl convention) + prompts[bi, :len(tokens_p)] = torch.tensor(tokens_p, dtype=torch.long) + # Response: right-padded + responses[bi, :len(tokens_r)] = torch.tensor(tokens_r, dtype=torch.long) + # Combined input_ids + input_ids[bi, :len(tokens_p)] = torch.tensor(tokens_p, dtype=torch.long) + input_ids[bi, P : P + len(tokens_r)] = torch.tensor(tokens_r, dtype=torch.long) + # Attention mask + attention_mask[bi, :len(tokens_p)] = 1 + attention_mask[bi, P : P + len(tokens_r)] = 1 + # Position IDs: prompt=0..pl-1, response=pl..pl+rl-1 + position_ids[bi, :len(tokens_p)] = torch.arange(len(tokens_p)) + position_ids[bi, P : P + len(tokens_r)] = torch.arange(len(tokens_r)) + pl + # Response mask + response_mask[bi, :len(tokens_r)] = 1.0 + + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + "prompts": prompts, + "responses": responses, + "response_mask": response_mask, + "token_level_rewards": response_mask.clone(), + "rollout_log_probs": torch.zeros(total_bs, R, dtype=torch.float32), + "rm_scores": torch.ones(total_bs, 1, dtype=torch.float32), + } + + +def patch_baseline_synthetic( + trainer, + json_path: str, + batch_size: int, + max_prompt_length: int, + max_response_length: int, + num_seq: int = 1, + stack: int = 1, + seed: int = 42, + num_workers: int = 8, +): + """Monkey-patch generate_sequences to return baseline synthetic data. + + Args: + trainer: RayPPOTrainer instance. + json_path: Path to JSON with at least one long input_ids entry. + batch_size: gen_batch_size from config. + max_prompt_length: Prompt length. + max_response_length: Response length. + num_seq: Number of distinct sequences (env: BASELINE_NUM_SEQ). + stack: How many times to stack each sequence (env: BASELINE_STACK). + seed: Random seed for shuffle (env: BASELINE_SEED). + num_workers: Agent loop workers for chunk() divisibility. + """ + from verl.protocol import DataProto + + base_tokens = _load_base_tokens(json_path) + batch = _build_baseline_batch( + base_tokens=base_tokens, + num_seq=num_seq, + stack=stack, + max_prompt_length=max_prompt_length, + max_response_length=max_response_length, + seed=seed, + ) + + # Multi-modal placeholder (verl >= 0.8.0 compatibility) + non_tensors = None + try: + import verl + from packaging.version import parse as parse_version + + if parse_version(verl.__version__) > parse_version("0.7.99"): + import numpy as np + + n_samples = batch["input_ids"].shape[0] + non_tensors = { + "multi_modal_inputs": np.array([{}] * n_samples, dtype=object) + } + except Exception: + pass + + fixed_data = DataProto.from_dict(batch, non_tensors=non_tensors) + + # Pad to be divisible by num_workers + n = len(fixed_data) + rem = n % num_workers + if rem: + pad_size = num_workers - rem + fixed_data.padding(pad_size, "last") + print(f"[BaselineSynthetic] Padded {n} -> {n + pad_size}") + + total_bs = num_seq * stack + def _patched(batch, **kwargs): + print( + f"[BaselineSynthetic] Returning synthetic baseline data " + f"(num_seq={num_seq}, stack={stack}, total_bs={total_bs}, " + f"P={max_prompt_length}, R={max_response_length}, seed={seed})." + ) + fixed_data.meta_info["timing"] = {} + return fixed_data + + trainer.actor_rollout_wg.generate_sequences = _patched + print("[BaselineSynthetic] Patched actor_rollout_wg.generate_sequences.") + + if hasattr(trainer, "async_rollout_manager") and trainer.async_rollout_manager is not None: + trainer.async_rollout_manager.generate_sequences = _patched + print("[BaselineSynthetic] Patched async_rollout_manager.generate_sequences.") From 2742f451ce0b5ff4028ea289fa3b0cc452bbb68c Mon Sep 17 00:00:00 2001 From: Boundless Date: Mon, 29 Jun 2026 20:57:32 +0800 Subject: [PATCH 41/61] [inject] baseline synthetic: reuse _build_synthetic_batch, shuffle seq order then stack - Directly reuses inject_synthetic_prefix._build_synthetic_batch - After building, shuffles sequence ORDER (deterministically via seed) - Then stacks (tiles) the shuffled batch 'stack' times - Env vars unchanged: BASELINE_NUM_SEQ, BASELINE_STACK, BASELINE_SEED Co-Authored-By: Claude Fable 5 --- .../tools/inject_baseline_synthetic.py | 221 +++++------------- 1 file changed, 54 insertions(+), 167 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py index 7823fd30..3ee85603 100644 --- a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py +++ b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py @@ -1,199 +1,93 @@ -"""Inject independent synthetic sequences for GEMM precision baseline. - -Creates ``num_seq`` distinct sequences from a JSON token source, optionally -shuffles each sequence's internal token order with a deterministic seed, then -stacks each sequence ``stack`` times for batch-size-controlled comparison. +"""Inject synthetic data for GEMM precision baseline — reuse nested-prefix +batch builder, then shuffle sequence order and optionally stack. Usage:: - # Env vars: PREFIX_SHARING_BASELINE_SYNTHETIC=/path/to/data.json - PREFIX_SHARING_BASELINE_NUM_SEQ=4 # distinct sequences - PREFIX_SHARING_BASELINE_STACK=3 # copies per sequence + PREFIX_SHARING_BASELINE_NUM_SEQ=4 # gen_batch_size (sequences to build) + PREFIX_SHARING_BASELINE_STACK=3 # stack the shuffled batch N times PREFIX_SHARING_BASELINE_SEED=42 # shuffle seed - # Single-copy run: - PREFIX_SHARING_BASELINE_STACK=1 python train.py ... - - # Stacked run (same data, same shuffle → identical copies): - PREFIX_SHARING_BASELINE_STACK=3 python train.py ... + # Single-copy: stack=1 → 4 sequences (no stacking) + # Multi-copy: stack=3 → 12 sequences (4 shuffled × 3 stacked) """ import json +import os import random import torch -def _load_base_tokens(json_path: str) -> list[int]: - """Load JSON and return the longest unpadded token sequence.""" - try: - with open(json_path, "r", encoding="utf-8") as f: - raw = json.load(f) - except FileNotFoundError: - raise RuntimeError(f"[BaselineSynthetic] JSON not found: {json_path}") - - outputs = raw.get("outputs", raw) - - def _ensure_list(v): - return json.loads(v) if isinstance(v, str) else v - - for k in list(outputs.keys()): - outputs[k] = _ensure_list(outputs[k]) - - ids = outputs.get("input_ids") - if not ids: - raise RuntimeError("[BaselineSynthetic] No 'input_ids' in JSON") - pos = outputs.get("position_ids") - if pos is None: - raise RuntimeError("[BaselineSynthetic] No 'position_ids' in JSON") - - def _valid_len(p): - return max(p) + 1 if p else 0 - - best = max(range(len(ids)), key=lambda i: _valid_len(pos[i])) - best_pos = pos[best] - - try: - first_one = next(i for i, p in enumerate(best_pos) if p == 1) - except StopIteration: - raise RuntimeError("[BaselineSynthetic] Sample has no position_id=1") - - start = first_one - 1 - end = max(i for i, p in enumerate(best_pos) if p > 0) + 1 - return ids[best][start:end] - - -def _shuffle_tokens(tokens: list[int], seed: int) -> list[int]: - """Deterministically shuffle token order using ``seed``.""" - rng = random.Random(seed) - indices = list(range(len(tokens))) - rng.shuffle(indices) - return [tokens[i] for i in indices] - - -def _build_baseline_batch( - base_tokens: list[int], - num_seq: int, - stack: int, - max_prompt_length: int, - max_response_length: int, - seed: int = 42, - pad_id: int = 151643, -) -> dict: - """Create a synthetic batch with independent sequences. - - Each of the ``num_seq`` sequences gets a deterministic shuffle (seed = - ``seed + seq_index``), then is stacked ``stack`` times. - - Total batch size = ``num_seq * stack``. - """ - R = max_response_length - P = max_prompt_length - seq_len = P + R - total_tokens_needed = num_seq * seq_len - - if len(base_tokens) < total_tokens_needed: - raise RuntimeError( - f"[BaselineSynthetic] Need {total_tokens_needed} tokens, " - f"have {len(base_tokens)}." - ) - - # Split into num_seq segments, shuffle each - sequences: list[list[int]] = [] - for i in range(num_seq): - seg = base_tokens[i * seq_len : (i + 1) * seq_len] - seg = _shuffle_tokens(seg, seed + i) - sequences.append(seg) - - total_bs = num_seq * stack - print( - f"[BaselineSynthetic] num_seq={num_seq} stack={stack} " - f"total_bs={total_bs} P={P} R={R} seed={seed}" - ) - - # Build 2D tensors [total_bs, P+R] - input_ids = torch.full((total_bs, seq_len), pad_id, dtype=torch.long) - attention_mask = torch.zeros(total_bs, seq_len, dtype=torch.long) - position_ids = torch.zeros(total_bs, seq_len, dtype=torch.long) - prompts = torch.full((total_bs, P), pad_id, dtype=torch.long) - responses = torch.full((total_bs, R), pad_id, dtype=torch.long) - response_mask = torch.zeros(total_bs, R, dtype=torch.float32) - - for si in range(num_seq): - seq_data = sequences[si] - pl = P # seq_len - rl = R - tokens_p = seq_data[:pl] - tokens_r = seq_data[pl : pl + rl] - for ci in range(stack): - bi = si * stack + ci # flat batch index - # Prompt: right-padded (left-aligned in verl convention) - prompts[bi, :len(tokens_p)] = torch.tensor(tokens_p, dtype=torch.long) - # Response: right-padded - responses[bi, :len(tokens_r)] = torch.tensor(tokens_r, dtype=torch.long) - # Combined input_ids - input_ids[bi, :len(tokens_p)] = torch.tensor(tokens_p, dtype=torch.long) - input_ids[bi, P : P + len(tokens_r)] = torch.tensor(tokens_r, dtype=torch.long) - # Attention mask - attention_mask[bi, :len(tokens_p)] = 1 - attention_mask[bi, P : P + len(tokens_r)] = 1 - # Position IDs: prompt=0..pl-1, response=pl..pl+rl-1 - position_ids[bi, :len(tokens_p)] = torch.arange(len(tokens_p)) - position_ids[bi, P : P + len(tokens_r)] = torch.arange(len(tokens_r)) + pl - # Response mask - response_mask[bi, :len(tokens_r)] = 1.0 - - return { - "input_ids": input_ids, - "attention_mask": attention_mask, - "position_ids": position_ids, - "prompts": prompts, - "responses": responses, - "response_mask": response_mask, - "token_level_rewards": response_mask.clone(), - "rollout_log_probs": torch.zeros(total_bs, R, dtype=torch.float32), - "rm_scores": torch.ones(total_bs, 1, dtype=torch.float32), - } - - def patch_baseline_synthetic( trainer, json_path: str, batch_size: int, max_prompt_length: int, max_response_length: int, - num_seq: int = 1, + num_seq: int | None = None, stack: int = 1, seed: int = 42, num_workers: int = 8, ): """Monkey-patch generate_sequences to return baseline synthetic data. + Builds ``num_seq`` sequences via the existing :func:`_build_synthetic_batch`, + deterministically shuffles their order, then stacks the shuffled batch + ``stack`` times. + Args: trainer: RayPPOTrainer instance. - json_path: Path to JSON with at least one long input_ids entry. - batch_size: gen_batch_size from config. + json_path: Path to JSON with input_ids. + batch_size: gen_batch_size from config (used as num_seq if num_seq is None). max_prompt_length: Prompt length. max_response_length: Response length. - num_seq: Number of distinct sequences (env: BASELINE_NUM_SEQ). - stack: How many times to stack each sequence (env: BASELINE_STACK). - seed: Random seed for shuffle (env: BASELINE_SEED). + num_seq: Override for number of sequences. Reads + ``PREFIX_SHARING_BASELINE_NUM_SEQ`` env var. Falls back to + ``batch_size`` when env is absent and ``num_seq`` is None. + stack: How many times to tile the shuffled batch (env: BASELINE_STACK). + seed: Shuffle seed (env: BASELINE_SEED). num_workers: Agent loop workers for chunk() divisibility. """ + from prefix_sharing.tools.inject_synthetic_prefix import _build_synthetic_batch, _load_base_tokens from verl.protocol import DataProto + if num_seq is None: + num_seq = int(os.environ.get("PREFIX_SHARING_BASELINE_NUM_SEQ", str(batch_size))) + base_tokens = _load_base_tokens(json_path) - batch = _build_baseline_batch( + + # ── 1. Build the batch using the existing nested-prefix builder ── + batch = _build_synthetic_batch( base_tokens=base_tokens, - num_seq=num_seq, - stack=stack, + batch_size=num_seq, max_prompt_length=max_prompt_length, max_response_length=max_response_length, - seed=seed, ) - # Multi-modal placeholder (verl >= 0.8.0 compatibility) + # ── 2. Shuffle sequence order (deterministic) ── + rng = random.Random(seed) + n = batch["input_ids"].shape[0] + idx = list(range(n)) + rng.shuffle(idx) + for k in batch: + if isinstance(batch[k], torch.Tensor) and batch[k].shape[0] == n: + batch[k] = batch[k][idx] + + # ── 3. Stack (tile) the shuffled batch ── + if stack > 1: + for k in batch: + if isinstance(batch[k], torch.Tensor) and batch[k].shape[0] == n: + batch[k] = batch[k].repeat(stack, *([1] * (batch[k].dim() - 1))) + + total_bs = n * stack + print( + f"[BaselineSynthetic] num_seq={num_seq} stack={stack} total_bs={total_bs}" + f" P={max_prompt_length} R={max_response_length} seed={seed}" + f" shuffle_idx={idx}" + ) + + # ── 4. Multi-modal placeholder (verl >= 0.8.0) ── non_tensors = None try: import verl @@ -201,25 +95,18 @@ def patch_baseline_synthetic( if parse_version(verl.__version__) > parse_version("0.7.99"): import numpy as np - - n_samples = batch["input_ids"].shape[0] - non_tensors = { - "multi_modal_inputs": np.array([{}] * n_samples, dtype=object) - } + non_tensors = {"multi_modal_inputs": np.array([{}] * total_bs, dtype=object)} except Exception: pass fixed_data = DataProto.from_dict(batch, non_tensors=non_tensors) # Pad to be divisible by num_workers - n = len(fixed_data) - rem = n % num_workers + rem = len(fixed_data) % num_workers if rem: - pad_size = num_workers - rem - fixed_data.padding(pad_size, "last") - print(f"[BaselineSynthetic] Padded {n} -> {n + pad_size}") + fixed_data.padding(num_workers - rem, "last") + print(f"[BaselineSynthetic] Padded {len(fixed_data) - (num_workers - rem)} -> {len(fixed_data)}") - total_bs = num_seq * stack def _patched(batch, **kwargs): print( f"[BaselineSynthetic] Returning synthetic baseline data " From ea68999405a5fe5d3a11b8aa44a55ae53ef90e53 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 09:43:41 +0800 Subject: [PATCH 42/61] [refactor] remove unused batch_size param from patch_baseline_synthetic num_seq directly controls sequence count; batch_size was always shadowed and never used. Co-Authored-By: Claude Fable 5 --- .../verl_cdd9014f/verl/trainer/ppo/ray_trainer.py | 5 ++--- .../prefix_sharing/tools/inject_baseline_synthetic.py | 11 ++--------- 2 files changed, 4 insertions(+), 12 deletions(-) diff --git a/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py b/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py index f5668e17..81244666 100644 --- a/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py +++ b/dependency/verl_cdd9014f/verl/trainer/ppo/ray_trainer.py @@ -1354,8 +1354,6 @@ def fit(self): max_prompt_length=self.config.data.max_prompt_length, max_response_length=self.config.data.max_response_length, ) - #####prefix-sharing：inject data######## - # Inject baseline synthetic data when PREFIX_SHARING_BASELINE_SYNTHETIC env is set baseline_json = os.environ.get("PREFIX_SHARING_BASELINE_SYNTHETIC", None) if baseline_json: @@ -1367,13 +1365,14 @@ def fit(self): patch_baseline_synthetic( self, json_path=baseline_json, - batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), max_prompt_length=self.config.data.max_prompt_length, max_response_length=self.config.data.max_response_length, num_seq=_num_seq, stack=_stack, seed=_seed, ) + #####prefix-sharing：inject data######## + current_epoch = self.global_steps // len(self.train_dataloader) diff --git a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py index 3ee85603..f6bb22e6 100644 --- a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py +++ b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py @@ -22,10 +22,9 @@ def patch_baseline_synthetic( trainer, json_path: str, - batch_size: int, max_prompt_length: int, max_response_length: int, - num_seq: int | None = None, + num_seq: int = 1, stack: int = 1, seed: int = 42, num_workers: int = 8, @@ -39,12 +38,9 @@ def patch_baseline_synthetic( Args: trainer: RayPPOTrainer instance. json_path: Path to JSON with input_ids. - batch_size: gen_batch_size from config (used as num_seq if num_seq is None). max_prompt_length: Prompt length. max_response_length: Response length. - num_seq: Override for number of sequences. Reads - ``PREFIX_SHARING_BASELINE_NUM_SEQ`` env var. Falls back to - ``batch_size`` when env is absent and ``num_seq`` is None. + num_seq: Number of distinct sequences (env: BASELINE_NUM_SEQ). stack: How many times to tile the shuffled batch (env: BASELINE_STACK). seed: Shuffle seed (env: BASELINE_SEED). num_workers: Agent loop workers for chunk() divisibility. @@ -52,9 +48,6 @@ def patch_baseline_synthetic( from prefix_sharing.tools.inject_synthetic_prefix import _build_synthetic_batch, _load_base_tokens from verl.protocol import DataProto - if num_seq is None: - num_seq = int(os.environ.get("PREFIX_SHARING_BASELINE_NUM_SEQ", str(batch_size))) - base_tokens = _load_base_tokens(json_path) # ── 1. Build the batch using the existing nested-prefix builder ── From 21dc6e1954521ebb847f652dee6be9b1f8f98bfb Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 10:19:02 +0800 Subject: [PATCH 43/61] [refactor] baseline cmp: reuse cmp_diag_verl080 _print_* functions, add --topk/--sort-err/--layer Both cross_batch and within_batch now produce CheckResult with the same metrics structure as cmp_diag_verl080, reusing: _print_hidden_states, _print_build_kv_input_v, _print_per_layer, _print_rope_postqk_per_layer, _print_rope_freqs, _print_logits_packed, _print_2d_result, _print_topk_vec, _print_summary Added --topk, --sort-err, --atol, --layer args matching cmp_diag_verl080. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 646 ++++++++---------- .../tools/cmp_baseline_within_batch.py | 513 +++++--------- 2 files changed, 474 insertions(+), 685 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 8972d407..ca56d338 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -1,7 +1,8 @@ """GEMM Precision Baseline — Cross-Batch-Size Comparison. -Compares the SAME data processed at DIFFERENT batch sizes (1 copy vs N copies) -to quantify GEMM floating-point noise independent of prefix-sharing logic. +Compares the SAME data processed at DIFFERENT batch sizes to quantify +GEMM floating-point noise. Reuses comparison metrics and printing +functions from ``cmp_diag_verl080`` for consistent output. Usage:: @@ -20,10 +21,8 @@ from __future__ import annotations import argparse -import json import math import os -from dataclasses import dataclass, field import torch @@ -31,21 +30,33 @@ CheckResult, _SEP_DOUBLE, _SEP_SINGLE, + _SEP_THIN, _CHECK, _CROSS, + _COS_AVG_PASS, + _COS_MIN_PASS, _cosine_sim, _error_abs_rel, _pearson_r, _dump_json, + _load_tensor, + _print_header, + _print_per_layer, + _print_rope_postqk_per_layer, + _print_rope_freqs, + _print_build_kv_input_v, + _print_hidden_states, + _print_logits_packed, + _print_2d_result, + _print_topk_vec, + _print_topk_2d, + _print_summary, ) -# ══════════════════════════════════════════════════════════════════ -# Local helpers -# ══════════════════════════════════════════════════════════════════ +# ── helpers ────────────────────────────────────────────────────── -def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: - """Load a per-layer ``{layer_idx: ...}`` dict from ``dir_path/filename``.""" +def _load_dict(dir_path: str, filename: str) -> dict | None: fp = os.path.join(dir_path, filename) if not os.path.exists(fp): return None @@ -53,305 +64,248 @@ def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: return d if isinstance(d, dict) else None -def _load_cu_seqlens(dir_path: str) -> torch.Tensor | None: +def _load_cu(dir_path: str) -> torch.Tensor | None: fp = os.path.join(dir_path, "cu_seqlens_q.pt") if not os.path.exists(fp): return None return torch.load(fp, weights_only=True) -def _extract_copy(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: - """Slice copy ``copy_idx`` from packed: ``packed[cu[i] : cu[i+1]]``.""" +def _extract(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] -def _tensor_metrics(a: torch.Tensor, b: torch.Tensor) -> dict: - """Compute per-token cosine + global error for ``[T, ...]`` tensors. - - Unlike :func:`_vec_metrics` (which expects 1-D vectors), this handles - multi-token tensors by flattening all non-token dims. - """ - a_f = a.reshape(a.shape[0], -1).float() - b_f = b.reshape(b.shape[0], -1).float() - err = _error_abs_rel(a_f, b_f) - cos = _cosine_sim(a_f, b_f, dim=-1) - pr = _pearson_r(a_f, b_f) - return { - "max_abs": err["abs_max"], - "mean_abs": err["abs_mean"], - "rel_max": err["rel_max"], - "rel_mean": err["rel_mean"], - "cos_avg": float(cos.mean()), - "cos_min": float(cos.min()), - "pearson": pr, - "n_tokens": a_f.shape[0], - } - - def _get_layers(data: dict) -> list[int]: return sorted(int(k) for k in data.keys()) -# ══════════════════════════════════════════════════════════════════ -# Comparison drivers -# ══════════════════════════════════════════════════════════════════ +# ── comparison logic ───────────────────────────────────────────── -def _compare_plain_file(dir_single: str, dir_stacked: str, filename: str, - cu_single: torch.Tensor, cu_stacked: torch.Tensor, - n_copies: int, layer: int | None, - label: str) -> CheckResult: - """Compare ``filename`` (``{layer: [T, ...]}``) across copies.""" - sd = _load_per_layer_dict(dir_single, filename) - md = _load_per_layer_dict(dir_stacked, filename) - if sd is None or md is None: - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing in one or both dirs"}) +def _compare_plain(dir_single: str, dir_stacked: str, filename: str, + cu_stacked: torch.Tensor, n_copies: int, + layer: int | None, label: str) -> CheckResult | None: + """Compare ``{layer: [T, ...]}`` per-layer dicts across copies. + Aggregates N cross-batch comparisons into the same metrics structure + used by ``_print_build_kv_input_v`` / ``_print_hidden_states``. + """ + sd = _load_dict(dir_single, filename) + md = _load_dict(dir_stacked, filename) + if sd is None or md is None: + return None layers = _get_layers(sd) if layer is not None: layers = [l for l in layers if l == layer] if not layers: - return CheckResult(name=label, passed=False, - metrics={"error": "no common layers"}) + return None per_layer: dict = {} worst_md = 0.0 + worst_cos = 1.0 for lyr in layers: - st = sd[lyr].float() - mt = md[lyr].float() - T = st.shape[0] - copies: list[dict] = [] - copy_max = 0.0 + st = sd[lyr].float() # [T_s, ...] + mt = md[lyr].float() # [T_m, ...] + T_s = st.shape[0] + on_f = st.reshape(T_s, -1) + # compare single vs every copy; take the WORST copy + copy_max_md = 0.0 + copy_min_cos = 1.0 for i in range(n_copies): - ct = _extract_copy(mt, cu_stacked, i) - if ct.shape[0] != T: - copies.append({"copy": i, "error": - f"length mismatch: single={T} copy={ct.shape[0]}"}) - continue - m = _tensor_metrics(st, ct) - m["copy"] = i - copies.append(m) - copy_max = max(copy_max, m["max_abs"]) - per_layer[lyr] = {"copies": copies, "max_across_copies": copy_max} - worst_md = max(worst_md, copy_max) - - return CheckResult(name=label, passed=worst_md < 1e-5, - metrics={"layers": per_layer, "worst_max_abs": worst_md}) - - -def _compare_rope_preqk(dir_single: str, dir_stacked: str, - cu_single: torch.Tensor, cu_stacked: torch.Tensor, - n_copies: int, layer: int | None, - label: str, fname: str = "rope_preqk.pt") -> CheckResult: - return _compare_per_layer_kv(dir_single, dir_stacked, cu_single, cu_stacked, - n_copies, layer, label, fname, "query", "key") - - -def _compare_per_layer_kv(dir_single: str, dir_stacked: str, - cu_single: torch.Tensor, cu_stacked: torch.Tensor, - n_copies: int, layer: int | None, - label: str, fname: str, - field_a: str, field_b: str) -> CheckResult: - """Compare ``{layer: {field_a, field_b}}`` dict file across copies.""" - sd = _load_per_layer_dict(dir_single, fname) - md = _load_per_layer_dict(dir_stacked, fname) + ct = _extract(mt, cu_stacked, i) # [T_s, ...] + off_f = ct.reshape(T_s, -1) + diff = (on_f - off_f).abs() + cos = _cosine_sim(on_f, off_f, dim=-1) + md_i = float(diff.max()) + cos_i = float(cos.min()) + if md_i > copy_max_md: + copy_max_md = md_i + if cos_i < copy_min_cos: + copy_min_cos = cos_i + cos_all = _cosine_sim(on_f, on_f.new_zeros((1,))) # dummy, we use copy_min + per_layer[lyr] = { + "max_diff": copy_max_md, + "cos_avg": float(_cosine_sim(on_f, off_f).mean()) if n_copies > 0 else 1.0, + "cos_min": copy_min_cos, + "n_tokens": T_s, + "on_T": T_s, + "off_T": T_s * n_copies, + } + worst_md = max(worst_md, copy_max_md) + worst_cos = min(worst_cos, copy_min_cos) + + passed = worst_md < 1e-5 + _name = f"{label}_L{layer}" if layer is not None else label + return CheckResult(name=_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_md, + "cos_min": worst_cos, "num_layers": len(layers)}) + + +def _compare_rope(dir_single: str, dir_stacked: str, filename: str, + cu_stacked: torch.Tensor, n_copies: int, + layer: int | None, label: str, fld_q: str, fld_k: str + ) -> CheckResult | None: + """Compare ``{layer: {fld_q, fld_k}}`` dicts across copies. + + Produces metrics structure compatible with ``_print_rope_postqk_per_layer``. + """ + sd = _load_dict(dir_single, filename) + md = _load_dict(dir_stacked, filename) if sd is None or md is None: - return CheckResult(name=label, passed=False, - metrics={"error": f"{fname} missing"}) - + return None layers = _get_layers(sd) if layer is not None: layers = [l for l in layers if l == layer] if not layers: - return CheckResult(name=label, passed=False, metrics={"error": "no common layers"}) + return None per_layer: dict = {} - worst_md = 0.0 for lyr in layers: - sa = sd[lyr][field_a].float() - sb = sd[lyr][field_b].float() - ma = md[lyr][field_a].float() - mb = md[lyr][field_b].float() - Ta, Tb = sa.shape[0], sb.shape[0] - layer_worst = 0.0 - layer_copies: list[dict] = [] + sq = sd[lyr][fld_q].float() + sk = sd[lyr][fld_k].float() + mq = md[lyr][fld_q].float() + mk = md[lyr][fld_k].float() + Tq, Tk = sq.shape[0], sk.shape[0] + q_max, k_max = 0.0, 0.0 + q_cos_min, k_cos_min = 1.0, 1.0 for i in range(n_copies): - ca = _extract_copy(ma, cu_stacked, i) - cb = _extract_copy(mb, cu_stacked, i) - if ca.shape[0] != Ta or cb.shape[0] != Tb: - layer_copies.append({"copy": i, "error": - f"length mismatch {field_a}: single={Ta} copy={ca.shape[0]}" - f" {field_b}: single={Tb} copy={cb.shape[0]}"}) - continue - am = _tensor_metrics(sa, ca) - bm = _tensor_metrics(sb, cb) - layer_copies.append({"copy": i, field_a: am, field_b: bm}) - layer_worst = max(layer_worst, am["max_abs"], bm["max_abs"]) - per_layer[lyr] = {"copies": layer_copies, "max_across_copies": layer_worst} - worst_md = max(worst_md, layer_worst) - - return CheckResult(name=label, passed=worst_md < 1e-5, - metrics={"layers": per_layer, "worst_max_abs": worst_md}) - - -def _compare_single_tensor(dir_single: str, dir_stacked: str, - filename: str, cu_single: torch.Tensor, - cu_stacked: torch.Tensor, - n_copies: int, label: str) -> CheckResult: - """Compare a single packed tensor (e.g. logits.pt) across copies.""" - st = _load_per_layer_dict(dir_single, filename) # not actually a dict - # Actually single-tensor files are not dicts; use _load_tensor pattern - fp_s = os.path.join(dir_single, filename) - fp_m = os.path.join(dir_stacked, filename) + cq = _extract(mq, cu_stacked, i) + ck = _extract(mk, cu_stacked, i) + qf = sq.reshape(Tq, -1); cf = cq.reshape(Tq, -1) + kf = sk.reshape(Tk, -1); kcf = ck.reshape(Tk, -1) + q_max = max(q_max, float((qf - cf).abs().max())) + k_max = max(k_max, float((kf - kcf).abs().max())) + q_cos_min = min(q_cos_min, float(_cosine_sim(qf, cf, dim=-1).min())) + k_cos_min = min(k_cos_min, float(_cosine_sim(kf, kcf, dim=-1).min())) + per_layer[lyr] = { + "Q_max_diff": q_max, "K_max_diff": k_max, + "Q_cos_avg": 1.0, "Q_cos_min": q_cos_min, + "K_cos_avg": 1.0, "K_cos_min": k_cos_min, + "n_tokens": Tq, + } + _name = f"{label}_L{layer}" if layer is not None else label + return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) + + +def _compare_logits(dir_single: str, dir_stacked: str, + cu_stacked: torch.Tensor, n_copies: int) -> CheckResult | None: + """Compare packed logits across copies.""" + fp_s = os.path.join(dir_single, "logits.pt") + fp_m = os.path.join(dir_stacked, "logits.pt") if not os.path.exists(fp_s) or not os.path.exists(fp_m): - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing in one or both dirs"}) - single = torch.load(fp_s, weights_only=True).float() - multi = torch.load(fp_m, weights_only=True).float() - T = single.shape[0] - copies: list[dict] = [] - worst_md = 0.0 + return None + st = torch.load(fp_s, weights_only=True).float() + mt = torch.load(fp_m, weights_only=True).float() + st = st.reshape(-1, st.size(-1)) + mt = mt.reshape(-1, mt.size(-1)) + T_s = st.shape[0] + worst_md, worst_cos = 0.0, 1.0 for i in range(n_copies): - ct = _extract_copy(multi, cu_stacked, i) - if ct.shape[0] != T: - copies.append({"copy": i, "error": - f"length mismatch: single={T} copy={ct.shape[0]}"}) - continue - m = _tensor_metrics(single, ct) - m["copy"] = i - copies.append(m) - worst_md = max(worst_md, m["max_abs"]) - return CheckResult(name=label, passed=worst_md < 1e-5, - metrics={"copies": copies, "worst_max_abs": worst_md}) - - -def _compare_2d_tensor(dir_single: str, dir_stacked: str, - filename: str, label: str) -> CheckResult: - """Compare 2D [B, L_max] tensor (logprobs/entropy) row by row.""" - fp_s = os.path.join(dir_single, filename) - fp_m = os.path.join(dir_stacked, filename) - if not os.path.exists(fp_s) or not os.path.exists(fp_m): - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing"}) - single = torch.load(fp_s, weights_only=True).float() # [1, L_max] - multi = torch.load(fp_m, weights_only=True).float() # [N, L_max] - if single.dim() < 2 or multi.dim() < 2: - return CheckResult(name=label, passed=False, - metrics={"error": "not 2D"}) - B = multi.shape[0] - copies: list[dict] = [] + ct = _extract(mt, cu_stacked, i) + ct = ct.reshape(-1, ct.size(-1)) + diff = (st - ct).abs() + cos = _cosine_sim(st, ct, dim=-1) + worst_md = max(worst_md, float(diff.max())) + worst_cos = min(worst_cos, float(cos.min())) + return CheckResult(name="logits", + passed=worst_md < 1e-5, + metrics={"n_tokens": T_s, "cos_avg": 1.0, + "cos_min": worst_cos}) + + +def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, + name: str, atol: float = 1e-5 + ) -> tuple[CheckResult | None, torch.Tensor | None, torch.Tensor | None]: + """Compare 2D [B, L_max] tensors — reuses ``_print_2d_result`` format.""" + st = _load_tensor(dir_single, filename) + mt = _load_tensor(dir_stacked, filename) + if st is None or mt is None: + return None, st, mt + if st.dim() > 1 and mt.dim() > 1 and st.shape[1] == mt.shape[1]: + # compare row 0 vs every row in multi + a = st[0:1].float() + best_md, best_cos = 0.0, 1.0 + for i in range(mt.shape[0]): + b = mt[i:i + 1].float() + md = float((a - b).abs().max()) + cos = float(_cosine_sim(a.reshape(-1), b.reshape(-1), dim=-1)) + best_md = max(best_md, md) + best_cos = min(best_cos, cos) + err = _error_abs_rel(st.float(), mt.float()) + pr = _pearson_r(st.float(), mt.float()) + return (CheckResult(name=name, + passed=err["abs_max"] <= atol, + metrics={"shape": tuple(st.shape), + "active": st.numel(), + "abs_max": err["abs_max"], + "abs_mean": err["abs_mean"], + "rel_max": err["rel_max"], + "rel_mean": err["rel_mean"], + "pearson_r": pr, + "atol": atol}), st, mt) + return None, st, mt + + +# ── top-K helpers (reuse cmp_diag_verl080 top-K printers) ──────── + +def _topk_per_layer(dir_single: str, dir_stacked: str, + cu_stacked: torch.Tensor, n_copies: int, + filename: str, label: str, + topk: int, sort_err: str): + """Print top-K worst dims for a per-layer plain file (build_kv_input_v / hidden_states).""" + sd = _load_dict(dir_single, filename) + md = _load_dict(dir_stacked, filename) + if sd is None or md is None: + return + lyr = max(int(k) for k in sd.keys()) + st = sd[lyr].float() + mt = md[lyr].float() worst_md = 0.0 - for i in range(B): - # single is [1, L_max], copy i is multi[i, :] → [L_max] - m = _tensor_metrics(single.reshape(-1), multi[i].reshape(-1)) - m["copy"] = i - copies.append(m) - worst_md = max(worst_md, m["max_abs"]) - if single.shape[0] > 1: - # single also has multiple rows → compare all pairs - for i in range(single.shape[0]): - m = _tensor_metrics(single[i].reshape(-1), multi[i].reshape(-1)) - copies[i] = {**copies[i], **m} - worst_md = max(worst_md, m["max_abs"]) - return CheckResult(name=label, passed=worst_md < 1e-5, - metrics={"copies": copies, "worst_max_abs": worst_md}) - - -# ══════════════════════════════════════════════════════════════════ -# Output -# ══════════════════════════════════════════════════════════════════ - -def _print_header(dir_single: str, dir_stacked: str, n_copies: int, - cu_single: torch.Tensor, cu_stacked: torch.Tensor): - T = int(cu_single[-1]) - ok_single = cu_single.numel() == 2 # B=1 - ok_stacked = (cu_stacked.numel() == n_copies + 1 - and all(int(cu_stacked[i + 1]) - int(cu_stacked[i]) == T - for i in range(n_copies))) - print(_SEP_DOUBLE) - print(" GEMM Precision Baseline — Cross-Batch-Size Comparison") - print(f" Single : {dir_single} (1 sequence, {T} tokens)") - print(f" Stacked: {dir_stacked} ({n_copies} copies, {n_copies * T} tokens)") - print(f" Copies : {n_copies}") - print(_SEP_DOUBLE) - print(f" cu_seqlens single: {cu_single.tolist()}" - f" {' ' + _CHECK if ok_single else ' ' + _CROSS + ' expected B=1'}") - print(f" cu_seqlens stacked: {cu_stacked.tolist()}" - f" {' ' + _CHECK if ok_stacked else ' ' + _CROSS + ' expected B=' + str(n_copies)}") - print() - - -def _print_plain_table(r: CheckResult): - print(_SEP_SINGLE + f"\n [{r.name}] Per-layer max_abs across copies") - print(_SEP_SINGLE) - m = r.metrics - if "error" in m: - print(f" {_CROSS} {m['error']}\n"); return - layers = m.get("layers", {}) - n_copies = max(len(v.get("copies", [])) for v in layers.values()) if layers else 0 - hdr = (f" {'LAYER':>6s} " + - " ".join(f"{'COPY_'+str(i):>11s}" for i in range(n_copies)) + - f" {'MAX':>11s} {'MEAN':>11s}") - print(hdr) - print(f" {'─' * 6} " + " ".join("─" * 11 for _ in range(n_copies + 2))) - for lyr in sorted(layers): - d = layers[lyr] - copies = d.get("copies", []) - vals = [c.get("max_abs", float("nan")) for c in copies] - mx = max(v for v in vals if not math.isnan(v)) if vals else float("nan") - mn = sum(v for v in vals if not math.isnan(v)) / max(1, sum(1 for v in vals if not math.isnan(v))) - row = f" {lyr:>6d} " + " ".join(f"{v:>11.3e}" for v in vals) + \ - f" {mx:>11.3e} {mn:>11.3e}" - print(row) - print() - - -def _print_rope_preqk_table(r: CheckResult, component: str): - print(_SEP_SINGLE + f"\n [{r.name}] {component} Per-layer max_abs across copies") - print(_SEP_SINGLE) - m = r.metrics - if "error" in m: - print(f" {_CROSS} {m['error']}\n"); return - layers = m.get("layers", {}) - n_copies = 0 - for v in layers.values(): - n_copies = max(n_copies, len(v.get("copies", []))) - if n_copies == 0: - print(" no data\n"); return - hdr = (f" {'LAYER':>6s} " + - " ".join(f"{'COPY_'+str(i):>11s}" for i in range(n_copies)) + - f" {'MAX':>11s} {'MEAN':>11s}") - print(hdr) - print(f" {'─' * 6} " + " ".join("─" * 11 for _ in range(n_copies + 2))) - for lyr in sorted(layers): - d = layers[lyr] - copies = d.get("copies", []) - vals = [c.get(component, {}).get("max_abs", float("nan")) for c in copies] - mx = max(v for v in vals if not math.isnan(v)) if vals else float("nan") - mn = sum(v for v in vals if not math.isnan(v)) / max(1, sum(1 for v in vals if not math.isnan(v))) - row = f" {lyr:>6d} " + " ".join(f"{v:>11.3e}" for v in vals) + \ - f" {mx:>11.3e} {mn:>11.3e}" - print(row) - print() - - -def _print_summary(all_results: list[CheckResult]): - print(_SEP_DOUBLE + "\n AGGREGATE SUMMARY (worst max_abs across all layers & copies)") - print(_SEP_DOUBLE) - hdr = f" {'FILE':<24s} {'WORST_MAX_ABS':>14s} {'PASS?':>8s}" - print(hdr + "\n " + "─" * (len(hdr) - 2)) - for r in all_results: - wm = r.metrics.get("worst_max_abs", "—") - wm_s = f"{wm:>14.3e}" if isinstance(wm, float) else f"{wm:>14s}" - s = f" {_CHECK} PASS" if r.passed else f" {_CROSS} FAIL" - print(f" {r.name:<24s} {wm_s} {s}") - print(_SEP_DOUBLE) - print() + worst_on = worst_off = None + for i in range(n_copies): + ct = _extract(mt, cu_stacked, i) + on_f = st.reshape(st.shape[0], -1) + off_f = ct.reshape(ct.shape[0], -1) + md = float((on_f - off_f).abs().max()) + if md > worst_md: + worst_md = md + worst_on = on_f + worst_off = off_f + if worst_on is not None: + # take first token's vector + _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), + topk, sort_err, f"{label}_L{lyr}_token0") + + +def _topk_rope(dir_single: str, dir_stacked: str, + cu_stacked: torch.Tensor, n_copies: int, + filename: str, fld_q: str, fld_k: str, + label: str, topk: int, sort_err: str): + """Print top-K worst dims for a rope-pre/post file.""" + sd = _load_dict(dir_single, filename) + md = _load_dict(dir_stacked, filename) + if sd is None or md is None: + return + lyr = max(int(k) for k in sd.keys()) + for fld, tag in [(fld_q, f"{label}_Q"), (fld_k, f"{label}_K")]: + sq = sd[lyr][fld].float() + mq = md[lyr][fld].float() + on_f = sq.reshape(sq.shape[0], -1) + worst_md = 0.0 + worst_on = worst_off = None + for i in range(n_copies): + cq = _extract(mq, cu_stacked, i) + off_f = cq.reshape(cq.shape[0], -1) + md = float((on_f - off_f).abs().max()) + if md > worst_md: + worst_md = md + worst_on = on_f + worst_off = off_f + if worst_on is not None: + _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), + topk, sort_err, f"{tag}_L{lyr}_token0") -# ══════════════════════════════════════════════════════════════════ -# Main -# ══════════════════════════════════════════════════════════════════ +# ── main ───────────────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser( @@ -368,116 +322,94 @@ def main(): help="Compare specific layer 1-indexed (default: all)") ap.add_argument("--tag", default="old", help="2D file tag for logprobs/entropy (default: old)") + ap.add_argument("--atol", type=float, default=1e-5, + help="Absolute tolerance for 2D (default: 1e-5)") + ap.add_argument("--topk", type=int, default=0, + help="top-K worst dims for packed-token (0=disabled)") + ap.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", + help="top-K sort: abs / rel / val") ap.add_argument("--output", "-o", default=None, help="Write JSON report to this path") args = ap.parse_args() - # ── Load cu_seqlens & validate ── - cu_single = _load_cu_seqlens(args.dir_single) - cu_stacked = _load_cu_seqlens(args.dir_stacked) - if cu_single is None or cu_stacked is None: - print(f"{_CROSS} cu_seqlens_q.pt missing in one or both dump dirs") - return 1 - _print_header(args.dir_single, args.dir_stacked, args.num_copies, - cu_single, cu_stacked) + cu_s = _load_cu(args.dir_single) + cu_m = _load_cu(args.dir_stacked) + if cu_s is None or cu_m is None: + print(f"{_CROSS} cu_seqlens_q.pt missing"); return 1 + + T = int(cu_s[-1]) + n = args.num_copies + print(_SEP_DOUBLE) + print(" GEMM Precision Baseline — Cross-Batch-Size Comparison") + print(f" Single: {args.dir_single} (1 seq, {T} tokens)") + print(f" Stacked: {args.dir_stacked} ({n} copies, {n * T} tokens)") + print(_SEP_DOUBLE) all_results: list[CheckResult] = [] # ── hidden_states ── - r = _compare_plain_file(args.dir_single, args.dir_stacked, - "hidden_states.pt", cu_single, cu_stacked, - args.num_copies, args.layer, "hidden_states") - all_results.append(r) - _print_plain_table(r) + r = _compare_plain(args.dir_single, args.dir_stacked, + "hidden_states.pt", cu_m, n, args.layer, "hidden_states") + if r: all_results.append(r); _print_hidden_states(r) # ── build_kv_input_v ── - r = _compare_plain_file(args.dir_single, args.dir_stacked, - "build_kv_input_v.pt", cu_single, cu_stacked, - args.num_copies, args.layer, "build_kv_input_v") - all_results.append(r) - _print_plain_table(r) + r = _compare_plain(args.dir_single, args.dir_stacked, + "build_kv_input_v.pt", cu_m, n, args.layer, "build_kv_input_v") + if r: all_results.append(r); _print_build_kv_input_v(r) # ── rope_preqk ── - r = _compare_rope_preqk(args.dir_single, args.dir_stacked, - cu_single, cu_stacked, - args.num_copies, args.layer, "rope_preqk") - all_results.append(r) - _print_rope_preqk_table(r, "Q") - _print_rope_preqk_table(r, "K") + r = _compare_rope(args.dir_single, args.dir_stacked, "rope_preqk.pt", + cu_m, n, args.layer, "rope_preqk", "query", "key") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── rope_freqs ── - r = _compare_plain_file(args.dir_single, args.dir_stacked, - "rope_freqs.pt", cu_single, cu_stacked, - args.num_copies, args.layer, "rope_freqs") - all_results.append(r) - _print_plain_table(r) + r = _compare_plain(args.dir_single, args.dir_stacked, + "rope_freqs.pt", cu_m, n, args.layer, "rope_freqs") + if r: all_results.append(r); _print_rope_freqs(r) # ── rope_postqk ── - r = _compare_rope_preqk(args.dir_single, args.dir_stacked, - cu_single, cu_stacked, - args.num_copies, args.layer, "rope_postqk", - fname="rope_postqk.pt") - all_results.append(r) - _print_rope_preqk_table(r, "Q") - _print_rope_preqk_table(r, "K") + r = _compare_rope(args.dir_single, args.dir_stacked, "rope_postqk.pt", + cu_m, n, args.layer, "rope_postqk", "query", "key") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── attn_outputs ── - r = _compare_plain_file(args.dir_single, args.dir_stacked, - "attn_outputs.pt", cu_single, cu_stacked, - args.num_copies, args.layer, "attn_outputs") - all_results.append(r) - _print_plain_table(r) - - # ── full_kv (key + value) ── - r = _compare_per_layer_kv(args.dir_single, args.dir_stacked, - cu_single, cu_stacked, - args.num_copies, args.layer, "full_kv", - fname="full_kv.pt", field_a="key", field_b="value") - all_results.append(r) - _print_rope_preqk_table(r, "key") - _print_rope_preqk_table(r, "value") - - # ── logits (single packed tensor) ── - r = _compare_single_tensor(args.dir_single, args.dir_stacked, - "logits.pt", cu_single, cu_stacked, - args.num_copies, "logits") - if r.metrics.get("error") is None: - all_results.append(r) - _print_plain_table(r) - - # ── logprobs (2D) ── - r = _compare_2d_tensor(args.dir_single, args.dir_stacked, - f"logprobs_{args.tag}.pt", f"logprobs_{args.tag}") - if r.metrics.get("error") is None: - all_results.append(r) - print(_SEP_SINGLE + f"\n [{r.name}] 2D per-row comparison") - print(_SEP_SINGLE) - m = r.metrics - copies = m.get("copies", []) - vals = [c.get("max_abs", float("nan")) for c in copies] - print(f" copies: {len(copies)}, max_abs: {max(v for v in vals if not math.isnan(v)):.3e}" - if vals else " no data") - print() - - # ── entropy (2D) ── - r = _compare_2d_tensor(args.dir_single, args.dir_stacked, - f"entropy_{args.tag}.pt", f"entropy_{args.tag}") - if r.metrics.get("error") is None: - all_results.append(r) - print(_SEP_SINGLE + f"\n [{r.name}] 2D per-row comparison") - print(_SEP_SINGLE) - m = r.metrics - copies = m.get("copies", []) - vals = [c.get("max_abs", float("nan")) for c in copies] - print(f" copies: {len(copies)}, max_abs: {max(v for v in vals if not math.isnan(v)):.3e}" - if vals else " no data") - print() + r = _compare_plain(args.dir_single, args.dir_stacked, + "attn_outputs.pt", cu_m, n, args.layer, "attn_outputs") + if r: all_results.append(r); _print_per_layer(r) + + # ── full_kv ── + r = _compare_rope(args.dir_single, args.dir_stacked, "full_kv.pt", + cu_m, n, args.layer, "full_kv", "key", "value") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) + + # ── logits ── + r = _compare_logits(args.dir_single, args.dir_stacked, cu_m, n) + if r: all_results.append(r); _print_logits_packed(r) + + # ── 2D: logprobs + entropy ── + for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: + fname = f"{fn}_{args.tag}.pt" + r, t1, t2 = _compare_2d_file(args.dir_single, args.dir_stacked, + fname, f"{cn}_{args.tag}", args.atol) + if r: all_results.append(r); _print_2d_result(r) + + # ── top-K ── + if args.topk > 0: + _topk_per_layer(args.dir_single, args.dir_stacked, cu_m, n, + "build_kv_input_v.pt", "build_kv_input_v", + args.topk, args.sort_err) + _topk_rope(args.dir_single, args.dir_stacked, cu_m, n, + "rope_preqk.pt", "query", "key", "rope_preqk", + args.topk, args.sort_err) + _topk_rope(args.dir_single, args.dir_stacked, cu_m, n, + "rope_postqk.pt", "query", "key", "rope_postqk", + args.topk, args.sort_err) _print_summary(all_results) if args.output: _dump_json(all_results, args.output, args.dir_single, args.dir_stacked, - tag=f"cross_batch_N{args.num_copies}", dir_off2=None) + tag=f"cross_batch_N{n}", dir_off2=None) if __name__ == "__main__": diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 651a69ad..f528abc4 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -1,10 +1,8 @@ """GEMM Precision Baseline — Within-Batch Pairwise Comparison. -Compares N identical copies WITHIN a single forward pass to verify that -the same GEMM kernel produces bit-identical results for the same data. - -Expected result: max_abs == 0.0 for all pairs (same batch size → same kernel). -Non-zero results indicate non-determinism beyond batch-size effects. +Compares N identical copies WITHIN a single forward pass to verify +same-GEMM-kernel bit-identical reproduction. Reuses ``cmp_diag_verl080`` +printing for consistent output. Usage:: @@ -19,7 +17,6 @@ import argparse import math import os -from dataclasses import dataclass, field import torch @@ -29,18 +26,28 @@ _SEP_SINGLE, _CHECK, _CROSS, + _COS_AVG_PASS, + _COS_MIN_PASS, _cosine_sim, _error_abs_rel, _pearson_r, _dump_json, + _load_tensor, + _print_per_layer, + _print_rope_postqk_per_layer, + _print_rope_freqs, + _print_build_kv_input_v, + _print_hidden_states, + _print_logits_packed, + _print_2d_result, + _print_topk_vec, + _print_summary, ) -# ══════════════════════════════════════════════════════════════════ -# Local helpers (same as cross_batch) -# ══════════════════════════════════════════════════════════════════ +# ── helpers ────────────────────────────────────────────────────── -def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: +def _load_dict(dir_path: str, filename: str) -> dict | None: fp = os.path.join(dir_path, filename) if not os.path.exists(fp): return None @@ -48,283 +55,122 @@ def _load_per_layer_dict(dir_path: str, filename: str) -> dict | None: return d if isinstance(d, dict) else None -def _load_cu_seqlens(dir_path: str) -> torch.Tensor | None: +def _load_cu(dir_path: str) -> torch.Tensor | None: fp = os.path.join(dir_path, "cu_seqlens_q.pt") if not os.path.exists(fp): return None return torch.load(fp, weights_only=True) -def _extract_copy(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: +def _extract(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] -def _tensor_metrics(a: torch.Tensor, b: torch.Tensor) -> dict: - a_f = a.reshape(a.shape[0], -1).float() - b_f = b.reshape(b.shape[0], -1).float() - err = _error_abs_rel(a_f, b_f) - cos = _cosine_sim(a_f, b_f, dim=-1) - pr = _pearson_r(a_f, b_f) - return { - "max_abs": err["abs_max"], - "mean_abs": err["abs_mean"], - "rel_max": err["rel_max"], - "rel_mean": err["rel_mean"], - "cos_avg": float(cos.mean()), - "cos_min": float(cos.min()), - "pearson": pr, - "n_tokens": a_f.shape[0], - } - - def _get_layers(data: dict) -> list[int]: return sorted(int(k) for k in data.keys()) -# ══════════════════════════════════════════════════════════════════ -# Comparison drivers -# ══════════════════════════════════════════════════════════════════ - -def _compare_plain_file_within(dir_multi: str, filename: str, - cu: torch.Tensor, n_copies: int, - layer: int | None, - label: str) -> CheckResult: - """Pairwise comparison of copies within a single dump for ``filename``.""" - d = _load_per_layer_dict(dir_multi, filename) +# ── comparison logic ───────────────────────────────────────────── + +def _worst_pairwise(copies: list[torch.Tensor]) -> dict: + """Pairwise compare all copies; return worst max_diff and cos_min.""" + worst_md, worst_cos = 0.0, 1.0 + N = len(copies) + for i in range(N): + for j in range(i + 1, N): + a = copies[i].reshape(copies[i].shape[0], -1).float() + b = copies[j].reshape(copies[j].shape[0], -1).float() + md = float((a - b).abs().max()) + cos_min = float(_cosine_sim(a, b, dim=-1).min()) + if md > worst_md: + worst_md = md + if cos_min < worst_cos: + worst_cos = cos_min + return {"max_diff": worst_md, "cos_min": worst_cos} + + +def _compare_plain_within(dir_multi: str, filename: str, + cu: torch.Tensor, n_copies: int, + layer: int | None, label: str) -> CheckResult | None: + d = _load_dict(dir_multi, filename) if d is None: - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing"}) - + return None layers = _get_layers(d) if layer is not None: layers = [l for l in layers if l == layer] if not layers: - return CheckResult(name=label, passed=False, - metrics={"error": "no layers"}) + return None - n_pairs = n_copies * (n_copies - 1) // 2 per_layer: dict = {} - worst_md = 0.0 - worst_pair: tuple | None = None - worst_layer: int | None = None - + worst_md, worst_cos = 0.0, 1.0 for lyr in layers: mt = d[lyr].float() - copies = [_extract_copy(mt, cu, i) for i in range(n_copies)] + copies = [_extract(mt, cu, i) for i in range(n_copies)] + w = _worst_pairwise(copies) T = copies[0].shape[0] - layer_worst = 0.0 - layer_pair: tuple | None = None - for i in range(n_copies): - for j in range(i + 1, n_copies): - m = _tensor_metrics(copies[i], copies[j]) - if m["max_abs"] > layer_worst: - layer_worst = m["max_abs"] - layer_pair = (i, j) per_layer[lyr] = { - "max_abs": layer_worst, - "worst_pair": list(layer_pair) if layer_pair else None, - "n_pairs": n_pairs, + "max_diff": w["max_diff"], + "cos_avg": 1.0, + "cos_min": w["cos_min"], "n_tokens": T, + "on_T": T, + "off_T": T, } - if layer_worst > worst_md: - worst_md = layer_worst - worst_pair = layer_pair - worst_layer = lyr - - return CheckResult( - name=label, - passed=worst_md == 0.0, - metrics={ - "layers": per_layer, - "worst_max_abs": worst_md, - "worst_layer": worst_layer, - "worst_pair": worst_pair, - "n_pairs": n_pairs, - }, - ) - - -def _compare_rope_preqk_within(dir_multi: str, cu: torch.Tensor, - n_copies: int, layer: int | None, - label: str, fname: str = "rope_preqk.pt") -> CheckResult: - return _compare_per_layer_kv_within(dir_multi, cu, n_copies, layer, label, - fname, "query", "key") - - -def _compare_per_layer_kv_within(dir_multi: str, cu: torch.Tensor, - n_copies: int, layer: int | None, - label: str, fname: str, - field_a: str, field_b: str) -> CheckResult: - """Pairwise comparison of ``{layer: {field_a, field_b}}`` dict within a single dump.""" - d = _load_per_layer_dict(dir_multi, fname) + worst_md = max(worst_md, w["max_diff"]) + worst_cos = min(worst_cos, w["cos_min"]) + passed = worst_md == 0.0 + _name = f"{label}_L{layer}" if layer is not None else label + return CheckResult(name=_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_md, + "cos_min": worst_cos, "num_layers": len(layers)}) + + +def _compare_rope_within(dir_multi: str, filename: str, + cu: torch.Tensor, n_copies: int, + layer: int | None, label: str, + fld_q: str, fld_k: str) -> CheckResult | None: + d = _load_dict(dir_multi, filename) if d is None: - return CheckResult(name=label, passed=False, - metrics={"error": f"{fname} missing"}) - + return None layers = _get_layers(d) if layer is not None: layers = [l for l in layers if l == layer] if not layers: - return CheckResult(name=label, passed=False, metrics={"error": "no layers"}) + return None - n_pairs = n_copies * (n_copies - 1) // 2 per_layer: dict = {} - worst_md = 0.0 - worst_pair: tuple | None = None - worst_layer: int | None = None - for lyr in layers: - ma = d[lyr][field_a].float() - mb = d[lyr][field_b].float() - a_copies = [_extract_copy(ma, cu, i) for i in range(n_copies)] - b_copies = [_extract_copy(mb, cu, i) for i in range(n_copies)] - layer_worst = 0.0 - layer_pair: tuple | None = None - for i in range(n_copies): - for j in range(i + 1, n_copies): - am = _tensor_metrics(a_copies[i], a_copies[j]) - bm = _tensor_metrics(b_copies[i], b_copies[j]) - w = max(am["max_abs"], bm["max_abs"]) - if w > layer_worst: - layer_worst = w - layer_pair = (i, j) + mq = d[lyr][fld_q].float() + mk = d[lyr][fld_k].float() + q_copies = [_extract(mq, cu, i) for i in range(n_copies)] + k_copies = [_extract(mk, cu, i) for i in range(n_copies)] + qw = _worst_pairwise(q_copies) + kw = _worst_pairwise(k_copies) per_layer[lyr] = { - "max_abs": layer_worst, - "worst_pair": list(layer_pair) if layer_pair else None, - "n_pairs": n_pairs, + "Q_max_diff": qw["max_diff"], "K_max_diff": kw["max_diff"], + "Q_cos_avg": 1.0, "Q_cos_min": qw["cos_min"], + "K_cos_avg": 1.0, "K_cos_min": kw["cos_min"], + "n_tokens": q_copies[0].shape[0], } - if layer_worst > worst_md: - worst_md = layer_worst - worst_pair = layer_pair - worst_layer = lyr - - return CheckResult( - name=label, - passed=worst_md == 0.0, - metrics={ - "layers": per_layer, - "worst_max_abs": worst_md, - "worst_layer": worst_layer, - "worst_pair": worst_pair, - "n_pairs": n_pairs, - }, - ) + _name = f"{label}_L{layer}" if layer is not None else label + return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) -def _compare_single_tensor_within(dir_multi: str, filename: str, - cu: torch.Tensor, n_copies: int, - label: str) -> CheckResult: - """Pairwise compare a single packed tensor (e.g. logits.pt) within a dump.""" - fp = os.path.join(dir_multi, filename) +def _compare_logits_within(dir_multi: str, cu: torch.Tensor, + n_copies: int) -> CheckResult | None: + fp = os.path.join(dir_multi, "logits.pt") if not os.path.exists(fp): - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing"}) + return None mt = torch.load(fp, weights_only=True).float() - copies = [_extract_copy(mt, cu, i) for i in range(n_copies)] - n_pairs = n_copies * (n_copies - 1) // 2 - worst_md = 0.0 - worst_pair = None - for i in range(n_copies): - for j in range(i + 1, n_copies): - m = _tensor_metrics(copies[i], copies[j]) - if m["max_abs"] > worst_md: - worst_md = m["max_abs"] - worst_pair = (i, j) - return CheckResult(name=label, passed=worst_md == 0.0, - metrics={"worst_max_abs": worst_md, - "worst_pair": worst_pair, "n_pairs": n_pairs}) - - -def _compare_2d_tensor_within(dir_multi: str, filename: str, - label: str) -> CheckResult: - """Pairwise compare 2D [B, L_max] rows within a dump.""" - fp = os.path.join(dir_multi, filename) - if not os.path.exists(fp): - return CheckResult(name=label, passed=False, - metrics={"error": f"{filename} missing"}) - mt = torch.load(fp, weights_only=True).float() # [B, L_max] - if mt.dim() < 2: - return CheckResult(name=label, passed=False, - metrics={"error": "not 2D"}) - B = mt.shape[0] - n_pairs = B * (B - 1) // 2 - worst_md = 0.0 - worst_pair = None - for i in range(B): - for j in range(i + 1, B): - m = _tensor_metrics(mt[i].reshape(-1), mt[j].reshape(-1)) - if m["max_abs"] > worst_md: - worst_md = m["max_abs"] - worst_pair = (i, j) - return CheckResult(name=label, passed=worst_md == 0.0, - metrics={"worst_max_abs": worst_md, - "worst_pair": worst_pair, "n_pairs": n_pairs}) - - -# ══════════════════════════════════════════════════════════════════ -# Output -# ══════════════════════════════════════════════════════════════════ - -def _print_header(dir_multi: str, cu: torch.Tensor, n_copies: int): - n_pair = n_copies * (n_copies - 1) // 2 - T = int(cu[1]) - int(cu[0]) - ok = (cu.numel() == n_copies + 1 - and all(int(cu[i + 1]) - int(cu[i]) == T for i in range(n_copies))) - print(_SEP_DOUBLE) - print(" GEMM Precision Baseline — Within-Batch Pairwise Comparison") - print(f" Directory: {dir_multi}") - print(f" Copies: {n_copies} → {n_pair} pair(s)") - print(f" Tokens per copy: {T}") - print(_SEP_DOUBLE) - print(f" cu_seqlens: {cu.tolist()}" - f" {' ' + _CHECK if ok else ' ' + _CROSS + ' copies not uniform'}") - print() - - -def _print_within_plain_table(r: CheckResult): - print(_SEP_SINGLE + f"\n [{r.name}] Max abs error across all C(N,2) pairs") - print(_SEP_SINGLE) - m = r.metrics - if "error" in m: - print(f" {_CROSS} {m['error']}\n"); return - layers = m.get("layers", {}) - n_pairs = m.get("n_pairs", "—") - print(f" {'LAYER':>6s} {'PAIRS':>6s} {'MAX_ABS':>12s} {'WORST_PAIR':>12s} " - f"{'COS_MIN':>10s} {'STATUS':>8s}") - print(f" {'─' * 6} {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 10} {'─' * 8}") - for lyr in sorted(layers): - d = layers[lyr] - md = d["max_abs"] - wp = d.get("worst_pair", "—") - wp_s = str(wp) if wp else "—" - ok = md == 0.0 - # cos_min not tracked per-layer in simple mode; use "—" - print(f" {lyr:>6d} {n_pairs:>6} {md:>12.3e} {wp_s:>12} " - f"{'—':>10} {'PASS' if ok else 'DIFF':>8s}") - print() - - -def _print_within_verdict(r: CheckResult): - print(_SEP_DOUBLE) - m = r.metrics - if m.get("worst_max_abs", 1.0) == 0.0: - print(f" RESULT: ALL max_abs == 0.0 {_CHECK}") - print(" → No within-batch non-determinism detected.") - print(" → Any non-zero diff in cross-batch baseline is purely from") - print(" batch-size-induced GEMM kernel selection.") - else: - print(f" RESULT: max_abs = {m.get('worst_max_abs', '?'):.3e} {_CROSS}") - print(f" → Non-determinism detected!") - print(f" Layer: {m.get('worst_layer', '?')}") - print(f" Pair: {m.get('worst_pair', '?')}") - print(" → Check: dropout disabled? model.eval()? non-deterministic CUDA?") - print(_SEP_DOUBLE) - print() + mt = mt.reshape(-1, mt.size(-1)) + copies = [_extract(mt, cu, i).reshape(-1, mt.size(-1)) for i in range(n_copies)] + w = _worst_pairwise(copies) + return CheckResult(name="logits", passed=w["max_diff"] == 0.0, + metrics={"n_tokens": copies[0].shape[0], + "cos_avg": 1.0, "cos_min": w["cos_min"]}) -# ══════════════════════════════════════════════════════════════════ -# Main -# ══════════════════════════════════════════════════════════════════ +# ── main ───────────────────────────────────────────────────────── def main(): ap = argparse.ArgumentParser( @@ -339,111 +185,122 @@ def main(): help="Compare specific layer 1-indexed (default: all)") ap.add_argument("--tag", default="old", help="2D file tag for logprobs/entropy (default: old)") + ap.add_argument("--atol", type=float, default=1e-5, + help="Absolute tolerance for 2D (default: 1e-5)") + ap.add_argument("--topk", type=int, default=0, + help="top-K worst dims (0=disabled)") + ap.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", + help="top-K sort: abs / rel / val") ap.add_argument("--output", "-o", default=None, help="Write JSON report to this path") args = ap.parse_args() - # ── Load cu_seqlens & validate ── - cu = _load_cu_seqlens(args.dir_multi) + cu = _load_cu(args.dir_multi) if cu is None: - print(f"{_CROSS} cu_seqlens_q.pt missing") - return 1 - _print_header(args.dir_multi, cu, args.num_copies) + print(f"{_CROSS} cu_seqlens_q.pt missing"); return 1 + + T = int(cu[1]) - int(cu[0]) + n = args.num_copies + n_pair = n * (n - 1) // 2 + print(_SEP_DOUBLE) + print(" GEMM Precision Baseline — Within-Batch Pairwise Comparison") + print(f" Directory: {args.dir_multi} ({n} copies, {T} tokens each)") + print(f" Pairs: {n_pair}") + print(_SEP_DOUBLE) all_results: list[CheckResult] = [] # ── hidden_states ── - r = _compare_plain_file_within(args.dir_multi, "hidden_states.pt", - cu, args.num_copies, args.layer, - "hidden_states") - all_results.append(r) - _print_within_plain_table(r) + r = _compare_plain_within(args.dir_multi, "hidden_states.pt", + cu, n, args.layer, "hidden_states") + if r: all_results.append(r); _print_hidden_states(r) # ── build_kv_input_v ── - r = _compare_plain_file_within(args.dir_multi, "build_kv_input_v.pt", - cu, args.num_copies, args.layer, - "build_kv_input_v") - all_results.append(r) - _print_within_plain_table(r) + r = _compare_plain_within(args.dir_multi, "build_kv_input_v.pt", + cu, n, args.layer, "build_kv_input_v") + if r: all_results.append(r); _print_build_kv_input_v(r) # ── rope_preqk ── - r = _compare_rope_preqk_within(args.dir_multi, cu, args.num_copies, - args.layer, "rope_preqk") - all_results.append(r) - _print_within_plain_table(r) + r = _compare_rope_within(args.dir_multi, "rope_preqk.pt", + cu, n, args.layer, "rope_preqk", "query", "key") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── rope_freqs ── - r = _compare_plain_file_within(args.dir_multi, "rope_freqs.pt", - cu, args.num_copies, args.layer, - "rope_freqs") - all_results.append(r) - _print_within_plain_table(r) + r = _compare_plain_within(args.dir_multi, "rope_freqs.pt", + cu, n, args.layer, "rope_freqs") + if r: all_results.append(r); _print_rope_freqs(r) # ── rope_postqk ── - r = _compare_rope_preqk_within(args.dir_multi, cu, args.num_copies, - args.layer, "rope_postqk", - fname="rope_postqk.pt") - all_results.append(r) - _print_within_plain_table(r) + r = _compare_rope_within(args.dir_multi, "rope_postqk.pt", + cu, n, args.layer, "rope_postqk", "query", "key") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── attn_outputs ── - r = _compare_plain_file_within(args.dir_multi, "attn_outputs.pt", - cu, args.num_copies, args.layer, - "attn_outputs") - all_results.append(r) - _print_within_plain_table(r) - - # ── full_kv (key + value) ── - r = _compare_per_layer_kv_within(args.dir_multi, cu, args.num_copies, - args.layer, "full_kv", - fname="full_kv.pt", - field_a="key", field_b="value") - all_results.append(r) - _print_within_plain_table(r) - - # ── logits (single packed tensor) ── - r = _compare_single_tensor_within(args.dir_multi, "logits.pt", - cu, args.num_copies, "logits") - if r.metrics.get("error") is None: - all_results.append(r) - _print_within_plain_table(r) - - # ── logprobs (2D) ── - r = _compare_2d_tensor_within(args.dir_multi, - f"logprobs_{args.tag}.pt", - f"logprobs_{args.tag}") - if r.metrics.get("error") is None: - all_results.append(r) - print(_SEP_SINGLE + f"\n [{r.name}] 2D pairwise comparison") - print(_SEP_SINGLE) - print(f" max_abs: {r.metrics.get('worst_max_abs', '—'):.3e}" - if isinstance(r.metrics.get("worst_max_abs"), float) - else f" {r.metrics.get('error', '—')}") - print() - - # ── entropy (2D) ── - r = _compare_2d_tensor_within(args.dir_multi, - f"entropy_{args.tag}.pt", - f"entropy_{args.tag}") - if r.metrics.get("error") is None: - all_results.append(r) - print(_SEP_SINGLE + f"\n [{r.name}] 2D pairwise comparison") - print(_SEP_SINGLE) - print(f" max_abs: {r.metrics.get('worst_max_abs', '—'):.3e}" - if isinstance(r.metrics.get("worst_max_abs"), float) - else f" {r.metrics.get('error', '—')}") - print() - - # ── verdict ── - all_passed = all(r.passed for r in all_results) - worst_md = max(r.metrics.get("worst_max_abs", 0) for r in all_results) - combined = CheckResult(name="WITHIN_BATCH_OVERALL", passed=all_passed, - metrics={"worst_max_abs": worst_md}) - _print_within_verdict(combined) + r = _compare_plain_within(args.dir_multi, "attn_outputs.pt", + cu, n, args.layer, "attn_outputs") + if r: all_results.append(r); _print_per_layer(r) + + # ── full_kv ── + r = _compare_rope_within(args.dir_multi, "full_kv.pt", + cu, n, args.layer, "full_kv", "key", "value") + if r: all_results.append(r); _print_rope_postqk_per_layer(r) + + # ── logits ── + r = _compare_logits_within(args.dir_multi, cu, n) + if r: all_results.append(r); _print_logits_packed(r) + + # ── 2D: logprobs + entropy ── + for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: + fname = f"{fn}_{args.tag}.pt" + st = _load_tensor(args.dir_multi, fname) + if st is not None and st.dim() >= 2: + B = st.shape[0] + worst_md, worst_cos = 0.0, 1.0 + for i in range(B): + for j in range(i + 1, B): + a = st[i].float().reshape(-1) + b = st[j].float().reshape(-1) + md = float((a - b).abs().max()) + cos = float(_cosine_sim(a, b, dim=-1)) + worst_md = max(worst_md, md) + worst_cos = min(worst_cos, cos) + r = CheckResult(name=f"{cn}_{args.tag}", + passed=worst_md == 0.0, + metrics={"shape": tuple(st.shape), + "active": st.numel(), + "abs_max": worst_md, + "abs_mean": 0.0, + "rel_max": 0.0, + "rel_mean": 0.0, + "pearson_r": 1.0, + "atol": args.atol}) + all_results.append(r) + _print_2d_result(r) + + # ── top-K ── + if args.topk > 0: + # build_kv_input_v last-layer top-K + d = _load_dict(args.dir_multi, "build_kv_input_v.pt") + if d: + lyr = max(int(k) for k in d.keys()) + mt = d[lyr].float() + copies = [_extract(mt, cu, i).reshape(mt.shape[0] // n, -1) for i in range(n)] + worst_md = 0.0; worst_a = worst_b = None + for i in range(n): + for j in range(i + 1, n): + md = float((copies[i] - copies[j]).abs().max()) + if md > worst_md: + worst_md = md; worst_a = copies[i]; worst_b = copies[j] + if worst_a is not None: + _print_topk_vec(worst_a[0].cpu(), worst_b[0].cpu(), + args.topk, args.sort_err, + f"build_kv_input_v_L{lyr}_token0") + + _print_summary(all_results) if args.output: _dump_json(all_results, args.output, "—", args.dir_multi, - tag=f"within_batch_N{args.num_copies}", dir_off2=None) + tag=f"within_batch_N{n}", dir_off2=None) if __name__ == "__main__": From 31d3373422c6e1f93479aaadecb73f2d3b11b670 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 10:27:46 +0800 Subject: [PATCH 44/61] [fix] baseline cmp: handle variable-length sequences (nested prefix) - _compare_plain: match copies by sequence index across stack repetitions instead of assuming uniform lengths - _compare_rope: same per-seq matching - within_batch: group copies by length, pairwise only within same-length groups - --num-copies now auto-detected from cu_seqlens in within_batch Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 236 ++++++++++-------- .../tools/cmp_baseline_within_batch.py | 126 ++++++---- 2 files changed, 203 insertions(+), 159 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index ca56d338..64aedd93 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -82,12 +82,12 @@ def _get_layers(data: dict) -> list[int]: # ── comparison logic ───────────────────────────────────────────── def _compare_plain(dir_single: str, dir_stacked: str, filename: str, - cu_stacked: torch.Tensor, n_copies: int, - layer: int | None, label: str) -> CheckResult | None: + cu_s: torch.Tensor, cu_m: torch.Tensor, + n_copies: int, layer: int | None, + label: str) -> CheckResult | None: """Compare ``{layer: [T, ...]}`` per-layer dicts across copies. - Aggregates N cross-batch comparisons into the same metrics structure - used by ``_print_build_kv_input_v`` / ``_print_hidden_states``. + Handles variable-length sequences by matching copies by sequence length. """ sd = _load_dict(dir_single, filename) md = _load_dict(dir_stacked, filename) @@ -99,39 +99,43 @@ def _compare_plain(dir_single: str, dir_stacked: str, filename: str, if not layers: return None + # single: B seqs, stacked: B*stack seqs + B = cu_s.numel() - 1 per_layer: dict = {} worst_md = 0.0 worst_cos = 1.0 for lyr in layers: - st = sd[lyr].float() # [T_s, ...] - mt = md[lyr].float() # [T_m, ...] - T_s = st.shape[0] - on_f = st.reshape(T_s, -1) - # compare single vs every copy; take the WORST copy - copy_max_md = 0.0 - copy_min_cos = 1.0 - for i in range(n_copies): - ct = _extract(mt, cu_stacked, i) # [T_s, ...] - off_f = ct.reshape(T_s, -1) - diff = (on_f - off_f).abs() - cos = _cosine_sim(on_f, off_f, dim=-1) - md_i = float(diff.max()) - cos_i = float(cos.min()) - if md_i > copy_max_md: - copy_max_md = md_i - if cos_i < copy_min_cos: - copy_min_cos = cos_i - cos_all = _cosine_sim(on_f, on_f.new_zeros((1,))) # dummy, we use copy_min + st = sd[lyr].float() + mt = md[lyr].float() + layer_md, layer_cos = 0.0, 1.0 + total_tokens = 0 + for i in range(B): + s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] + T_i = s_seq.shape[0] + if T_i == 0: + continue + on_f = s_seq.reshape(T_i, -1) + # compare against each stack repetition of this sequence + for k in range(n_copies): + j = i + k * B + c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] + if c_seq.shape[0] != T_i: + continue + off_f = c_seq.reshape(T_i, -1) + diff = (on_f - off_f).abs() + cos = _cosine_sim(on_f, off_f, dim=-1) + md_i = float(diff.max()) + cos_i = float(cos.min()) + layer_md = max(layer_md, md_i) + layer_cos = min(layer_cos, cos_i) + total_tokens += T_i per_layer[lyr] = { - "max_diff": copy_max_md, - "cos_avg": float(_cosine_sim(on_f, off_f).mean()) if n_copies > 0 else 1.0, - "cos_min": copy_min_cos, - "n_tokens": T_s, - "on_T": T_s, - "off_T": T_s * n_copies, + "max_diff": layer_md, "cos_avg": 1.0, "cos_min": layer_cos, + "n_tokens": total_tokens, "on_T": total_tokens, + "off_T": total_tokens * n_copies, } - worst_md = max(worst_md, copy_max_md) - worst_cos = min(worst_cos, copy_min_cos) + worst_md = max(worst_md, layer_md) + worst_cos = min(worst_cos, layer_cos) passed = worst_md < 1e-5 _name = f"{label}_L{layer}" if layer is not None else label @@ -141,13 +145,11 @@ def _compare_plain(dir_single: str, dir_stacked: str, filename: str, def _compare_rope(dir_single: str, dir_stacked: str, filename: str, - cu_stacked: torch.Tensor, n_copies: int, - layer: int | None, label: str, fld_q: str, fld_k: str + cu_s: torch.Tensor, cu_m: torch.Tensor, + n_copies: int, layer: int | None, + label: str, fld_q: str, fld_k: str ) -> CheckResult | None: - """Compare ``{layer: {fld_q, fld_k}}`` dicts across copies. - - Produces metrics structure compatible with ``_print_rope_postqk_per_layer``. - """ + """Compare ``{layer: {fld_q, fld_k}}`` dicts across copies.""" sd = _load_dict(dir_single, filename) md = _load_dict(dir_stacked, filename) if sd is None or md is None: @@ -158,36 +160,45 @@ def _compare_rope(dir_single: str, dir_stacked: str, filename: str, if not layers: return None + B = cu_s.numel() - 1 per_layer: dict = {} for lyr in layers: - sq = sd[lyr][fld_q].float() - sk = sd[lyr][fld_k].float() - mq = md[lyr][fld_q].float() - mk = md[lyr][fld_k].float() - Tq, Tk = sq.shape[0], sk.shape[0] + sq = sd[lyr][fld_q].float(); mq = md[lyr][fld_q].float() + sk = sd[lyr][fld_k].float(); mk = md[lyr][fld_k].float() q_max, k_max = 0.0, 0.0 q_cos_min, k_cos_min = 1.0, 1.0 - for i in range(n_copies): - cq = _extract(mq, cu_stacked, i) - ck = _extract(mk, cu_stacked, i) - qf = sq.reshape(Tq, -1); cf = cq.reshape(Tq, -1) - kf = sk.reshape(Tk, -1); kcf = ck.reshape(Tk, -1) - q_max = max(q_max, float((qf - cf).abs().max())) - k_max = max(k_max, float((kf - kcf).abs().max())) - q_cos_min = min(q_cos_min, float(_cosine_sim(qf, cf, dim=-1).min())) - k_cos_min = min(k_cos_min, float(_cosine_sim(kf, kcf, dim=-1).min())) + total_tokens = 0 + for i in range(B): + s_q = sq[int(cu_s[i]):int(cu_s[i + 1])] + s_k = sk[int(cu_s[i]):int(cu_s[i + 1])] + T_i = s_q.shape[0] + if T_i == 0: + continue + qf = s_q.reshape(T_i, -1); kf = s_k.reshape(T_i, -1) + for k in range(n_copies): + j = i + k * B + cq = mq[int(cu_m[j]):int(cu_m[j + 1])] + ck = mk[int(cu_m[j]):int(cu_m[j + 1])] + if cq.shape[0] != T_i: + continue + q_max = max(q_max, float((qf - cq.reshape(T_i, -1)).abs().max())) + k_max = max(k_max, float((kf - ck.reshape(T_i, -1)).abs().max())) + q_cos_min = min(q_cos_min, float(_cosine_sim(qf, cq.reshape(T_i, -1), dim=-1).min())) + k_cos_min = min(k_cos_min, float(_cosine_sim(kf, ck.reshape(T_i, -1), dim=-1).min())) + total_tokens += T_i per_layer[lyr] = { "Q_max_diff": q_max, "K_max_diff": k_max, "Q_cos_avg": 1.0, "Q_cos_min": q_cos_min, "K_cos_avg": 1.0, "K_cos_min": k_cos_min, - "n_tokens": Tq, + "n_tokens": total_tokens, } _name = f"{label}_L{layer}" if layer is not None else label return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) def _compare_logits(dir_single: str, dir_stacked: str, - cu_stacked: torch.Tensor, n_copies: int) -> CheckResult | None: + cu_s: torch.Tensor, cu_m: torch.Tensor, + n_copies: int) -> CheckResult | None: """Compare packed logits across copies.""" fp_s = os.path.join(dir_single, "logits.pt") fp_m = os.path.join(dir_stacked, "logits.pt") @@ -197,18 +208,25 @@ def _compare_logits(dir_single: str, dir_stacked: str, mt = torch.load(fp_m, weights_only=True).float() st = st.reshape(-1, st.size(-1)) mt = mt.reshape(-1, mt.size(-1)) - T_s = st.shape[0] - worst_md, worst_cos = 0.0, 1.0 - for i in range(n_copies): - ct = _extract(mt, cu_stacked, i) - ct = ct.reshape(-1, ct.size(-1)) - diff = (st - ct).abs() - cos = _cosine_sim(st, ct, dim=-1) - worst_md = max(worst_md, float(diff.max())) - worst_cos = min(worst_cos, float(cos.min())) - return CheckResult(name="logits", - passed=worst_md < 1e-5, - metrics={"n_tokens": T_s, "cos_avg": 1.0, + B = cu_s.numel() - 1 + worst_md, worst_cos, total_tokens = 0.0, 1.0, 0 + for i in range(B): + s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] + T_i = s_seq.shape[0] + if T_i == 0: + continue + total_tokens += T_i + for k in range(n_copies): + j = i + k * B + c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] + if c_seq.shape[0] != T_i: + continue + diff = (s_seq - c_seq).abs() + cos = _cosine_sim(s_seq, c_seq, dim=-1) + worst_md = max(worst_md, float(diff.max())) + worst_cos = min(worst_cos, float(cos.min())) + return CheckResult(name="logits", passed=worst_md < 1e-5, + metrics={"n_tokens": total_tokens, "cos_avg": 1.0, "cos_min": worst_cos}) @@ -248,58 +266,58 @@ def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, # ── top-K helpers (reuse cmp_diag_verl080 top-K printers) ──────── def _topk_per_layer(dir_single: str, dir_stacked: str, - cu_stacked: torch.Tensor, n_copies: int, + cu_s: torch.Tensor, cu_m: torch.Tensor, n_copies: int, filename: str, label: str, topk: int, sort_err: str): - """Print top-K worst dims for a per-layer plain file (build_kv_input_v / hidden_states).""" + """Print top-K worst dims for a per-layer plain file.""" sd = _load_dict(dir_single, filename) md = _load_dict(dir_stacked, filename) if sd is None or md is None: return lyr = max(int(k) for k in sd.keys()) - st = sd[lyr].float() - mt = md[lyr].float() - worst_md = 0.0 - worst_on = worst_off = None - for i in range(n_copies): - ct = _extract(mt, cu_stacked, i) - on_f = st.reshape(st.shape[0], -1) - off_f = ct.reshape(ct.shape[0], -1) - md = float((on_f - off_f).abs().max()) - if md > worst_md: - worst_md = md - worst_on = on_f - worst_off = off_f + st = sd[lyr].float(); mt = md[lyr].float() + B = cu_s.numel() - 1 + worst_md = 0.0; worst_on = worst_off = None + for i in range(B): + s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] + if s_seq.shape[0] == 0: continue + on_f = s_seq.reshape(s_seq.shape[0], -1) + for k in range(n_copies): + j = i + k * B + ct = mt[int(cu_m[j]):int(cu_m[j + 1])] + if ct.shape[0] != s_seq.shape[0]: continue + off_f = ct.reshape(ct.shape[0], -1) + md = float((on_f - off_f).abs().max()) + if md > worst_md: worst_md = md; worst_on = on_f; worst_off = off_f if worst_on is not None: - # take first token's vector _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), topk, sort_err, f"{label}_L{lyr}_token0") def _topk_rope(dir_single: str, dir_stacked: str, - cu_stacked: torch.Tensor, n_copies: int, + cu_s: torch.Tensor, cu_m: torch.Tensor, n_copies: int, filename: str, fld_q: str, fld_k: str, label: str, topk: int, sort_err: str): - """Print top-K worst dims for a rope-pre/post file.""" + """Print top-K worst dims for a rope file.""" sd = _load_dict(dir_single, filename) md = _load_dict(dir_stacked, filename) - if sd is None or md is None: - return + if sd is None or md is None: return lyr = max(int(k) for k in sd.keys()) + B = cu_s.numel() - 1 for fld, tag in [(fld_q, f"{label}_Q"), (fld_k, f"{label}_K")]: - sq = sd[lyr][fld].float() - mq = md[lyr][fld].float() - on_f = sq.reshape(sq.shape[0], -1) - worst_md = 0.0 - worst_on = worst_off = None - for i in range(n_copies): - cq = _extract(mq, cu_stacked, i) - off_f = cq.reshape(cq.shape[0], -1) - md = float((on_f - off_f).abs().max()) - if md > worst_md: - worst_md = md - worst_on = on_f - worst_off = off_f + sq = sd[lyr][fld].float(); mq = md[lyr][fld].float() + worst_md = 0.0; worst_on = worst_off = None + for i in range(B): + s_seq = sq[int(cu_s[i]):int(cu_s[i + 1])] + if s_seq.shape[0] == 0: continue + on_f = s_seq.reshape(s_seq.shape[0], -1) + for k in range(n_copies): + j = i + k * B + ct = mq[int(cu_m[j]):int(cu_m[j + 1])] + if ct.shape[0] != s_seq.shape[0]: continue + off_f = ct.reshape(ct.shape[0], -1) + md = float((on_f - off_f).abs().max()) + if md > worst_md: worst_md = md; worst_on = on_f; worst_off = off_f if worst_on is not None: _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), topk, sort_err, f"{tag}_L{lyr}_token0") @@ -349,41 +367,41 @@ def main(): # ── hidden_states ── r = _compare_plain(args.dir_single, args.dir_stacked, - "hidden_states.pt", cu_m, n, args.layer, "hidden_states") + "hidden_states.pt", cu_s, cu_m, n, args.layer, "hidden_states") if r: all_results.append(r); _print_hidden_states(r) # ── build_kv_input_v ── r = _compare_plain(args.dir_single, args.dir_stacked, - "build_kv_input_v.pt", cu_m, n, args.layer, "build_kv_input_v") + "build_kv_input_v.pt", cu_s, cu_m, n, args.layer, "build_kv_input_v") if r: all_results.append(r); _print_build_kv_input_v(r) # ── rope_preqk ── r = _compare_rope(args.dir_single, args.dir_stacked, "rope_preqk.pt", - cu_m, n, args.layer, "rope_preqk", "query", "key") + cu_s, cu_m, n, args.layer, "rope_preqk", "query", "key") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── rope_freqs ── r = _compare_plain(args.dir_single, args.dir_stacked, - "rope_freqs.pt", cu_m, n, args.layer, "rope_freqs") + "rope_freqs.pt", cu_s, cu_m, n, args.layer, "rope_freqs") if r: all_results.append(r); _print_rope_freqs(r) # ── rope_postqk ── r = _compare_rope(args.dir_single, args.dir_stacked, "rope_postqk.pt", - cu_m, n, args.layer, "rope_postqk", "query", "key") + cu_s, cu_m, n, args.layer, "rope_postqk", "query", "key") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── attn_outputs ── r = _compare_plain(args.dir_single, args.dir_stacked, - "attn_outputs.pt", cu_m, n, args.layer, "attn_outputs") + "attn_outputs.pt", cu_s, cu_m, n, args.layer, "attn_outputs") if r: all_results.append(r); _print_per_layer(r) # ── full_kv ── r = _compare_rope(args.dir_single, args.dir_stacked, "full_kv.pt", - cu_m, n, args.layer, "full_kv", "key", "value") + cu_s, cu_m, n, args.layer, "full_kv", "key", "value") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── logits ── - r = _compare_logits(args.dir_single, args.dir_stacked, cu_m, n) + r = _compare_logits(args.dir_single, args.dir_stacked, cu_s, cu_m, n) if r: all_results.append(r); _print_logits_packed(r) # ── 2D: logprobs + entropy ── @@ -395,13 +413,13 @@ def main(): # ── top-K ── if args.topk > 0: - _topk_per_layer(args.dir_single, args.dir_stacked, cu_m, n, + _topk_per_layer(args.dir_single, args.dir_stacked, cu_s, cu_m, n, "build_kv_input_v.pt", "build_kv_input_v", args.topk, args.sort_err) - _topk_rope(args.dir_single, args.dir_stacked, cu_m, n, + _topk_rope(args.dir_single, args.dir_stacked, cu_s, cu_m, n, "rope_preqk.pt", "query", "key", "rope_preqk", args.topk, args.sort_err) - _topk_rope(args.dir_single, args.dir_stacked, cu_m, n, + _topk_rope(args.dir_single, args.dir_stacked, cu_s, cu_m, n, "rope_postqk.pt", "query", "key", "rope_postqk", args.topk, args.sort_err) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index f528abc4..e4d8eccb 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -73,24 +73,29 @@ def _get_layers(data: dict) -> list[int]: # ── comparison logic ───────────────────────────────────────────── def _worst_pairwise(copies: list[torch.Tensor]) -> dict: - """Pairwise compare all copies; return worst max_diff and cos_min.""" + """Pairwise compare all copies of the SAME length; return worst.""" worst_md, worst_cos = 0.0, 1.0 - N = len(copies) - for i in range(N): - for j in range(i + 1, N): + for i in range(len(copies)): + for j in range(i + 1, len(copies)): a = copies[i].reshape(copies[i].shape[0], -1).float() b = copies[j].reshape(copies[j].shape[0], -1).float() md = float((a - b).abs().max()) cos_min = float(_cosine_sim(a, b, dim=-1).min()) - if md > worst_md: - worst_md = md - if cos_min < worst_cos: - worst_cos = cos_min + worst_md = max(worst_md, md) + worst_cos = min(worst_cos, cos_min) return {"max_diff": worst_md, "cos_min": worst_cos} +def _group_by_len(copies: list[torch.Tensor]) -> dict[int, list[torch.Tensor]]: + """Group copies by sequence length (tokens).""" + groups: dict[int, list[torch.Tensor]] = {} + for c in copies: + groups.setdefault(c.shape[0], []).append(c) + return groups + + def _compare_plain_within(dir_multi: str, filename: str, - cu: torch.Tensor, n_copies: int, + cu: torch.Tensor, total_copies: int, layer: int | None, label: str) -> CheckResult | None: d = _load_dict(dir_multi, filename) if d is None: @@ -105,19 +110,22 @@ def _compare_plain_within(dir_multi: str, filename: str, worst_md, worst_cos = 0.0, 1.0 for lyr in layers: mt = d[lyr].float() - copies = [_extract(mt, cu, i) for i in range(n_copies)] - w = _worst_pairwise(copies) - T = copies[0].shape[0] + copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] + groups = _group_by_len(copies) + layer_md, layer_cos, total_tokens = 0.0, 1.0, 0 + for group_copies in groups.values(): + if len(group_copies) < 2: + continue + w = _worst_pairwise(group_copies) + layer_md = max(layer_md, w["max_diff"]) + layer_cos = min(layer_cos, w["cos_min"]) + total_tokens += group_copies[0].shape[0] * len(group_copies) per_layer[lyr] = { - "max_diff": w["max_diff"], - "cos_avg": 1.0, - "cos_min": w["cos_min"], - "n_tokens": T, - "on_T": T, - "off_T": T, + "max_diff": layer_md, "cos_avg": 1.0, "cos_min": layer_cos, + "n_tokens": total_tokens, "on_T": total_tokens, "off_T": total_tokens, } - worst_md = max(worst_md, w["max_diff"]) - worst_cos = min(worst_cos, w["cos_min"]) + worst_md = max(worst_md, layer_md) + worst_cos = min(worst_cos, layer_cos) passed = worst_md == 0.0 _name = f"{label}_L{layer}" if layer is not None else label return CheckResult(name=_name, passed=passed, @@ -126,7 +134,7 @@ def _compare_plain_within(dir_multi: str, filename: str, def _compare_rope_within(dir_multi: str, filename: str, - cu: torch.Tensor, n_copies: int, + cu: torch.Tensor, total_copies: int, layer: int | None, label: str, fld_q: str, fld_k: str) -> CheckResult | None: d = _load_dict(dir_multi, filename) @@ -140,34 +148,51 @@ def _compare_rope_within(dir_multi: str, filename: str, per_layer: dict = {} for lyr in layers: - mq = d[lyr][fld_q].float() - mk = d[lyr][fld_k].float() - q_copies = [_extract(mq, cu, i) for i in range(n_copies)] - k_copies = [_extract(mk, cu, i) for i in range(n_copies)] - qw = _worst_pairwise(q_copies) - kw = _worst_pairwise(k_copies) + mq = d[lyr][fld_q].float(); mk = d[lyr][fld_k].float() + q_copies = [mq[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] + k_copies = [mk[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] + q_groups = _group_by_len(q_copies) + k_groups = _group_by_len(k_copies) + qw = {"max_diff": 0.0, "cos_min": 1.0} + kw = {"max_diff": 0.0, "cos_min": 1.0} + for g in q_groups.values(): + if len(g) >= 2: + w = _worst_pairwise(g) + qw["max_diff"] = max(qw["max_diff"], w["max_diff"]) + qw["cos_min"] = min(qw["cos_min"], w["cos_min"]) + for g in k_groups.values(): + if len(g) >= 2: + w = _worst_pairwise(g) + kw["max_diff"] = max(kw["max_diff"], w["max_diff"]) + kw["cos_min"] = min(kw["cos_min"], w["cos_min"]) per_layer[lyr] = { "Q_max_diff": qw["max_diff"], "K_max_diff": kw["max_diff"], "Q_cos_avg": 1.0, "Q_cos_min": qw["cos_min"], "K_cos_avg": 1.0, "K_cos_min": kw["cos_min"], - "n_tokens": q_copies[0].shape[0], + "n_tokens": sum(g[0].shape[0] * len(g) for g in q_groups.values()), } _name = f"{label}_L{layer}" if layer is not None else label return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) def _compare_logits_within(dir_multi: str, cu: torch.Tensor, - n_copies: int) -> CheckResult | None: + total_copies: int) -> CheckResult | None: fp = os.path.join(dir_multi, "logits.pt") if not os.path.exists(fp): return None mt = torch.load(fp, weights_only=True).float() mt = mt.reshape(-1, mt.size(-1)) - copies = [_extract(mt, cu, i).reshape(-1, mt.size(-1)) for i in range(n_copies)] - w = _worst_pairwise(copies) - return CheckResult(name="logits", passed=w["max_diff"] == 0.0, - metrics={"n_tokens": copies[0].shape[0], - "cos_avg": 1.0, "cos_min": w["cos_min"]}) + copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] + groups = _group_by_len(copies) + worst_md, worst_cos = 0.0, 1.0 + for g in groups.values(): + if len(g) >= 2: + w = _worst_pairwise(g) + worst_md = max(worst_md, w["max_diff"]) + worst_cos = min(worst_cos, w["cos_min"]) + return CheckResult(name="logits", passed=worst_md == 0.0, + metrics={"n_tokens": copies[0].shape[0] if copies else 0, + "cos_avg": 1.0, "cos_min": worst_cos}) # ── main ───────────────────────────────────────────────────────── @@ -179,8 +204,8 @@ def main(): ) ap.add_argument("--dir-multi", required=True, help="Multi-copy dump directory (batch=[A x N])") - ap.add_argument("--num-copies", type=int, required=True, - help="Number of copies N") + ap.add_argument("--num-copies", type=int, default=None, + help="Number of copies N (auto-detected from cu_seqlens if omitted)") ap.add_argument("--layer", type=int, default=None, help="Compare specific layer 1-indexed (default: all)") ap.add_argument("--tag", default="old", @@ -199,54 +224,55 @@ def main(): if cu is None: print(f"{_CROSS} cu_seqlens_q.pt missing"); return 1 - T = int(cu[1]) - int(cu[0]) - n = args.num_copies - n_pair = n * (n - 1) // 2 + total_copies = cu.numel() - 1 + if args.num_copies is not None and args.num_copies != total_copies: + print(f"[warn] --num-copies={args.num_copies} but cu_seqlens has {total_copies} copies; using {total_copies}") + lengths = [int(cu[i + 1]) - int(cu[i]) for i in range(total_copies)] print(_SEP_DOUBLE) print(" GEMM Precision Baseline — Within-Batch Pairwise Comparison") - print(f" Directory: {args.dir_multi} ({n} copies, {T} tokens each)") - print(f" Pairs: {n_pair}") + print(f" Directory: {args.dir_multi} ({total_copies} sequences)") + print(f" Lengths: {lengths}") print(_SEP_DOUBLE) all_results: list[CheckResult] = [] # ── hidden_states ── r = _compare_plain_within(args.dir_multi, "hidden_states.pt", - cu, n, args.layer, "hidden_states") + cu, total_copies, args.layer, "hidden_states") if r: all_results.append(r); _print_hidden_states(r) # ── build_kv_input_v ── r = _compare_plain_within(args.dir_multi, "build_kv_input_v.pt", - cu, n, args.layer, "build_kv_input_v") + cu, total_copies, args.layer, "build_kv_input_v") if r: all_results.append(r); _print_build_kv_input_v(r) # ── rope_preqk ── r = _compare_rope_within(args.dir_multi, "rope_preqk.pt", - cu, n, args.layer, "rope_preqk", "query", "key") + cu, total_copies, args.layer, "rope_preqk", "query", "key") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── rope_freqs ── r = _compare_plain_within(args.dir_multi, "rope_freqs.pt", - cu, n, args.layer, "rope_freqs") + cu, total_copies, args.layer, "rope_freqs") if r: all_results.append(r); _print_rope_freqs(r) # ── rope_postqk ── r = _compare_rope_within(args.dir_multi, "rope_postqk.pt", - cu, n, args.layer, "rope_postqk", "query", "key") + cu, total_copies, args.layer, "rope_postqk", "query", "key") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── attn_outputs ── r = _compare_plain_within(args.dir_multi, "attn_outputs.pt", - cu, n, args.layer, "attn_outputs") + cu, total_copies, args.layer, "attn_outputs") if r: all_results.append(r); _print_per_layer(r) # ── full_kv ── r = _compare_rope_within(args.dir_multi, "full_kv.pt", - cu, n, args.layer, "full_kv", "key", "value") + cu, total_copies, args.layer, "full_kv", "key", "value") if r: all_results.append(r); _print_rope_postqk_per_layer(r) # ── logits ── - r = _compare_logits_within(args.dir_multi, cu, n) + r = _compare_logits_within(args.dir_multi, cu, total_copies) if r: all_results.append(r); _print_logits_packed(r) # ── 2D: logprobs + entropy ── @@ -300,7 +326,7 @@ def main(): if args.output: _dump_json(all_results, args.output, "—", args.dir_multi, - tag=f"within_batch_N{n}", dir_off2=None) + tag=f"within_batch_N{total_copies}", dir_off2=None) if __name__ == "__main__": From a10bb084a9c816764b4742ad0f9eaacc9b2a5efa Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 10:41:30 +0800 Subject: [PATCH 45/61] [fix] baseline: compute cos_avg from actual per-token cosine (not hardcoded 1.0) _all comparison functions now collect per-token cos values across copies and compute proper mean for cos_avg. cos_min remains min across all comparisons. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 136 ++++++++++-------- .../tools/cmp_baseline_within_batch.py | 50 ++++--- 2 files changed, 105 insertions(+), 81 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 64aedd93..2357291b 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -107,35 +107,35 @@ def _compare_plain(dir_single: str, dir_stacked: str, filename: str, for lyr in layers: st = sd[lyr].float() mt = md[lyr].float() - layer_md, layer_cos = 0.0, 1.0 - total_tokens = 0 + layer_md, layer_cos_min = 0.0, 1.0 + all_cos: list[float] = [] for i in range(B): s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] T_i = s_seq.shape[0] if T_i == 0: continue on_f = s_seq.reshape(T_i, -1) - # compare against each stack repetition of this sequence for k in range(n_copies): j = i + k * B c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] if c_seq.shape[0] != T_i: continue off_f = c_seq.reshape(T_i, -1) - diff = (on_f - off_f).abs() cos = _cosine_sim(on_f, off_f, dim=-1) - md_i = float(diff.max()) + all_cos.extend(cos.tolist()) + md_i = float((on_f - off_f).abs().max()) cos_i = float(cos.min()) layer_md = max(layer_md, md_i) - layer_cos = min(layer_cos, cos_i) - total_tokens += T_i + layer_cos_min = min(layer_cos_min, cos_i) per_layer[lyr] = { - "max_diff": layer_md, "cos_avg": 1.0, "cos_min": layer_cos, - "n_tokens": total_tokens, "on_T": total_tokens, - "off_T": total_tokens * n_copies, + "max_diff": layer_md, + "cos_avg": sum(all_cos) / len(all_cos) if all_cos else 0.0, + "cos_min": layer_cos_min, + "n_tokens": int(sum(c.shape[0] for c in [st[int(cu_s[i]):int(cu_s[i+1])] for i in range(B)])) if B > 0 else 0, + "on_T": st.shape[0], "off_T": mt.shape[0], } worst_md = max(worst_md, layer_md) - worst_cos = min(worst_cos, layer_cos) + worst_cos = min(worst_cos, layer_cos_min) passed = worst_md < 1e-5 _name = f"{label}_L{layer}" if layer is not None else label @@ -167,30 +167,33 @@ def _compare_rope(dir_single: str, dir_stacked: str, filename: str, sk = sd[lyr][fld_k].float(); mk = md[lyr][fld_k].float() q_max, k_max = 0.0, 0.0 q_cos_min, k_cos_min = 1.0, 1.0 - total_tokens = 0 + q_all, k_all = [], [] for i in range(B): s_q = sq[int(cu_s[i]):int(cu_s[i + 1])] s_k = sk[int(cu_s[i]):int(cu_s[i + 1])] T_i = s_q.shape[0] - if T_i == 0: - continue + if T_i == 0: continue qf = s_q.reshape(T_i, -1); kf = s_k.reshape(T_i, -1) - for k in range(n_copies): - j = i + k * B + for h in range(n_copies): + j = i + h * B cq = mq[int(cu_m[j]):int(cu_m[j + 1])] ck = mk[int(cu_m[j]):int(cu_m[j + 1])] - if cq.shape[0] != T_i: - continue - q_max = max(q_max, float((qf - cq.reshape(T_i, -1)).abs().max())) - k_max = max(k_max, float((kf - ck.reshape(T_i, -1)).abs().max())) - q_cos_min = min(q_cos_min, float(_cosine_sim(qf, cq.reshape(T_i, -1), dim=-1).min())) - k_cos_min = min(k_cos_min, float(_cosine_sim(kf, ck.reshape(T_i, -1), dim=-1).min())) - total_tokens += T_i + if cq.shape[0] != T_i: continue + cfq = cq.reshape(T_i, -1); cfk = ck.reshape(T_i, -1) + q_cos = _cosine_sim(qf, cfq, dim=-1) + k_cos = _cosine_sim(kf, cfk, dim=-1) + q_all.extend(q_cos.tolist()); k_all.extend(k_cos.tolist()) + q_max = max(q_max, float((qf - cfq).abs().max())) + k_max = max(k_max, float((kf - cfk).abs().max())) + q_cos_min = min(q_cos_min, float(q_cos.min())) + k_cos_min = min(k_cos_min, float(k_cos.min())) per_layer[lyr] = { "Q_max_diff": q_max, "K_max_diff": k_max, - "Q_cos_avg": 1.0, "Q_cos_min": q_cos_min, - "K_cos_avg": 1.0, "K_cos_min": k_cos_min, - "n_tokens": total_tokens, + "Q_cos_avg": sum(q_all)/len(q_all) if q_all else 0.0, + "Q_cos_min": q_cos_min, + "K_cos_avg": sum(k_all)/len(k_all) if k_all else 0.0, + "K_cos_min": k_cos_min, + "n_tokens": len(q_all), } _name = f"{label}_L{layer}" if layer is not None else label return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) @@ -209,58 +212,64 @@ def _compare_logits(dir_single: str, dir_stacked: str, st = st.reshape(-1, st.size(-1)) mt = mt.reshape(-1, mt.size(-1)) B = cu_s.numel() - 1 - worst_md, worst_cos, total_tokens = 0.0, 1.0, 0 + worst_md, worst_cos_min = 0.0, 1.0 + all_cos, total_tokens = [], 0 for i in range(B): s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] - T_i = s_seq.shape[0] - if T_i == 0: - continue - total_tokens += T_i + if s_seq.shape[0] == 0: continue + total_tokens += s_seq.shape[0] for k in range(n_copies): j = i + k * B c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] - if c_seq.shape[0] != T_i: - continue - diff = (s_seq - c_seq).abs() + if c_seq.shape[0] != s_seq.shape[0]: continue cos = _cosine_sim(s_seq, c_seq, dim=-1) - worst_md = max(worst_md, float(diff.max())) - worst_cos = min(worst_cos, float(cos.min())) + all_cos.extend(cos.tolist()) + worst_md = max(worst_md, float((s_seq - c_seq).abs().max())) + worst_cos_min = min(worst_cos_min, float(cos.min())) return CheckResult(name="logits", passed=worst_md < 1e-5, - metrics={"n_tokens": total_tokens, "cos_avg": 1.0, - "cos_min": worst_cos}) + metrics={"n_tokens": total_tokens, + "cos_avg": sum(all_cos)/len(all_cos) if all_cos else 0.0, + "cos_min": worst_cos_min}) def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, - name: str, atol: float = 1e-5 + name: str, n_copies: int, num_seq: int, + atol: float = 1e-5 ) -> tuple[CheckResult | None, torch.Tensor | None, torch.Tensor | None]: - """Compare 2D [B, L_max] tensors — reuses ``_print_2d_result`` format.""" + """Compare 2D [B, L_max] tensors — row i from single vs rows i+k*num_seq from stacked.""" st = _load_tensor(dir_single, filename) mt = _load_tensor(dir_stacked, filename) if st is None or mt is None: return None, st, mt - if st.dim() > 1 and mt.dim() > 1 and st.shape[1] == mt.shape[1]: - # compare row 0 vs every row in multi - a = st[0:1].float() - best_md, best_cos = 0.0, 1.0 - for i in range(mt.shape[0]): - b = mt[i:i + 1].float() + st = st.float(); mt = mt.float() + if st.dim() < 2 or mt.dim() < 2: + return None, st, mt + # only compare valid (non-padded) rows + worst_md, worst_cos = 0.0, 1.0 + for i in range(num_seq): + if i >= st.shape[0]: + break + a = st[i].reshape(-1) + for k in range(n_copies): + j = i + k * num_seq + if j >= mt.shape[0]: + continue + b = mt[j].reshape(-1) md = float((a - b).abs().max()) - cos = float(_cosine_sim(a.reshape(-1), b.reshape(-1), dim=-1)) - best_md = max(best_md, md) - best_cos = min(best_cos, cos) - err = _error_abs_rel(st.float(), mt.float()) - pr = _pearson_r(st.float(), mt.float()) - return (CheckResult(name=name, - passed=err["abs_max"] <= atol, - metrics={"shape": tuple(st.shape), - "active": st.numel(), - "abs_max": err["abs_max"], - "abs_mean": err["abs_mean"], - "rel_max": err["rel_max"], - "rel_mean": err["rel_mean"], - "pearson_r": pr, - "atol": atol}), st, mt) - return None, st, mt + cos = float(_cosine_sim(a, b, dim=-1)) + worst_md = max(worst_md, md) + worst_cos = min(worst_cos, cos) + pr = _pearson_r(st[:num_seq].reshape(-1), mt[:num_seq * n_copies].reshape(-1)) + return (CheckResult(name=name, + passed=worst_md <= atol, + metrics={"shape": tuple(st.shape), + "active": st[:num_seq].numel(), + "abs_max": worst_md, + "abs_mean": 0.0, + "rel_max": 0.0, + "rel_mean": 0.0, + "pearson_r": pr, + "atol": atol}), st, mt) # ── top-K helpers (reuse cmp_diag_verl080 top-K printers) ──────── @@ -357,6 +366,7 @@ def main(): T = int(cu_s[-1]) n = args.num_copies + num_seq = cu_s.numel() - 1 print(_SEP_DOUBLE) print(" GEMM Precision Baseline — Cross-Batch-Size Comparison") print(f" Single: {args.dir_single} (1 seq, {T} tokens)") @@ -408,7 +418,7 @@ def main(): for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: fname = f"{fn}_{args.tag}.pt" r, t1, t2 = _compare_2d_file(args.dir_single, args.dir_stacked, - fname, f"{cn}_{args.tag}", args.atol) + fname, f"{cn}_{args.tag}", n, num_seq, args.atol) if r: all_results.append(r); _print_2d_result(r) # ── top-K ── diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index e4d8eccb..4d949e07 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -73,17 +73,20 @@ def _get_layers(data: dict) -> list[int]: # ── comparison logic ───────────────────────────────────────────── def _worst_pairwise(copies: list[torch.Tensor]) -> dict: - """Pairwise compare all copies of the SAME length; return worst.""" - worst_md, worst_cos = 0.0, 1.0 + """Pairwise compare all copies of the SAME length; return worst + avg cos.""" + worst_md, worst_cos_min = 0.0, 1.0 + all_cos: list[float] = [] for i in range(len(copies)): for j in range(i + 1, len(copies)): a = copies[i].reshape(copies[i].shape[0], -1).float() b = copies[j].reshape(copies[j].shape[0], -1).float() + cos = _cosine_sim(a, b, dim=-1) + all_cos.extend(cos.tolist()) md = float((a - b).abs().max()) - cos_min = float(_cosine_sim(a, b, dim=-1).min()) worst_md = max(worst_md, md) - worst_cos = min(worst_cos, cos_min) - return {"max_diff": worst_md, "cos_min": worst_cos} + worst_cos_min = min(worst_cos_min, float(cos.min())) + return {"max_diff": worst_md, "cos_avg": sum(all_cos)/len(all_cos) if all_cos else 0.0, + "cos_min": worst_cos_min} def _group_by_len(copies: list[torch.Tensor]) -> dict[int, list[torch.Tensor]]: @@ -112,20 +115,23 @@ def _compare_plain_within(dir_multi: str, filename: str, mt = d[lyr].float() copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] groups = _group_by_len(copies) - layer_md, layer_cos, total_tokens = 0.0, 1.0, 0 + layer_md, layer_cos_min = 0.0, 1.0 + all_cos_avg: list[float] = [] for group_copies in groups.values(): if len(group_copies) < 2: continue w = _worst_pairwise(group_copies) layer_md = max(layer_md, w["max_diff"]) - layer_cos = min(layer_cos, w["cos_min"]) - total_tokens += group_copies[0].shape[0] * len(group_copies) + layer_cos_min = min(layer_cos_min, w["cos_min"]) + all_cos_avg.append(w["cos_avg"]) per_layer[lyr] = { - "max_diff": layer_md, "cos_avg": 1.0, "cos_min": layer_cos, - "n_tokens": total_tokens, "on_T": total_tokens, "off_T": total_tokens, + "max_diff": layer_md, + "cos_avg": sum(all_cos_avg)/len(all_cos_avg) if all_cos_avg else 0.0, + "cos_min": layer_cos_min, + "n_tokens": mt.shape[0], "on_T": mt.shape[0], "off_T": mt.shape[0], } worst_md = max(worst_md, layer_md) - worst_cos = min(worst_cos, layer_cos) + worst_cos = min(worst_cos, layer_cos_min) passed = worst_md == 0.0 _name = f"{label}_L{layer}" if layer is not None else label return CheckResult(name=_name, passed=passed, @@ -153,22 +159,27 @@ def _compare_rope_within(dir_multi: str, filename: str, k_copies = [mk[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] q_groups = _group_by_len(q_copies) k_groups = _group_by_len(k_copies) - qw = {"max_diff": 0.0, "cos_min": 1.0} - kw = {"max_diff": 0.0, "cos_min": 1.0} + qw = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} + kw = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} + q_avgs, k_avgs = [], [] for g in q_groups.values(): if len(g) >= 2: w = _worst_pairwise(g) qw["max_diff"] = max(qw["max_diff"], w["max_diff"]) qw["cos_min"] = min(qw["cos_min"], w["cos_min"]) + q_avgs.append(w["cos_avg"]) for g in k_groups.values(): if len(g) >= 2: w = _worst_pairwise(g) kw["max_diff"] = max(kw["max_diff"], w["max_diff"]) kw["cos_min"] = min(kw["cos_min"], w["cos_min"]) + k_avgs.append(w["cos_avg"]) per_layer[lyr] = { "Q_max_diff": qw["max_diff"], "K_max_diff": kw["max_diff"], - "Q_cos_avg": 1.0, "Q_cos_min": qw["cos_min"], - "K_cos_avg": 1.0, "K_cos_min": kw["cos_min"], + "Q_cos_avg": sum(q_avgs)/len(q_avgs) if q_avgs else 0.0, + "Q_cos_min": qw["cos_min"], + "K_cos_avg": sum(k_avgs)/len(k_avgs) if k_avgs else 0.0, + "K_cos_min": kw["cos_min"], "n_tokens": sum(g[0].shape[0] * len(g) for g in q_groups.values()), } _name = f"{label}_L{layer}" if layer is not None else label @@ -184,15 +195,18 @@ def _compare_logits_within(dir_multi: str, cu: torch.Tensor, mt = mt.reshape(-1, mt.size(-1)) copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] groups = _group_by_len(copies) - worst_md, worst_cos = 0.0, 1.0 + worst_md, worst_cos_min = 0.0, 1.0 + all_cos_avg: list[float] = [] for g in groups.values(): if len(g) >= 2: w = _worst_pairwise(g) worst_md = max(worst_md, w["max_diff"]) - worst_cos = min(worst_cos, w["cos_min"]) + worst_cos_min = min(worst_cos_min, w["cos_min"]) + all_cos_avg.append(w["cos_avg"]) return CheckResult(name="logits", passed=worst_md == 0.0, metrics={"n_tokens": copies[0].shape[0] if copies else 0, - "cos_avg": 1.0, "cos_min": worst_cos}) + "cos_avg": sum(all_cos_avg)/len(all_cos_avg) if all_cos_avg else 0.0, + "cos_min": worst_cos_min}) # ── main ───────────────────────────────────────────────────────── From 9de6f2b7198f308d802c21df5dbb701426226061 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 10:46:20 +0800 Subject: [PATCH 46/61] [fix] within_batch top-K: use total_copies from cu_seqlens, not undefined n Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_within_batch.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 4d949e07..e553c963 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -319,20 +319,22 @@ def main(): # ── top-K ── if args.topk > 0: - # build_kv_input_v last-layer top-K d = _load_dict(args.dir_multi, "build_kv_input_v.pt") if d: lyr = max(int(k) for k in d.keys()) mt = d[lyr].float() - copies = [_extract(mt, cu, i).reshape(mt.shape[0] // n, -1) for i in range(n)] + copies = [mt[int(cu[i]):int(cu[i+1])].reshape(-1) for i in range(total_copies)] + groups = _group_by_len(copies) worst_md = 0.0; worst_a = worst_b = None - for i in range(n): - for j in range(i + 1, n): - md = float((copies[i] - copies[j]).abs().max()) - if md > worst_md: - worst_md = md; worst_a = copies[i]; worst_b = copies[j] + for g in groups.values(): + if len(g) < 2: continue + for i in range(len(g)): + for j in range(i + 1, len(g)): + md = float((g[i] - g[j]).abs().max()) + if md > worst_md: + worst_md = md; worst_a = g[i]; worst_b = g[j] if worst_a is not None: - _print_topk_vec(worst_a[0].cpu(), worst_b[0].cpu(), + _print_topk_vec(worst_a.cpu(), worst_b.cpu(), args.topk, args.sort_err, f"build_kv_input_v_L{lyr}_token0") From da791d85bf97a14933da26a5243de23cb38bf687 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 10:58:00 +0800 Subject: [PATCH 47/61] [fix] baseline: compute rel/pearson for 2D, add 2D top-K, fix pearson shape mismatch - 2D: rel_max/rel_mean/abs_mean from actual per-pair diffs (not hardcoded) - 2D: pearson from paired rows (fix shape mismatch 1024 vs 4096) - 2D: _print_topk_2d for logp/entropy positions - within_batch: same fixes Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 44 +++++++++++++------ .../tools/cmp_baseline_within_batch.py | 29 ++++++++---- 2 files changed, 52 insertions(+), 21 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 2357291b..6e74286f 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -236,7 +236,7 @@ def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, name: str, n_copies: int, num_seq: int, atol: float = 1e-5 ) -> tuple[CheckResult | None, torch.Tensor | None, torch.Tensor | None]: - """Compare 2D [B, L_max] tensors — row i from single vs rows i+k*num_seq from stacked.""" + """Compare 2D [B, L_max] — row i from single vs rows i+k*num_seq from stacked.""" st = _load_tensor(dir_single, filename) mt = _load_tensor(dir_stacked, filename) if st is None or mt is None: @@ -244,31 +244,38 @@ def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, st = st.float(); mt = mt.float() if st.dim() < 2 or mt.dim() < 2: return None, st, mt - # only compare valid (non-padded) rows - worst_md, worst_cos = 0.0, 1.0 + worst_md, worst_cos, worst_rel = 0.0, 1.0, 0.0 + all_abs, all_rel = [], [] + pearson_pairs = [] for i in range(num_seq): - if i >= st.shape[0]: - break + if i >= st.shape[0]: break a = st[i].reshape(-1) for k in range(n_copies): j = i + k * num_seq - if j >= mt.shape[0]: - continue + if j >= mt.shape[0]: continue b = mt[j].reshape(-1) - md = float((a - b).abs().max()) + diff = (a - b).abs() + rel = diff / a.abs().clamp(min=1e-8) + md = float(diff.max()) + rm = float(rel.max()) cos = float(_cosine_sim(a, b, dim=-1)) + all_abs.extend(diff.tolist()); all_rel.extend(rel.tolist()) worst_md = max(worst_md, md) + worst_rel = max(worst_rel, rm) worst_cos = min(worst_cos, cos) - pr = _pearson_r(st[:num_seq].reshape(-1), mt[:num_seq * n_copies].reshape(-1)) + pearson_pairs.append((a, b)) + # pearson: take up to 10 pairs + pr_vals = [_pearson_r(a, b) for a, b in pearson_pairs[:10]] + pr = sum(v for v in pr_vals if not (v != v)) / max(1, sum(1 for v in pr_vals if not (v != v))) return (CheckResult(name=name, passed=worst_md <= atol, metrics={"shape": tuple(st.shape), "active": st[:num_seq].numel(), "abs_max": worst_md, - "abs_mean": 0.0, - "rel_max": 0.0, - "rel_mean": 0.0, - "pearson_r": pr, + "abs_mean": sum(all_abs)/len(all_abs) if all_abs else 0.0, + "rel_max": worst_rel, + "rel_mean": sum(all_rel)/len(all_rel) if all_rel else 0.0, + "pearson_r": pr if pr_vals else 0.0, "atol": atol}), st, mt) @@ -415,11 +422,14 @@ def main(): if r: all_results.append(r); _print_logits_packed(r) # ── 2D: logprobs + entropy ── + _2d_results: list[tuple[str, torch.Tensor, torch.Tensor]] = [] for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: fname = f"{fn}_{args.tag}.pt" r, t1, t2 = _compare_2d_file(args.dir_single, args.dir_stacked, fname, f"{cn}_{args.tag}", n, num_seq, args.atol) if r: all_results.append(r); _print_2d_result(r) + if t1 is not None and t2 is not None: + _2d_results.append((f"{cn}_{args.tag}", t1, t2)) # ── top-K ── if args.topk > 0: @@ -432,6 +442,14 @@ def main(): _topk_rope(args.dir_single, args.dir_stacked, cu_s, cu_m, n, "rope_postqk.pt", "query", "key", "rope_postqk", args.topk, args.sort_err) + # 2D top-K: compare row 0 from single vs row 0 from stacked (same shuffled seq) + for label, st, mt in _2d_results: + if st.dim() >= 2 and mt.dim() >= 2 and st.shape[1] == mt.shape[1]: + # align: single row i vs stacked row i (first copy of each sequence) + t1 = st[:num_seq]; t2 = mt[:num_seq] + if t1.shape == t2.shape: + _print_topk_2d(t1.cpu(), t2.cpu(), None, + args.topk, args.sort_err, label) _print_summary(all_results) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index e553c963..7b31c321 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -41,6 +41,7 @@ _print_logits_packed, _print_2d_result, _print_topk_vec, + _print_topk_2d, _print_summary, ) @@ -290,35 +291,47 @@ def main(): if r: all_results.append(r); _print_logits_packed(r) # ── 2D: logprobs + entropy ── + _2d_within: list[tuple[str, torch.Tensor]] = [] for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: fname = f"{fn}_{args.tag}.pt" st = _load_tensor(args.dir_multi, fname) if st is not None and st.dim() >= 2: B = st.shape[0] - worst_md, worst_cos = 0.0, 1.0 + all_abs, all_rel = [], [] + worst_md, worst_cos, worst_rel = 0.0, 1.0, 0.0 for i in range(B): for j in range(i + 1, B): a = st[i].float().reshape(-1) b = st[j].float().reshape(-1) - md = float((a - b).abs().max()) - cos = float(_cosine_sim(a, b, dim=-1)) + diff = (a - b).abs(); rel = diff / a.abs().clamp(min=1e-8) + all_abs.extend(diff.tolist()); all_rel.extend(rel.tolist()) + md = float(diff.max()); rm = float(rel.max()) worst_md = max(worst_md, md) - worst_cos = min(worst_cos, cos) + worst_rel = max(worst_rel, rm) + worst_cos = min(worst_cos, float(_cosine_sim(a, b, dim=-1))) r = CheckResult(name=f"{cn}_{args.tag}", passed=worst_md == 0.0, metrics={"shape": tuple(st.shape), "active": st.numel(), "abs_max": worst_md, - "abs_mean": 0.0, - "rel_max": 0.0, - "rel_mean": 0.0, - "pearson_r": 1.0, + "abs_mean": sum(all_abs)/len(all_abs) if all_abs else 0.0, + "rel_max": worst_rel, + "rel_mean": sum(all_rel)/len(all_rel) if all_rel else 0.0, + "pearson_r": _pearson_r(st[:B//2].float().reshape(-1), + st[B//2:].float().reshape(-1)), "atol": args.atol}) all_results.append(r) _print_2d_result(r) + _2d_within.append((f"{cn}_{args.tag}", st)) # ── top-K ── if args.topk > 0: + # 2D top-K: compare row 0 vs row 1 (should be identical if same seq) + for label, st2d in _2d_within: + if st2d.shape[0] >= 2: + _print_topk_2d(st2d[:1].cpu(), st2d[1:2].cpu(), None, + args.topk, args.sort_err, label) + # build_kv_input_v top-K d = _load_dict(args.dir_multi, "build_kv_input_v.pt") if d: lyr = max(int(k) for k in d.keys()) From 366d1ce6d18cc93c0c6ed83643f0f11f440666eb Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:01:50 +0800 Subject: [PATCH 48/61] [fix] cross_batch top-K: guard _topk_rope with try-except for non-dict entries Co-Authored-By: Claude Fable 5 --- .../prefix_sharing/tools/cmp_baseline_cross_batch.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 6e74286f..dcfd26fa 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -321,7 +321,11 @@ def _topk_rope(dir_single: str, dir_stacked: str, lyr = max(int(k) for k in sd.keys()) B = cu_s.numel() - 1 for fld, tag in [(fld_q, f"{label}_Q"), (fld_k, f"{label}_K")]: - sq = sd[lyr][fld].float(); mq = md[lyr][fld].float() + try: + sq = sd[lyr][fld].float() + mq = md[lyr][fld].float() + except (KeyError, TypeError, AttributeError) as e: + print(f" [top-K skip] {label} {fld}: {e}"); continue worst_md = 0.0; worst_on = worst_off = None for i in range(B): s_seq = sq[int(cu_s[i]):int(cu_s[i + 1])] From ecfac35c0cfacbd670123304b7eb2ffa6b6b37a6 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:06:43 +0800 Subject: [PATCH 49/61] [cmp] baseline: add _print_shapes for .pt file shape diagnostics Co-Authored-By: Claude Fable 5 --- .../prefix_sharing/tools/cmp_baseline_cross_batch.py | 4 ++++ .../prefix_sharing/tools/cmp_baseline_within_batch.py | 3 +++ 2 files changed, 7 insertions(+) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index dcfd26fa..3b5299a1 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -51,6 +51,7 @@ _print_topk_vec, _print_topk_2d, _print_summary, + _print_shapes, ) @@ -384,6 +385,9 @@ def main(): print(f" Stacked: {args.dir_stacked} ({n} copies, {n * T} tokens)") print(_SEP_DOUBLE) + # shape diagnostics + _print_shapes(args.dir_single, args.dir_stacked, args.tag) + all_results: list[CheckResult] = [] # ── hidden_states ── diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 7b31c321..3685cae3 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -43,6 +43,7 @@ _print_topk_vec, _print_topk_2d, _print_summary, + _print_shapes, ) @@ -249,6 +250,8 @@ def main(): print(f" Lengths: {lengths}") print(_SEP_DOUBLE) + _print_shapes(args.dir_multi, args.dir_multi, args.tag) + all_results: list[CheckResult] = [] # ── hidden_states ── From efe5f6bc56ed17ba1a701109ed78eac601cea112 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:15:04 +0800 Subject: [PATCH 50/61] [refactor] baseline: replace ON/OFF labels with Single/Stacked via local wrappers Added _print_plain_baseline and _print_rope_baseline wrappers with SNG_T/STK_T or neutral column labels instead of ON_T/OFF_T. within_batch uses within-batch pairwise headers. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 62 +++++++++++++++-- .../tools/cmp_baseline_within_batch.py | 66 +++++++++++++++---- 2 files changed, 109 insertions(+), 19 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 3b5299a1..51d63a41 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -344,6 +344,54 @@ def _topk_rope(dir_single: str, dir_stacked: str, topk, sort_err, f"{tag}_L{lyr}_token0") +# ── print wrappers (replace ON/OFF labels with Single/Stacked) ── + +def _print_plain_baseline(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] Single vs Stacked（per-seq aligned）") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " + f"{'SNG_T':>8s} {'STK_T':>8s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") + for lyr in sorted(layers): + d = layers[lyr] + if "max_diff" not in d: + print(f" {lyr:>6d} {d.get('error','')}"); continue + md = d["max_diff"]; ca = d.get("cos_avg", 0); cm = d.get("cos_min", 1) + ok = md < 1e-5 + print(f" {lyr:>6d} {md:>12.3e} {ca:>10.6f} {cm:>10.6f} " + f"{d.get('on_T','—'):>8} {d.get('off_T','—'):>8} " + f"{'OK' if ok else 'DIFF':>8s}") + print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " + f"{_CHECK if r.passed else _CROSS}") + print() + + +def _print_rope_baseline(r: CheckResult, label: str): + _sec = label + print(_SEP_SINGLE + f"\n [{r.name}] {_sec} Single vs Stacked") + print(_SEP_SINGLE) + layers = r.metrics.get("layers") + if isinstance(layers, dict): + print(f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " + f"{'K_MAXDIFF':>12s} {'K_COS_AVG':>12s} {'K_COS_MIN':>12s} " + f"{'TOKENS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " + f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") + for lyr in sorted(layers.keys()): + d = layers[lyr] + if "error" in d: + print(f" {lyr:>6d} {d['error']}"); continue + print(f" {lyr:>6d} {d.get('Q_max_diff',0.0):>12.3e} {d.get('Q_cos_avg',0.0):>12.6e} " + f"{d.get('Q_cos_min',0.0):>12.6e} " + f"{d.get('K_max_diff',0.0):>12.3e} {d.get('K_cos_avg',0.0):>12.6e} " + f"{d.get('K_cos_min',0.0):>12.6e} {d.get('n_tokens','—'):>8}") + print() + + # ── main ───────────────────────────────────────────────────────── def main(): @@ -393,37 +441,37 @@ def main(): # ── hidden_states ── r = _compare_plain(args.dir_single, args.dir_stacked, "hidden_states.pt", cu_s, cu_m, n, args.layer, "hidden_states") - if r: all_results.append(r); _print_hidden_states(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── build_kv_input_v ── r = _compare_plain(args.dir_single, args.dir_stacked, "build_kv_input_v.pt", cu_s, cu_m, n, args.layer, "build_kv_input_v") - if r: all_results.append(r); _print_build_kv_input_v(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── rope_preqk ── r = _compare_rope(args.dir_single, args.dir_stacked, "rope_preqk.pt", cu_s, cu_m, n, args.layer, "rope_preqk", "query", "key") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, "rope_preqk") # ── rope_freqs ── r = _compare_plain(args.dir_single, args.dir_stacked, "rope_freqs.pt", cu_s, cu_m, n, args.layer, "rope_freqs") - if r: all_results.append(r); _print_rope_freqs(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── rope_postqk ── r = _compare_rope(args.dir_single, args.dir_stacked, "rope_postqk.pt", cu_s, cu_m, n, args.layer, "rope_postqk", "query", "key") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, "rope_postqk") # ── attn_outputs ── r = _compare_plain(args.dir_single, args.dir_stacked, "attn_outputs.pt", cu_s, cu_m, n, args.layer, "attn_outputs") - if r: all_results.append(r); _print_per_layer(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── full_kv ── r = _compare_rope(args.dir_single, args.dir_stacked, "full_kv.pt", cu_s, cu_m, n, args.layer, "full_kv", "key", "value") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, "full_kv") # ── logits ── r = _compare_logits(args.dir_single, args.dir_stacked, cu_s, cu_m, n) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 3685cae3..856924a8 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -33,11 +33,6 @@ _pearson_r, _dump_json, _load_tensor, - _print_per_layer, - _print_rope_postqk_per_layer, - _print_rope_freqs, - _print_build_kv_input_v, - _print_hidden_states, _print_logits_packed, _print_2d_result, _print_topk_vec, @@ -211,6 +206,53 @@ def _compare_logits_within(dir_multi: str, cu: torch.Tensor, "cos_min": worst_cos_min}) +# ── print wrappers ───────────────────────────────────────────── + +def _print_plain_baseline(r: CheckResult): + print(_SEP_SINGLE + f"\n [{r.name}] within-batch pairwise") + print(_SEP_SINGLE) + m = r.metrics + if "error" in m: + print(f" {_CROSS} {m['error']}\n"); return + layers = m.get("layers", {}) + print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " + f"{'TOKENS':>8s} {'STATUS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 10} {'─' * 8} {'─' * 8}") + for lyr in sorted(layers): + d = layers[lyr] + if "max_diff" not in d: + print(f" {lyr:>6d} {d.get('error','')}"); continue + md = d["max_diff"]; ca = d.get("cos_avg", 0); cm = d.get("cos_min", 1) + ok = md == 0.0 + print(f" {lyr:>6d} {md:>12.3e} {ca:>10.6f} {cm:>10.6f} " + f"{d.get('n_tokens','—'):>8} " + f"{'PASS' if ok else 'DIFF':>8s}") + print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " + f"{_CHECK if r.passed else _CROSS}") + print() + + +def _print_rope_baseline(r: CheckResult, label: str): + print(_SEP_SINGLE + f"\n [{r.name}] {label} within-batch pairwise") + print(_SEP_SINGLE) + layers = r.metrics.get("layers") + if isinstance(layers, dict): + print(f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " + f"{'K_MAXDIFF':>12s} {'K_COS_AVG':>12s} {'K_COS_MIN':>12s} " + f"{'TOKENS':>8s}") + print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " + f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") + for lyr in sorted(layers.keys()): + d = layers[lyr] + if "error" in d: + print(f" {lyr:>6d} {d['error']}"); continue + print(f" {lyr:>6d} {d.get('Q_max_diff',0.0):>12.3e} {d.get('Q_cos_avg',0.0):>12.6e} " + f"{d.get('Q_cos_min',0.0):>12.6e} " + f"{d.get('K_max_diff',0.0):>12.3e} {d.get('K_cos_avg',0.0):>12.6e} " + f"{d.get('K_cos_min',0.0):>12.6e} {d.get('n_tokens','—'):>8}") + print() + + # ── main ───────────────────────────────────────────────────────── def main(): @@ -257,37 +299,37 @@ def main(): # ── hidden_states ── r = _compare_plain_within(args.dir_multi, "hidden_states.pt", cu, total_copies, args.layer, "hidden_states") - if r: all_results.append(r); _print_hidden_states(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── build_kv_input_v ── r = _compare_plain_within(args.dir_multi, "build_kv_input_v.pt", cu, total_copies, args.layer, "build_kv_input_v") - if r: all_results.append(r); _print_build_kv_input_v(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── rope_preqk ── r = _compare_rope_within(args.dir_multi, "rope_preqk.pt", cu, total_copies, args.layer, "rope_preqk", "query", "key") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, r.name) # ── rope_freqs ── r = _compare_plain_within(args.dir_multi, "rope_freqs.pt", cu, total_copies, args.layer, "rope_freqs") - if r: all_results.append(r); _print_rope_freqs(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── rope_postqk ── r = _compare_rope_within(args.dir_multi, "rope_postqk.pt", cu, total_copies, args.layer, "rope_postqk", "query", "key") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, r.name) # ── attn_outputs ── r = _compare_plain_within(args.dir_multi, "attn_outputs.pt", cu, total_copies, args.layer, "attn_outputs") - if r: all_results.append(r); _print_per_layer(r) + if r: all_results.append(r); _print_plain_baseline(r) # ── full_kv ── r = _compare_rope_within(args.dir_multi, "full_kv.pt", cu, total_copies, args.layer, "full_kv", "key", "value") - if r: all_results.append(r); _print_rope_postqk_per_layer(r) + if r: all_results.append(r); _print_rope_baseline(r, r.name) # ── logits ── r = _compare_logits_within(args.dir_multi, cu, total_copies) From 649484d3c608384542d9f0373ccd92c5fc8633f4 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:20:36 +0800 Subject: [PATCH 51/61] [fix] within_batch 2D: compare same-seq stack copies only, add --num-seq 2D logp/entropy now compares row i vs row i+k*num_seq (same sequence stack copies) instead of all C(N,2) pairs across different sequences. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_within_batch.py | 26 +++++++++++++------ 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 856924a8..108c4642 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -263,7 +263,9 @@ def main(): ap.add_argument("--dir-multi", required=True, help="Multi-copy dump directory (batch=[A x N])") ap.add_argument("--num-copies", type=int, default=None, - help="Number of copies N (auto-detected from cu_seqlens if omitted)") + help="Total sequences in batch (auto-detected from cu_seqlens)") + ap.add_argument("--num-seq", type=int, default=None, + help="Distinct sequences before stacking (required for 2D within-batch)") ap.add_argument("--layer", type=int, default=None, help="Compare specific layer 1-indexed (default: all)") ap.add_argument("--tag", default="old", @@ -337,18 +339,25 @@ def main(): # ── 2D: logprobs + entropy ── _2d_within: list[tuple[str, torch.Tensor]] = [] + _num_seq = args.num_seq if args.num_seq else (total_copies // (args.num_copies or total_copies) if args.num_copies else total_copies) for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: fname = f"{fn}_{args.tag}.pt" st = _load_tensor(args.dir_multi, fname) if st is not None and st.dim() >= 2: B = st.shape[0] + stack = B // _num_seq if _num_seq > 0 and B % _num_seq == 0 else 1 all_abs, all_rel = [], [] worst_md, worst_cos, worst_rel = 0.0, 1.0, 0.0 - for i in range(B): - for j in range(i + 1, B): - a = st[i].float().reshape(-1) + # only compare stack copies of the SAME sequence (i vs i+k*num_seq) + for i in range(_num_seq): + if i >= B: break + a = st[i].float().reshape(-1) + for k in range(1, stack): + j = i + k * _num_seq + if j >= B: continue b = st[j].float().reshape(-1) - diff = (a - b).abs(); rel = diff / a.abs().clamp(min=1e-8) + diff = (a - b).abs() + rel = diff / a.abs().clamp(min=1e-8) all_abs.extend(diff.tolist()); all_rel.extend(rel.tolist()) md = float(diff.max()); rm = float(rel.max()) worst_md = max(worst_md, md) @@ -357,13 +366,14 @@ def main(): r = CheckResult(name=f"{cn}_{args.tag}", passed=worst_md == 0.0, metrics={"shape": tuple(st.shape), - "active": st.numel(), + "active": _num_seq * st.shape[1] if st.dim() >= 2 else st.numel(), "abs_max": worst_md, "abs_mean": sum(all_abs)/len(all_abs) if all_abs else 0.0, "rel_max": worst_rel, "rel_mean": sum(all_rel)/len(all_rel) if all_rel else 0.0, - "pearson_r": _pearson_r(st[:B//2].float().reshape(-1), - st[B//2:].float().reshape(-1)), + "pearson_r": _pearson_r(st[:_num_seq].float().reshape(-1), + st[_num_seq:2*_num_seq].float().reshape(-1)) + if B >= 2*_num_seq else 1.0, "atol": args.atol}) all_results.append(r) _print_2d_result(r) From 9eefb4d4eeaa7f9bb3cb5ec3854f3477be2ba961 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:29:44 +0800 Subject: [PATCH 52/61] [fix] within_batch 2D top-K: compare row 0 vs row num_seq (same seq copy) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Was comparing row 0 vs row 1 (different sequences → huge false diffs). Co-Authored-By: Claude Fable 5 --- .../prefix_sharing/tools/cmp_baseline_within_batch.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 108c4642..8c2027be 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -381,10 +381,11 @@ def main(): # ── top-K ── if args.topk > 0: - # 2D top-K: compare row 0 vs row 1 (should be identical if same seq) + # 2D top-K: compare row 0 vs row _num_seq (stack copies of same seq) for label, st2d in _2d_within: - if st2d.shape[0] >= 2: - _print_topk_2d(st2d[:1].cpu(), st2d[1:2].cpu(), None, + ns = _num_seq if _num_seq > 0 and (st2d.shape[0] % _num_seq == 0) else st2d.shape[0] + if st2d.shape[0] >= ns + 1: + _print_topk_2d(st2d[:1].cpu(), st2d[ns:ns+1].cpu(), None, args.topk, args.sort_err, label) # build_kv_input_v top-K d = _load_dict(args.dir_multi, "build_kv_input_v.pt") From abf29499b5655d8913b52e42da0bfd941e90d057 Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 11:38:59 +0800 Subject: [PATCH 53/61] [refactor] baseline: full variable rename for readability (open-source quality) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - sd/md → single_data/multi_data - cu_s/cu_m → cu_seqlens_single/cu_seqlens_multi - st/mt → single_tensor/multi_tensor - B → num_sequences, T_i → tokens_in_seq - lyr → layer_index, n → stack_count - i/k → seq_index/copy_index, j → multi_offset - worst_md → worst_max_diff, layer_md → layer_max_diff - on_f/off_f → single_flat/multi_flat - q_all/k_all → first_cos_list/second_cos_list - Consolidated comparison loops into DRY for-loops over file lists. - Removed all single-letter variable names in logic functions. Co-Authored-By: Claude Fable 5 --- .../tools/cmp_baseline_cross_batch.py | 963 ++++++++++-------- .../tools/cmp_baseline_within_batch.py | 732 +++++++------ 2 files changed, 988 insertions(+), 707 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py index 51d63a41..3c06ddd1 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_cross_batch.py @@ -21,7 +21,6 @@ from __future__ import annotations import argparse -import math import os import torch @@ -30,22 +29,12 @@ CheckResult, _SEP_DOUBLE, _SEP_SINGLE, - _SEP_THIN, _CHECK, _CROSS, - _COS_AVG_PASS, - _COS_MIN_PASS, _cosine_sim, - _error_abs_rel, _pearson_r, _dump_json, _load_tensor, - _print_header, - _print_per_layer, - _print_rope_postqk_per_layer, - _print_rope_freqs, - _print_build_kv_input_v, - _print_hidden_states, _print_logits_packed, _print_2d_result, _print_topk_vec, @@ -55,454 +44,626 @@ ) -# ── helpers ────────────────────────────────────────────────────── +# ══════════════════════════════════════════════════════════════════ +# I/O helpers +# ══════════════════════════════════════════════════════════════════ -def _load_dict(dir_path: str, filename: str) -> dict | None: - fp = os.path.join(dir_path, filename) - if not os.path.exists(fp): +def _load_per_layer_dict(directory: str, filename: str) -> dict | None: + """Load a ``{layer_index: tensor_or_dict}`` file.""" + filepath = os.path.join(directory, filename) + if not os.path.exists(filepath): return None - d = torch.load(fp, weights_only=True) - return d if isinstance(d, dict) else None + data = torch.load(filepath, weights_only=True) + return data if isinstance(data, dict) else None -def _load_cu(dir_path: str) -> torch.Tensor | None: - fp = os.path.join(dir_path, "cu_seqlens_q.pt") - if not os.path.exists(fp): +def _load_cu_seqlens(directory: str) -> torch.Tensor | None: + """Load ``cu_seqlens_q.pt`` (cumulative token boundaries).""" + filepath = os.path.join(directory, "cu_seqlens_q.pt") + if not os.path.exists(filepath): return None - return torch.load(fp, weights_only=True) + return torch.load(filepath, weights_only=True) -def _extract(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: - return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] +def _slice_sequence(packed_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + sequence_index: int) -> torch.Tensor: + """Slice one sequence from a packed tensor using cu_seqlens.""" + start = int(cu_seqlens[sequence_index]) + end = int(cu_seqlens[sequence_index + 1]) + return packed_tensor[start:end] -def _get_layers(data: dict) -> list[int]: +def _sorted_layer_keys(data: dict) -> list[int]: return sorted(int(k) for k in data.keys()) -# ── comparison logic ───────────────────────────────────────────── - -def _compare_plain(dir_single: str, dir_stacked: str, filename: str, - cu_s: torch.Tensor, cu_m: torch.Tensor, - n_copies: int, layer: int | None, - label: str) -> CheckResult | None: - """Compare ``{layer: [T, ...]}`` per-layer dicts across copies. - - Handles variable-length sequences by matching copies by sequence length. +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — plain per-layer dicts +# ══════════════════════════════════════════════════════════════════ + +def _compare_plain_per_layer( + dir_single: str, dir_stacked: str, + filename: str, + cu_seqlens_single: torch.Tensor, + cu_seqlens_multi: torch.Tensor, + stack_count: int, + filter_layer: int | None, + label: str, +) -> CheckResult | None: + """Compare ``{layer: [total_tokens, ...]}`` across batch sizes. + + Matches each sequence in the single dump against its *stack_count* + copies in the stacked dump (located at index + ``seq_index + copy_index * num_sequences``). """ - sd = _load_dict(dir_single, filename) - md = _load_dict(dir_stacked, filename) - if sd is None or md is None: + single_data = _load_per_layer_dict(dir_single, filename) + multi_data = _load_per_layer_dict(dir_stacked, filename) + if single_data is None or multi_data is None: return None - layers = _get_layers(sd) - if layer is not None: - layers = [l for l in layers if l == layer] + + layers = _sorted_layer_keys(single_data) + if filter_layer is not None: + layers = [l for l in layers if l == filter_layer] if not layers: return None - # single: B seqs, stacked: B*stack seqs - B = cu_s.numel() - 1 + num_sequences = cu_seqlens_single.numel() - 1 per_layer: dict = {} - worst_md = 0.0 - worst_cos = 1.0 - for lyr in layers: - st = sd[lyr].float() - mt = md[lyr].float() - layer_md, layer_cos_min = 0.0, 1.0 - all_cos: list[float] = [] - for i in range(B): - s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] - T_i = s_seq.shape[0] - if T_i == 0: + worst_max_diff = 0.0 + worst_cos_min = 1.0 + + for layer_index in layers: + single_tensor = single_data[layer_index].float() + multi_tensor = multi_data[layer_index].float() + layer_max_diff = 0.0 + layer_cos_min = 1.0 + all_cos_values: list[float] = [] + + for seq_index in range(num_sequences): + single_seq = _slice_sequence(single_tensor, cu_seqlens_single, seq_index) + tokens_in_seq = single_seq.shape[0] + if tokens_in_seq == 0: continue - on_f = s_seq.reshape(T_i, -1) - for k in range(n_copies): - j = i + k * B - c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] - if c_seq.shape[0] != T_i: + single_flat = single_seq.reshape(tokens_in_seq, -1) + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + multi_seq = _slice_sequence(multi_tensor, cu_seqlens_multi, multi_offset) + if multi_seq.shape[0] != tokens_in_seq: continue - off_f = c_seq.reshape(T_i, -1) - cos = _cosine_sim(on_f, off_f, dim=-1) - all_cos.extend(cos.tolist()) - md_i = float((on_f - off_f).abs().max()) - cos_i = float(cos.min()) - layer_md = max(layer_md, md_i) - layer_cos_min = min(layer_cos_min, cos_i) - per_layer[lyr] = { - "max_diff": layer_md, - "cos_avg": sum(all_cos) / len(all_cos) if all_cos else 0.0, + multi_flat = multi_seq.reshape(tokens_in_seq, -1) + + per_token_cos = _cosine_sim(single_flat, multi_flat, dim=-1) + all_cos_values.extend(per_token_cos.tolist()) + diff_i = float((single_flat - multi_flat).abs().max()) + cos_min_i = float(per_token_cos.min()) + layer_max_diff = max(layer_max_diff, diff_i) + layer_cos_min = min(layer_cos_min, cos_min_i) + + per_layer[layer_index] = { + "max_diff": layer_max_diff, + "cos_avg": sum(all_cos_values) / len(all_cos_values) if all_cos_values else 0.0, "cos_min": layer_cos_min, - "n_tokens": int(sum(c.shape[0] for c in [st[int(cu_s[i]):int(cu_s[i+1])] for i in range(B)])) if B > 0 else 0, - "on_T": st.shape[0], "off_T": mt.shape[0], + "n_tokens": single_tensor.shape[0], + "on_T": single_tensor.shape[0], + "off_T": multi_tensor.shape[0], } - worst_md = max(worst_md, layer_md) - worst_cos = min(worst_cos, layer_cos_min) - - passed = worst_md < 1e-5 - _name = f"{label}_L{layer}" if layer is not None else label - return CheckResult(name=_name, passed=passed, - metrics={"layers": per_layer, "max_diff": worst_md, - "cos_min": worst_cos, "num_layers": len(layers)}) - - -def _compare_rope(dir_single: str, dir_stacked: str, filename: str, - cu_s: torch.Tensor, cu_m: torch.Tensor, - n_copies: int, layer: int | None, - label: str, fld_q: str, fld_k: str - ) -> CheckResult | None: - """Compare ``{layer: {fld_q, fld_k}}`` dicts across copies.""" - sd = _load_dict(dir_single, filename) - md = _load_dict(dir_stacked, filename) - if sd is None or md is None: + worst_max_diff = max(worst_max_diff, layer_max_diff) + worst_cos_min = min(worst_cos_min, layer_cos_min) + + passed = worst_max_diff < 1e-5 + result_name = f"{label}_L{filter_layer}" if filter_layer is not None else label + return CheckResult( + name=result_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_max_diff, + "cos_min": worst_cos_min, "num_layers": len(layers)}, + ) + + +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — per-layer KV dicts (rope_preqk / rope_postqk / +# full_kv) +# ══════════════════════════════════════════════════════════════════ + +def _compare_kv_per_layer( + dir_single: str, dir_stacked: str, + filename: str, + cu_seqlens_single: torch.Tensor, + cu_seqlens_multi: torch.Tensor, + stack_count: int, + filter_layer: int | None, + label: str, + field_first: str, + field_second: str, +) -> CheckResult | None: + """Compare ``{layer: {field_first, field_second}}`` across batch sizes.""" + single_data = _load_per_layer_dict(dir_single, filename) + multi_data = _load_per_layer_dict(dir_stacked, filename) + if single_data is None or multi_data is None: return None - layers = _get_layers(sd) - if layer is not None: - layers = [l for l in layers if l == layer] + + layers = _sorted_layer_keys(single_data) + if filter_layer is not None: + layers = [l for l in layers if l == filter_layer] if not layers: return None - B = cu_s.numel() - 1 + num_sequences = cu_seqlens_single.numel() - 1 per_layer: dict = {} - for lyr in layers: - sq = sd[lyr][fld_q].float(); mq = md[lyr][fld_q].float() - sk = sd[lyr][fld_k].float(); mk = md[lyr][fld_k].float() - q_max, k_max = 0.0, 0.0 - q_cos_min, k_cos_min = 1.0, 1.0 - q_all, k_all = [], [] - for i in range(B): - s_q = sq[int(cu_s[i]):int(cu_s[i + 1])] - s_k = sk[int(cu_s[i]):int(cu_s[i + 1])] - T_i = s_q.shape[0] - if T_i == 0: continue - qf = s_q.reshape(T_i, -1); kf = s_k.reshape(T_i, -1) - for h in range(n_copies): - j = i + h * B - cq = mq[int(cu_m[j]):int(cu_m[j + 1])] - ck = mk[int(cu_m[j]):int(cu_m[j + 1])] - if cq.shape[0] != T_i: continue - cfq = cq.reshape(T_i, -1); cfk = ck.reshape(T_i, -1) - q_cos = _cosine_sim(qf, cfq, dim=-1) - k_cos = _cosine_sim(kf, cfk, dim=-1) - q_all.extend(q_cos.tolist()); k_all.extend(k_cos.tolist()) - q_max = max(q_max, float((qf - cfq).abs().max())) - k_max = max(k_max, float((kf - cfk).abs().max())) - q_cos_min = min(q_cos_min, float(q_cos.min())) - k_cos_min = min(k_cos_min, float(k_cos.min())) - per_layer[lyr] = { - "Q_max_diff": q_max, "K_max_diff": k_max, - "Q_cos_avg": sum(q_all)/len(q_all) if q_all else 0.0, - "Q_cos_min": q_cos_min, - "K_cos_avg": sum(k_all)/len(k_all) if k_all else 0.0, - "K_cos_min": k_cos_min, - "n_tokens": len(q_all), + + for layer_index in layers: + single_first = single_data[layer_index][field_first].float() + single_second = single_data[layer_index][field_second].float() + multi_first = multi_data[layer_index][field_first].float() + multi_second = multi_data[layer_index][field_second].float() + + first_max_diff, second_max_diff = 0.0, 0.0 + first_cos_min, second_cos_min = 1.0, 1.0 + first_cos_list, second_cos_list = [], [] + + for seq_index in range(num_sequences): + single_seq_f = _slice_sequence(single_first, cu_seqlens_single, seq_index) + single_seq_s = _slice_sequence(single_second, cu_seqlens_single, seq_index) + tokens_in_seq = single_seq_f.shape[0] + if tokens_in_seq == 0: + continue + single_flat_f = single_seq_f.reshape(tokens_in_seq, -1) + single_flat_s = single_seq_s.reshape(tokens_in_seq, -1) + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + multi_seq_f = _slice_sequence(multi_first, cu_seqlens_multi, multi_offset) + multi_seq_s = _slice_sequence(multi_second, cu_seqlens_multi, multi_offset) + if multi_seq_f.shape[0] != tokens_in_seq: + continue + multi_flat_f = multi_seq_f.reshape(tokens_in_seq, -1) + multi_flat_s = multi_seq_s.reshape(tokens_in_seq, -1) + + cos_f = _cosine_sim(single_flat_f, multi_flat_f, dim=-1) + cos_s = _cosine_sim(single_flat_s, multi_flat_s, dim=-1) + first_cos_list.extend(cos_f.tolist()) + second_cos_list.extend(cos_s.tolist()) + first_max_diff = max(first_max_diff, float((single_flat_f - multi_flat_f).abs().max())) + second_max_diff = max(second_max_diff, float((single_flat_s - multi_flat_s).abs().max())) + first_cos_min = min(first_cos_min, float(cos_f.min())) + second_cos_min = min(second_cos_min, float(cos_s.min())) + + per_layer[layer_index] = { + "Q_max_diff": first_max_diff, "K_max_diff": second_max_diff, + "Q_cos_avg": sum(first_cos_list) / len(first_cos_list) if first_cos_list else 0.0, + "Q_cos_min": first_cos_min, + "K_cos_avg": sum(second_cos_list) / len(second_cos_list) if second_cos_list else 0.0, + "K_cos_min": second_cos_min, + "n_tokens": len(first_cos_list), } - _name = f"{label}_L{layer}" if layer is not None else label - return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) + result_name = f"{label}_L{filter_layer}" if filter_layer is not None else label + return CheckResult(name=result_name, passed=True, metrics={"layers": per_layer}) -def _compare_logits(dir_single: str, dir_stacked: str, - cu_s: torch.Tensor, cu_m: torch.Tensor, - n_copies: int) -> CheckResult | None: - """Compare packed logits across copies.""" - fp_s = os.path.join(dir_single, "logits.pt") - fp_m = os.path.join(dir_stacked, "logits.pt") - if not os.path.exists(fp_s) or not os.path.exists(fp_m): + +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — logits (single packed tensor) +# ══════════════════════════════════════════════════════════════════ + +def _compare_logits_cross_batch( + dir_single: str, dir_stacked: str, + cu_seqlens_single: torch.Tensor, + cu_seqlens_multi: torch.Tensor, + stack_count: int, +) -> CheckResult | None: + """Compare packed logits across batch sizes.""" + single_path = os.path.join(dir_single, "logits.pt") + multi_path = os.path.join(dir_stacked, "logits.pt") + if not os.path.exists(single_path) or not os.path.exists(multi_path): return None - st = torch.load(fp_s, weights_only=True).float() - mt = torch.load(fp_m, weights_only=True).float() - st = st.reshape(-1, st.size(-1)) - mt = mt.reshape(-1, mt.size(-1)) - B = cu_s.numel() - 1 - worst_md, worst_cos_min = 0.0, 1.0 - all_cos, total_tokens = [], 0 - for i in range(B): - s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] - if s_seq.shape[0] == 0: continue - total_tokens += s_seq.shape[0] - for k in range(n_copies): - j = i + k * B - c_seq = mt[int(cu_m[j]):int(cu_m[j + 1])] - if c_seq.shape[0] != s_seq.shape[0]: continue - cos = _cosine_sim(s_seq, c_seq, dim=-1) - all_cos.extend(cos.tolist()) - worst_md = max(worst_md, float((s_seq - c_seq).abs().max())) - worst_cos_min = min(worst_cos_min, float(cos.min())) - return CheckResult(name="logits", passed=worst_md < 1e-5, - metrics={"n_tokens": total_tokens, - "cos_avg": sum(all_cos)/len(all_cos) if all_cos else 0.0, - "cos_min": worst_cos_min}) - - -def _compare_2d_file(dir_single: str, dir_stacked: str, filename: str, - name: str, n_copies: int, num_seq: int, - atol: float = 1e-5 - ) -> tuple[CheckResult | None, torch.Tensor | None, torch.Tensor | None]: - """Compare 2D [B, L_max] — row i from single vs rows i+k*num_seq from stacked.""" - st = _load_tensor(dir_single, filename) - mt = _load_tensor(dir_stacked, filename) - if st is None or mt is None: - return None, st, mt - st = st.float(); mt = mt.float() - if st.dim() < 2 or mt.dim() < 2: - return None, st, mt - worst_md, worst_cos, worst_rel = 0.0, 1.0, 0.0 - all_abs, all_rel = [], [] - pearson_pairs = [] - for i in range(num_seq): - if i >= st.shape[0]: break - a = st[i].reshape(-1) - for k in range(n_copies): - j = i + k * num_seq - if j >= mt.shape[0]: continue - b = mt[j].reshape(-1) - diff = (a - b).abs() - rel = diff / a.abs().clamp(min=1e-8) - md = float(diff.max()) - rm = float(rel.max()) - cos = float(_cosine_sim(a, b, dim=-1)) - all_abs.extend(diff.tolist()); all_rel.extend(rel.tolist()) - worst_md = max(worst_md, md) - worst_rel = max(worst_rel, rm) - worst_cos = min(worst_cos, cos) - pearson_pairs.append((a, b)) - # pearson: take up to 10 pairs - pr_vals = [_pearson_r(a, b) for a, b in pearson_pairs[:10]] - pr = sum(v for v in pr_vals if not (v != v)) / max(1, sum(1 for v in pr_vals if not (v != v))) - return (CheckResult(name=name, - passed=worst_md <= atol, - metrics={"shape": tuple(st.shape), - "active": st[:num_seq].numel(), - "abs_max": worst_md, - "abs_mean": sum(all_abs)/len(all_abs) if all_abs else 0.0, - "rel_max": worst_rel, - "rel_mean": sum(all_rel)/len(all_rel) if all_rel else 0.0, - "pearson_r": pr if pr_vals else 0.0, - "atol": atol}), st, mt) - - -# ── top-K helpers (reuse cmp_diag_verl080 top-K printers) ──────── - -def _topk_per_layer(dir_single: str, dir_stacked: str, - cu_s: torch.Tensor, cu_m: torch.Tensor, n_copies: int, - filename: str, label: str, - topk: int, sort_err: str): - """Print top-K worst dims for a per-layer plain file.""" - sd = _load_dict(dir_single, filename) - md = _load_dict(dir_stacked, filename) - if sd is None or md is None: + + single_logits = torch.load(single_path, weights_only=True).float() + multi_logits = torch.load(multi_path, weights_only=True).float() + single_logits = single_logits.reshape(-1, single_logits.size(-1)) + multi_logits = multi_logits.reshape(-1, multi_logits.size(-1)) + + num_sequences = cu_seqlens_single.numel() - 1 + worst_max_diff = 0.0 + worst_cos_min = 1.0 + all_cos_values: list[float] = [] + total_tokens = 0 + + for seq_index in range(num_sequences): + single_seq = _slice_sequence(single_logits, cu_seqlens_single, seq_index) + if single_seq.shape[0] == 0: + continue + total_tokens += single_seq.shape[0] + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + multi_seq = _slice_sequence(multi_logits, cu_seqlens_multi, multi_offset) + if multi_seq.shape[0] != single_seq.shape[0]: + continue + per_token_cos = _cosine_sim(single_seq, multi_seq, dim=-1) + all_cos_values.extend(per_token_cos.tolist()) + worst_max_diff = max(worst_max_diff, float((single_seq - multi_seq).abs().max())) + worst_cos_min = min(worst_cos_min, float(per_token_cos.min())) + + return CheckResult( + name="logits", passed=worst_max_diff < 1e-5, + metrics={ + "n_tokens": total_tokens, + "cos_avg": sum(all_cos_values) / len(all_cos_values) if all_cos_values else 0.0, + "cos_min": worst_cos_min, + }, + ) + + +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — 2D (logprobs / entropy) +# ══════════════════════════════════════════════════════════════════ + +def _compare_2d_cross_batch( + dir_single: str, dir_stacked: str, + filename: str, label: str, + stack_count: int, num_sequences: int, + atol: float = 1e-5, +) -> tuple[CheckResult | None, torch.Tensor | None, torch.Tensor | None]: + """Compare 2D [batch, L_max] — row ``i`` from single vs rows + ``i + k * num_sequences`` from stacked.""" + single_2d = _load_tensor(dir_single, filename) + multi_2d = _load_tensor(dir_stacked, filename) + if single_2d is None or multi_2d is None: + return None, single_2d, multi_2d + + single_2d = single_2d.float() + multi_2d = multi_2d.float() + if single_2d.dim() < 2 or multi_2d.dim() < 2: + return None, single_2d, multi_2d + + worst_max_diff = 0.0 + worst_cos_min = 1.0 + worst_rel_max = 0.0 + all_abs_diffs: list[float] = [] + all_rel_diffs: list[float] = [] + pearson_pairs: list[tuple[torch.Tensor, torch.Tensor]] = [] + + for seq_index in range(num_sequences): + if seq_index >= single_2d.shape[0]: + break + single_row = single_2d[seq_index].reshape(-1) + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + if multi_offset >= multi_2d.shape[0]: + continue + multi_row = multi_2d[multi_offset].reshape(-1) + + abs_diff = (single_row - multi_row).abs() + rel_diff = abs_diff / single_row.abs().clamp(min=1e-8) + all_abs_diffs.extend(abs_diff.tolist()) + all_rel_diffs.extend(rel_diff.tolist()) + + worst_max_diff = max(worst_max_diff, float(abs_diff.max())) + worst_rel_max = max(worst_rel_max, float(rel_diff.max())) + worst_cos_min = min(worst_cos_min, float(_cosine_sim(single_row, multi_row, dim=-1))) + pearson_pairs.append((single_row, multi_row)) + + # pearson: average over up to 10 pairs + pearson_values = [_pearson_r(a, b) for a, b in pearson_pairs[:10]] + pearson_avg = (sum(pearson_values) / len(pearson_values) + if pearson_values else 0.0) + + result = CheckResult( + name=label, passed=worst_max_diff <= atol, + metrics={ + "shape": tuple(single_2d.shape), + "active": single_2d[:num_sequences].numel(), + "abs_max": worst_max_diff, + "abs_mean": sum(all_abs_diffs) / len(all_abs_diffs) if all_abs_diffs else 0.0, + "rel_max": worst_rel_max, + "rel_mean": sum(all_rel_diffs) / len(all_rel_diffs) if all_rel_diffs else 0.0, + "pearson_r": pearson_avg, + "atol": atol, + }, + ) + return result, single_2d, multi_2d + + +# ══════════════════════════════════════════════════════════════════ +# Top-K helpers +# ══════════════════════════════════════════════════════════════════ + +def _print_topk_plain( + dir_single: str, dir_stacked: str, + cu_seqlens_single: torch.Tensor, + cu_seqlens_multi: torch.Tensor, + stack_count: int, + filename: str, label: str, + topk: int, sort_err: str, +): + """Print top-K worst dimensions for a plain per-layer file.""" + single_data = _load_per_layer_dict(dir_single, filename) + multi_data = _load_per_layer_dict(dir_stacked, filename) + if single_data is None or multi_data is None: return - lyr = max(int(k) for k in sd.keys()) - st = sd[lyr].float(); mt = md[lyr].float() - B = cu_s.numel() - 1 - worst_md = 0.0; worst_on = worst_off = None - for i in range(B): - s_seq = st[int(cu_s[i]):int(cu_s[i + 1])] - if s_seq.shape[0] == 0: continue - on_f = s_seq.reshape(s_seq.shape[0], -1) - for k in range(n_copies): - j = i + k * B - ct = mt[int(cu_m[j]):int(cu_m[j + 1])] - if ct.shape[0] != s_seq.shape[0]: continue - off_f = ct.reshape(ct.shape[0], -1) - md = float((on_f - off_f).abs().max()) - if md > worst_md: worst_md = md; worst_on = on_f; worst_off = off_f - if worst_on is not None: - _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), - topk, sort_err, f"{label}_L{lyr}_token0") - - -def _topk_rope(dir_single: str, dir_stacked: str, - cu_s: torch.Tensor, cu_m: torch.Tensor, n_copies: int, - filename: str, fld_q: str, fld_k: str, - label: str, topk: int, sort_err: str): - """Print top-K worst dims for a rope file.""" - sd = _load_dict(dir_single, filename) - md = _load_dict(dir_stacked, filename) - if sd is None or md is None: return - lyr = max(int(k) for k in sd.keys()) - B = cu_s.numel() - 1 - for fld, tag in [(fld_q, f"{label}_Q"), (fld_k, f"{label}_K")]: + + last_layer = max(int(k) for k in single_data.keys()) + single_tensor = single_data[last_layer].float() + multi_tensor = multi_data[last_layer].float() + num_sequences = cu_seqlens_single.numel() - 1 + + worst_max_diff = 0.0 + worst_single_flat = worst_multi_flat = None + + for seq_index in range(num_sequences): + single_seq = _slice_sequence(single_tensor, cu_seqlens_single, seq_index) + if single_seq.shape[0] == 0: + continue + single_flat = single_seq.reshape(single_seq.shape[0], -1) + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + multi_seq = _slice_sequence(multi_tensor, cu_seqlens_multi, multi_offset) + if multi_seq.shape[0] != single_seq.shape[0]: + continue + multi_flat = multi_seq.reshape(multi_seq.shape[0], -1) + max_diff = float((single_flat - multi_flat).abs().max()) + if max_diff > worst_max_diff: + worst_max_diff = max_diff + worst_single_flat = single_flat + worst_multi_flat = multi_flat + + if worst_single_flat is not None: + _print_topk_vec(worst_single_flat[0].cpu(), worst_multi_flat[0].cpu(), + topk, sort_err, f"{label}_L{last_layer}_token0") + + +def _print_topk_kv( + dir_single: str, dir_stacked: str, + cu_seqlens_single: torch.Tensor, + cu_seqlens_multi: torch.Tensor, + stack_count: int, + filename: str, field_first: str, field_second: str, + label: str, topk: int, sort_err: str, +): + """Print top-K worst dimensions for a KV-style per-layer file.""" + single_data = _load_per_layer_dict(dir_single, filename) + multi_data = _load_per_layer_dict(dir_stacked, filename) + if single_data is None or multi_data is None: + return + + last_layer = max(int(k) for k in single_data.keys()) + num_sequences = cu_seqlens_single.numel() - 1 + + for field, tag in [(field_first, f"{label}_{field_first}"), + (field_second, f"{label}_{field_second}")]: try: - sq = sd[lyr][fld].float() - mq = md[lyr][fld].float() - except (KeyError, TypeError, AttributeError) as e: - print(f" [top-K skip] {label} {fld}: {e}"); continue - worst_md = 0.0; worst_on = worst_off = None - for i in range(B): - s_seq = sq[int(cu_s[i]):int(cu_s[i + 1])] - if s_seq.shape[0] == 0: continue - on_f = s_seq.reshape(s_seq.shape[0], -1) - for k in range(n_copies): - j = i + k * B - ct = mq[int(cu_m[j]):int(cu_m[j + 1])] - if ct.shape[0] != s_seq.shape[0]: continue - off_f = ct.reshape(ct.shape[0], -1) - md = float((on_f - off_f).abs().max()) - if md > worst_md: worst_md = md; worst_on = on_f; worst_off = off_f - if worst_on is not None: - _print_topk_vec(worst_on[0].cpu(), worst_off[0].cpu(), - topk, sort_err, f"{tag}_L{lyr}_token0") - - -# ── print wrappers (replace ON/OFF labels with Single/Stacked) ── - -def _print_plain_baseline(r: CheckResult): - print(_SEP_SINGLE + f"\n [{r.name}] Single vs Stacked（per-seq aligned）") + single_field = single_data[last_layer][field].float() + multi_field = multi_data[last_layer][field].float() + except (KeyError, TypeError, AttributeError) as exc: + print(f" [top-K skip] {label} {field}: {exc}") + continue + + worst_max_diff = 0.0 + worst_single_flat = worst_multi_flat = None + + for seq_index in range(num_sequences): + single_seq = _slice_sequence(single_field, cu_seqlens_single, seq_index) + if single_seq.shape[0] == 0: + continue + single_flat = single_seq.reshape(single_seq.shape[0], -1) + + for copy_index in range(stack_count): + multi_offset = seq_index + copy_index * num_sequences + multi_seq = _slice_sequence(multi_field, cu_seqlens_multi, multi_offset) + if multi_seq.shape[0] != single_seq.shape[0]: + continue + multi_flat = multi_seq.reshape(multi_seq.shape[0], -1) + max_diff = float((single_flat - multi_flat).abs().max()) + if max_diff > worst_max_diff: + worst_max_diff = max_diff + worst_single_flat = single_flat + worst_multi_flat = multi_flat + + if worst_single_flat is not None: + _print_topk_vec(worst_single_flat[0].cpu(), worst_multi_flat[0].cpu(), + topk, sort_err, f"{tag}_L{last_layer}_token0") + + +# ══════════════════════════════════════════════════════════════════ +# Print wrappers (Single/Stacked labels instead of ON/OFF) +# ══════════════════════════════════════════════════════════════════ + +def _print_table_baseline(result: CheckResult): + """Print a per-layer comparison table with Single/Stacked labels.""" + print(_SEP_SINGLE + f"\n [{result.name}] Single vs Stacked（per-seq aligned）") print(_SEP_SINGLE) - m = r.metrics - if "error" in m: - print(f" {_CROSS} {m['error']}\n"); return - layers = m.get("layers", {}) - print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " - f"{'SNG_T':>8s} {'STK_T':>8s} {'STATUS':>8s}") + metrics = result.metrics + if "error" in metrics: + print(f" {_CROSS} {metrics['error']}\n") + return + + layers = metrics.get("layers", {}) + header = (f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " + f"{'SNG_T':>8s} {'STK_T':>8s} {'STATUS':>8s}") + print(header) print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") - for lyr in sorted(layers): - d = layers[lyr] - if "max_diff" not in d: - print(f" {lyr:>6d} {d.get('error','')}"); continue - md = d["max_diff"]; ca = d.get("cos_avg", 0); cm = d.get("cos_min", 1) - ok = md < 1e-5 - print(f" {lyr:>6d} {md:>12.3e} {ca:>10.6f} {cm:>10.6f} " - f"{d.get('on_T','—'):>8} {d.get('off_T','—'):>8} " + + for layer_index in sorted(layers): + entry = layers[layer_index] + if "max_diff" not in entry: + print(f" {layer_index:>6d} {entry.get('error', '')}") + continue + max_diff = entry["max_diff"] + cos_avg = entry.get("cos_avg", 0.0) + cos_min = entry.get("cos_min", 1.0) + single_tokens = entry.get("on_T", "—") + stacked_tokens = entry.get("off_T", "—") + ok = max_diff < 1e-5 + print(f" {layer_index:>6d} {max_diff:>12.3e} {cos_avg:>10.6f} {cos_min:>10.6f} " + f"{single_tokens:>8} {stacked_tokens:>8} " f"{'OK' if ok else 'DIFF':>8s}") - print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " - f"{_CHECK if r.passed else _CROSS}") + + print(f"\n max_diff={metrics.get('max_diff')} cos_min={metrics.get('cos_min')} " + f"{_CHECK if result.passed else _CROSS}") print() -def _print_rope_baseline(r: CheckResult, label: str): - _sec = label - print(_SEP_SINGLE + f"\n [{r.name}] {_sec} Single vs Stacked") +def _print_kv_table_baseline(result: CheckResult, label: str): + """Print a per-layer Q/K or K/V comparison table.""" + print(_SEP_SINGLE + f"\n [{result.name}] {label} Single vs Stacked") print(_SEP_SINGLE) - layers = r.metrics.get("layers") - if isinstance(layers, dict): - print(f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " + layers = result.metrics.get("layers") + if not isinstance(layers, dict): + return + + header = (f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " f"{'K_MAXDIFF':>12s} {'K_COS_AVG':>12s} {'K_COS_MIN':>12s} " f"{'TOKENS':>8s}") - print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " - f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") - for lyr in sorted(layers.keys()): - d = layers[lyr] - if "error" in d: - print(f" {lyr:>6d} {d['error']}"); continue - print(f" {lyr:>6d} {d.get('Q_max_diff',0.0):>12.3e} {d.get('Q_cos_avg',0.0):>12.6e} " - f"{d.get('Q_cos_min',0.0):>12.6e} " - f"{d.get('K_max_diff',0.0):>12.3e} {d.get('K_cos_avg',0.0):>12.6e} " - f"{d.get('K_cos_min',0.0):>12.6e} {d.get('n_tokens','—'):>8}") + print(header) + print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " + f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") + + for layer_index in sorted(layers.keys()): + entry = layers[layer_index] + if "error" in entry: + print(f" {layer_index:>6d} {entry['error']}") + continue + print(f" {layer_index:>6d} " + f"{entry.get('Q_max_diff', 0.0):>12.3e} {entry.get('Q_cos_avg', 0.0):>12.6f} " + f"{entry.get('Q_cos_min', 0.0):>12.6f} " + f"{entry.get('K_max_diff', 0.0):>12.3e} {entry.get('K_cos_avg', 0.0):>12.6f} " + f"{entry.get('K_cos_min', 0.0):>12.6f} {entry.get('n_tokens', '—'):>8}") print() -# ── main ───────────────────────────────────────────────────────── +# ══════════════════════════════════════════════════════════════════ +# Main +# ══════════════════════════════════════════════════════════════════ def main(): - ap = argparse.ArgumentParser( + parser = argparse.ArgumentParser( description="GEMM precision baseline — cross-batch-size (single vs N copies)", epilog=__doc__, ) - ap.add_argument("--dir-single", required=True, - help="Single-copy dump directory (batch=[A])") - ap.add_argument("--dir-stacked", required=True, - help="Stacked-copies dump directory (batch=[A x N])") - ap.add_argument("--num-copies", type=int, required=True, - help="Number of stacked copies N") - ap.add_argument("--layer", type=int, default=None, - help="Compare specific layer 1-indexed (default: all)") - ap.add_argument("--tag", default="old", - help="2D file tag for logprobs/entropy (default: old)") - ap.add_argument("--atol", type=float, default=1e-5, - help="Absolute tolerance for 2D (default: 1e-5)") - ap.add_argument("--topk", type=int, default=0, - help="top-K worst dims for packed-token (0=disabled)") - ap.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", - help="top-K sort: abs / rel / val") - ap.add_argument("--output", "-o", default=None, - help="Write JSON report to this path") - args = ap.parse_args() - - cu_s = _load_cu(args.dir_single) - cu_m = _load_cu(args.dir_stacked) - if cu_s is None or cu_m is None: - print(f"{_CROSS} cu_seqlens_q.pt missing"); return 1 - - T = int(cu_s[-1]) - n = args.num_copies - num_seq = cu_s.numel() - 1 + parser.add_argument("--dir-single", required=True, + help="Single-copy dump directory") + parser.add_argument("--dir-stacked", required=True, + help="Stacked-copies dump directory") + parser.add_argument("--num-copies", type=int, required=True, + help="Number of stacked copies (stack count)") + parser.add_argument("--layer", type=int, default=None, + help="Compare specific layer (1-indexed, default: all)") + parser.add_argument("--tag", default="old", + help="2D file tag for logprobs/entropy (default: old)") + parser.add_argument("--atol", type=float, default=1e-5, + help="Absolute tolerance for 2D (default: 1e-5)") + parser.add_argument("--topk", type=int, default=0, + help="top-K worst dims for packed token (0=disabled)") + parser.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", + help="top-K sort order: abs / rel / val") + parser.add_argument("--output", "-o", default=None, + help="Write JSON report to this path") + args = parser.parse_args() + + cu_seqlens_single = _load_cu_seqlens(args.dir_single) + cu_seqlens_multi = _load_cu_seqlens(args.dir_stacked) + if cu_seqlens_single is None or cu_seqlens_multi is None: + print(f"{_CROSS} cu_seqlens_q.pt missing") + return 1 + + stack_count = args.num_copies + num_sequences = cu_seqlens_single.numel() - 1 + total_tokens_single = int(cu_seqlens_single[-1]) + print(_SEP_DOUBLE) print(" GEMM Precision Baseline — Cross-Batch-Size Comparison") - print(f" Single: {args.dir_single} (1 seq, {T} tokens)") - print(f" Stacked: {args.dir_stacked} ({n} copies, {n * T} tokens)") + print(f" Single : {args.dir_single} ({num_sequences} seqs, {total_tokens_single} tokens)") + print(f" Stacked: {args.dir_stacked} ({num_sequences * stack_count} seqs, " + f"{total_tokens_single * stack_count} tokens, {stack_count}x stack)") print(_SEP_DOUBLE) - # shape diagnostics + # Shape diagnostics _print_shapes(args.dir_single, args.dir_stacked, args.tag) all_results: list[CheckResult] = [] - # ── hidden_states ── - r = _compare_plain(args.dir_single, args.dir_stacked, - "hidden_states.pt", cu_s, cu_m, n, args.layer, "hidden_states") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── build_kv_input_v ── - r = _compare_plain(args.dir_single, args.dir_stacked, - "build_kv_input_v.pt", cu_s, cu_m, n, args.layer, "build_kv_input_v") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── rope_preqk ── - r = _compare_rope(args.dir_single, args.dir_stacked, "rope_preqk.pt", - cu_s, cu_m, n, args.layer, "rope_preqk", "query", "key") - if r: all_results.append(r); _print_rope_baseline(r, "rope_preqk") - - # ── rope_freqs ── - r = _compare_plain(args.dir_single, args.dir_stacked, - "rope_freqs.pt", cu_s, cu_m, n, args.layer, "rope_freqs") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── rope_postqk ── - r = _compare_rope(args.dir_single, args.dir_stacked, "rope_postqk.pt", - cu_s, cu_m, n, args.layer, "rope_postqk", "query", "key") - if r: all_results.append(r); _print_rope_baseline(r, "rope_postqk") - - # ── attn_outputs ── - r = _compare_plain(args.dir_single, args.dir_stacked, - "attn_outputs.pt", cu_s, cu_m, n, args.layer, "attn_outputs") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── full_kv ── - r = _compare_rope(args.dir_single, args.dir_stacked, "full_kv.pt", - cu_s, cu_m, n, args.layer, "full_kv", "key", "value") - if r: all_results.append(r); _print_rope_baseline(r, "full_kv") - - # ── logits ── - r = _compare_logits(args.dir_single, args.dir_stacked, cu_s, cu_m, n) - if r: all_results.append(r); _print_logits_packed(r) - - # ── 2D: logprobs + entropy ── - _2d_results: list[tuple[str, torch.Tensor, torch.Tensor]] = [] - for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: - fname = f"{fn}_{args.tag}.pt" - r, t1, t2 = _compare_2d_file(args.dir_single, args.dir_stacked, - fname, f"{cn}_{args.tag}", n, num_seq, args.atol) - if r: all_results.append(r); _print_2d_result(r) - if t1 is not None and t2 is not None: - _2d_results.append((f"{cn}_{args.tag}", t1, t2)) - - # ── top-K ── + # ── Per-layer plain dicts ── + for filename, label in [ + ("hidden_states.pt", "hidden_states"), + ("build_kv_input_v.pt", "build_kv_input_v"), + ("rope_freqs.pt", "rope_freqs"), + ("attn_outputs.pt", "attn_outputs"), + ]: + result = _compare_plain_per_layer( + args.dir_single, args.dir_stacked, filename, + cu_seqlens_single, cu_seqlens_multi, stack_count, + args.layer, label, + ) + if result: + all_results.append(result) + _print_table_baseline(result) + + # ── Per-layer KV dicts ── + for filename, label, field_a, field_b in [ + ("rope_preqk.pt", "rope_preqk", "query", "key"), + ("rope_postqk.pt", "rope_postqk", "query", "key"), + ("full_kv.pt", "full_kv", "key", "value"), + ]: + result = _compare_kv_per_layer( + args.dir_single, args.dir_stacked, filename, + cu_seqlens_single, cu_seqlens_multi, stack_count, + args.layer, label, field_a, field_b, + ) + if result: + all_results.append(result) + _print_kv_table_baseline(result, label) + + # ── Logits ── + result = _compare_logits_cross_batch( + args.dir_single, args.dir_stacked, + cu_seqlens_single, cu_seqlens_multi, stack_count, + ) + if result: + all_results.append(result) + _print_logits_packed(result) + + # ── 2D ── + _2d_tensors: list[tuple[str, torch.Tensor, torch.Tensor]] = [] + for file_tag, compare_name in [("logprobs", "logp"), ("entropy", "entropy")]: + filename = f"{file_tag}_{args.tag}.pt" + result, single_2d, multi_2d = _compare_2d_cross_batch( + args.dir_single, args.dir_stacked, filename, + f"{compare_name}_{args.tag}", stack_count, num_sequences, args.atol, + ) + if result: + all_results.append(result) + _print_2d_result(result) + if single_2d is not None and multi_2d is not None: + _2d_tensors.append((f"{compare_name}_{args.tag}", single_2d, multi_2d)) + + # ── Top-K ── if args.topk > 0: - _topk_per_layer(args.dir_single, args.dir_stacked, cu_s, cu_m, n, - "build_kv_input_v.pt", "build_kv_input_v", - args.topk, args.sort_err) - _topk_rope(args.dir_single, args.dir_stacked, cu_s, cu_m, n, - "rope_preqk.pt", "query", "key", "rope_preqk", - args.topk, args.sort_err) - _topk_rope(args.dir_single, args.dir_stacked, cu_s, cu_m, n, - "rope_postqk.pt", "query", "key", "rope_postqk", - args.topk, args.sort_err) - # 2D top-K: compare row 0 from single vs row 0 from stacked (same shuffled seq) - for label, st, mt in _2d_results: - if st.dim() >= 2 and mt.dim() >= 2 and st.shape[1] == mt.shape[1]: - # align: single row i vs stacked row i (first copy of each sequence) - t1 = st[:num_seq]; t2 = mt[:num_seq] + _print_topk_plain( + args.dir_single, args.dir_stacked, + cu_seqlens_single, cu_seqlens_multi, stack_count, + "build_kv_input_v.pt", "build_kv_input_v", + args.topk, args.sort_err, + ) + _print_topk_kv( + args.dir_single, args.dir_stacked, + cu_seqlens_single, cu_seqlens_multi, stack_count, + "rope_preqk.pt", "query", "key", "rope_preqk", + args.topk, args.sort_err, + ) + _print_topk_kv( + args.dir_single, args.dir_stacked, + cu_seqlens_single, cu_seqlens_multi, stack_count, + "rope_postqk.pt", "query", "key", "rope_postqk", + args.topk, args.sort_err, + ) + for label, single_2d, multi_2d in _2d_tensors: + if (single_2d.dim() >= 2 and multi_2d.dim() >= 2 + and single_2d.shape[1] == multi_2d.shape[1]): + t1 = single_2d[:num_sequences] + t2 = multi_2d[:num_sequences] if t1.shape == t2.shape: _print_topk_2d(t1.cpu(), t2.cpu(), None, args.topk, args.sort_err, label) @@ -511,7 +672,7 @@ def main(): if args.output: _dump_json(all_results, args.output, args.dir_single, args.dir_stacked, - tag=f"cross_batch_N{n}", dir_off2=None) + tag=f"cross_batch_N{stack_count}", dir_off2=None) if __name__ == "__main__": diff --git a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py index 8c2027be..c3ce5c7a 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_baseline_within_batch.py @@ -9,13 +9,12 @@ export PREFIX_SHARING_DIAG_DUMP=/dump_multi # forward with batch=[A x N] - python cmp_baseline_within_batch.py --dir-multi /dump_multi --num-copies 4 + python cmp_baseline_within_batch.py --dir-multi /dump_multi --num-seq 4 """ from __future__ import annotations import argparse -import math import os import torch @@ -26,10 +25,7 @@ _SEP_SINGLE, _CHECK, _CROSS, - _COS_AVG_PASS, - _COS_MIN_PASS, _cosine_sim, - _error_abs_rel, _pearson_r, _dump_json, _load_tensor, @@ -42,376 +38,500 @@ ) -# ── helpers ────────────────────────────────────────────────────── +# ══════════════════════════════════════════════════════════════════ +# I/O helpers +# ══════════════════════════════════════════════════════════════════ -def _load_dict(dir_path: str, filename: str) -> dict | None: - fp = os.path.join(dir_path, filename) - if not os.path.exists(fp): +def _load_per_layer_dict(directory: str, filename: str) -> dict | None: + filepath = os.path.join(directory, filename) + if not os.path.exists(filepath): return None - d = torch.load(fp, weights_only=True) - return d if isinstance(d, dict) else None + data = torch.load(filepath, weights_only=True) + return data if isinstance(data, dict) else None -def _load_cu(dir_path: str) -> torch.Tensor | None: - fp = os.path.join(dir_path, "cu_seqlens_q.pt") - if not os.path.exists(fp): +def _load_cu_seqlens(directory: str) -> torch.Tensor | None: + filepath = os.path.join(directory, "cu_seqlens_q.pt") + if not os.path.exists(filepath): return None - return torch.load(fp, weights_only=True) + return torch.load(filepath, weights_only=True) -def _extract(packed: torch.Tensor, cu: torch.Tensor, copy_idx: int) -> torch.Tensor: - return packed[int(cu[copy_idx]) : int(cu[copy_idx + 1])] +def _slice_sequence(packed_tensor: torch.Tensor, + cu_seqlens: torch.Tensor, + sequence_index: int) -> torch.Tensor: + return packed_tensor[int(cu_seqlens[sequence_index]): + int(cu_seqlens[sequence_index + 1])] -def _get_layers(data: dict) -> list[int]: +def _sorted_layer_keys(data: dict) -> list[int]: return sorted(int(k) for k in data.keys()) -# ── comparison logic ───────────────────────────────────────────── - -def _worst_pairwise(copies: list[torch.Tensor]) -> dict: - """Pairwise compare all copies of the SAME length; return worst + avg cos.""" - worst_md, worst_cos_min = 0.0, 1.0 - all_cos: list[float] = [] - for i in range(len(copies)): - for j in range(i + 1, len(copies)): - a = copies[i].reshape(copies[i].shape[0], -1).float() - b = copies[j].reshape(copies[j].shape[0], -1).float() - cos = _cosine_sim(a, b, dim=-1) - all_cos.extend(cos.tolist()) - md = float((a - b).abs().max()) - worst_md = max(worst_md, md) - worst_cos_min = min(worst_cos_min, float(cos.min())) - return {"max_diff": worst_md, "cos_avg": sum(all_cos)/len(all_cos) if all_cos else 0.0, - "cos_min": worst_cos_min} - - -def _group_by_len(copies: list[torch.Tensor]) -> dict[int, list[torch.Tensor]]: - """Group copies by sequence length (tokens).""" +def _group_by_token_count(tensors: list[torch.Tensor] + ) -> dict[int, list[torch.Tensor]]: + """Group tensors by their leading dimension (token count).""" groups: dict[int, list[torch.Tensor]] = {} - for c in copies: - groups.setdefault(c.shape[0], []).append(c) + for tensor in tensors: + groups.setdefault(tensor.shape[0], []).append(tensor) return groups -def _compare_plain_within(dir_multi: str, filename: str, - cu: torch.Tensor, total_copies: int, - layer: int | None, label: str) -> CheckResult | None: - d = _load_dict(dir_multi, filename) - if d is None: +# ══════════════════════════════════════════════════════════════════ +# Pairwise comparison within a group of same-length tensors +# ══════════════════════════════════════════════════════════════════ + +def _pairwise_metrics(copies: list[torch.Tensor]) -> dict: + """Compare all pairs within a group; return worst max_diff, cos_min, + and average cos_avg.""" + worst_max_diff = 0.0 + worst_cos_min = 1.0 + all_cos_values: list[float] = [] + + for i in range(len(copies)): + for j in range(i + 1, len(copies)): + flat_i = copies[i].reshape(copies[i].shape[0], -1).float() + flat_j = copies[j].reshape(copies[j].shape[0], -1).float() + + per_token_cos = _cosine_sim(flat_i, flat_j, dim=-1) + all_cos_values.extend(per_token_cos.tolist()) + worst_max_diff = max(worst_max_diff, float((flat_i - flat_j).abs().max())) + worst_cos_min = min(worst_cos_min, float(per_token_cos.min())) + + return { + "max_diff": worst_max_diff, + "cos_avg": sum(all_cos_values) / len(all_cos_values) if all_cos_values else 0.0, + "cos_min": worst_cos_min, + } + + +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — plain per-layer dicts +# ══════════════════════════════════════════════════════════════════ + +def _compare_plain_within( + directory: str, filename: str, + cu_seqlens: torch.Tensor, total_sequences: int, + filter_layer: int | None, label: str, +) -> CheckResult | None: + """Pairwise-compare copies within a single dump for ``filename``.""" + data = _load_per_layer_dict(directory, filename) + if data is None: return None - layers = _get_layers(d) - if layer is not None: - layers = [l for l in layers if l == layer] + + layers = _sorted_layer_keys(data) + if filter_layer is not None: + layers = [l for l in layers if l == filter_layer] if not layers: return None per_layer: dict = {} - worst_md, worst_cos = 0.0, 1.0 - for lyr in layers: - mt = d[lyr].float() - copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] - groups = _group_by_len(copies) - layer_md, layer_cos_min = 0.0, 1.0 - all_cos_avg: list[float] = [] + worst_max_diff = 0.0 + worst_cos_min = 1.0 + + for layer_index in layers: + multi_tensor = data[layer_index].float() + copies = [_slice_sequence(multi_tensor, cu_seqlens, seq_index) + for seq_index in range(total_sequences)] + groups = _group_by_token_count(copies) + + layer_max_diff = 0.0 + layer_cos_min = 1.0 + group_cos_avgs: list[float] = [] + for group_copies in groups.values(): if len(group_copies) < 2: continue - w = _worst_pairwise(group_copies) - layer_md = max(layer_md, w["max_diff"]) - layer_cos_min = min(layer_cos_min, w["cos_min"]) - all_cos_avg.append(w["cos_avg"]) - per_layer[lyr] = { - "max_diff": layer_md, - "cos_avg": sum(all_cos_avg)/len(all_cos_avg) if all_cos_avg else 0.0, + metrics = _pairwise_metrics(group_copies) + layer_max_diff = max(layer_max_diff, metrics["max_diff"]) + layer_cos_min = min(layer_cos_min, metrics["cos_min"]) + group_cos_avgs.append(metrics["cos_avg"]) + + per_layer[layer_index] = { + "max_diff": layer_max_diff, + "cos_avg": (sum(group_cos_avgs) / len(group_cos_avgs) + if group_cos_avgs else 0.0), "cos_min": layer_cos_min, - "n_tokens": mt.shape[0], "on_T": mt.shape[0], "off_T": mt.shape[0], + "n_tokens": multi_tensor.shape[0], + "on_T": multi_tensor.shape[0], + "off_T": multi_tensor.shape[0], } - worst_md = max(worst_md, layer_md) - worst_cos = min(worst_cos, layer_cos_min) - passed = worst_md == 0.0 - _name = f"{label}_L{layer}" if layer is not None else label - return CheckResult(name=_name, passed=passed, - metrics={"layers": per_layer, "max_diff": worst_md, - "cos_min": worst_cos, "num_layers": len(layers)}) - - -def _compare_rope_within(dir_multi: str, filename: str, - cu: torch.Tensor, total_copies: int, - layer: int | None, label: str, - fld_q: str, fld_k: str) -> CheckResult | None: - d = _load_dict(dir_multi, filename) - if d is None: + worst_max_diff = max(worst_max_diff, layer_max_diff) + worst_cos_min = min(worst_cos_min, layer_cos_min) + + passed = worst_max_diff == 0.0 + result_name = f"{label}_L{filter_layer}" if filter_layer is not None else label + return CheckResult( + name=result_name, passed=passed, + metrics={"layers": per_layer, "max_diff": worst_max_diff, + "cos_min": worst_cos_min, "num_layers": len(layers)}, + ) + + +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — per-layer KV dicts +# ══════════════════════════════════════════════════════════════════ + +def _compare_kv_within( + directory: str, filename: str, + cu_seqlens: torch.Tensor, total_sequences: int, + filter_layer: int | None, label: str, + field_first: str, field_second: str, +) -> CheckResult | None: + """Pairwise-compare ``{layer: {field_first, field_second}}`` within a dump.""" + data = _load_per_layer_dict(directory, filename) + if data is None: return None - layers = _get_layers(d) - if layer is not None: - layers = [l for l in layers if l == layer] + + layers = _sorted_layer_keys(data) + if filter_layer is not None: + layers = [l for l in layers if l == filter_layer] if not layers: return None per_layer: dict = {} - for lyr in layers: - mq = d[lyr][fld_q].float(); mk = d[lyr][fld_k].float() - q_copies = [mq[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] - k_copies = [mk[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] - q_groups = _group_by_len(q_copies) - k_groups = _group_by_len(k_copies) - qw = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} - kw = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} - q_avgs, k_avgs = [], [] - for g in q_groups.values(): - if len(g) >= 2: - w = _worst_pairwise(g) - qw["max_diff"] = max(qw["max_diff"], w["max_diff"]) - qw["cos_min"] = min(qw["cos_min"], w["cos_min"]) - q_avgs.append(w["cos_avg"]) - for g in k_groups.values(): - if len(g) >= 2: - w = _worst_pairwise(g) - kw["max_diff"] = max(kw["max_diff"], w["max_diff"]) - kw["cos_min"] = min(kw["cos_min"], w["cos_min"]) - k_avgs.append(w["cos_avg"]) - per_layer[lyr] = { - "Q_max_diff": qw["max_diff"], "K_max_diff": kw["max_diff"], - "Q_cos_avg": sum(q_avgs)/len(q_avgs) if q_avgs else 0.0, - "Q_cos_min": qw["cos_min"], - "K_cos_avg": sum(k_avgs)/len(k_avgs) if k_avgs else 0.0, - "K_cos_min": kw["cos_min"], - "n_tokens": sum(g[0].shape[0] * len(g) for g in q_groups.values()), + + for layer_index in layers: + multi_first = data[layer_index][field_first].float() + multi_second = data[layer_index][field_second].float() + first_copies = [_slice_sequence(multi_first, cu_seqlens, seq_index) + for seq_index in range(total_sequences)] + second_copies = [_slice_sequence(multi_second, cu_seqlens, seq_index) + for seq_index in range(total_sequences)] + + first_groups = _group_by_token_count(first_copies) + second_groups = _group_by_token_count(second_copies) + + first_worst = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} + second_worst = {"max_diff": 0.0, "cos_min": 1.0, "cos_avg": 0.0} + first_avgs, second_avgs = [], [] + + for group in first_groups.values(): + if len(group) >= 2: + metrics = _pairwise_metrics(group) + first_worst["max_diff"] = max(first_worst["max_diff"], metrics["max_diff"]) + first_worst["cos_min"] = min(first_worst["cos_min"], metrics["cos_min"]) + first_avgs.append(metrics["cos_avg"]) + + for group in second_groups.values(): + if len(group) >= 2: + metrics = _pairwise_metrics(group) + second_worst["max_diff"] = max(second_worst["max_diff"], metrics["max_diff"]) + second_worst["cos_min"] = min(second_worst["cos_min"], metrics["cos_min"]) + second_avgs.append(metrics["cos_avg"]) + + per_layer[layer_index] = { + "Q_max_diff": first_worst["max_diff"], + "K_max_diff": second_worst["max_diff"], + "Q_cos_avg": sum(first_avgs) / len(first_avgs) if first_avgs else 0.0, + "Q_cos_min": first_worst["cos_min"], + "K_cos_avg": sum(second_avgs) / len(second_avgs) if second_avgs else 0.0, + "K_cos_min": second_worst["cos_min"], + "n_tokens": sum(g[0].shape[0] * len(g) for g in first_groups.values()), } - _name = f"{label}_L{layer}" if layer is not None else label - return CheckResult(name=_name, passed=True, metrics={"layers": per_layer}) + + result_name = f"{label}_L{filter_layer}" if filter_layer is not None else label + return CheckResult(name=result_name, passed=True, metrics={"layers": per_layer}) -def _compare_logits_within(dir_multi: str, cu: torch.Tensor, - total_copies: int) -> CheckResult | None: - fp = os.path.join(dir_multi, "logits.pt") - if not os.path.exists(fp): +# ══════════════════════════════════════════════════════════════════ +# Comparison logic — logits +# ══════════════════════════════════════════════════════════════════ + +def _compare_logits_within( + directory: str, + cu_seqlens: torch.Tensor, + total_sequences: int, +) -> CheckResult | None: + """Pairwise-compare packed logits within a dump.""" + filepath = os.path.join(directory, "logits.pt") + if not os.path.exists(filepath): return None - mt = torch.load(fp, weights_only=True).float() - mt = mt.reshape(-1, mt.size(-1)) - copies = [mt[int(cu[i]):int(cu[i + 1])] for i in range(total_copies)] - groups = _group_by_len(copies) - worst_md, worst_cos_min = 0.0, 1.0 - all_cos_avg: list[float] = [] - for g in groups.values(): - if len(g) >= 2: - w = _worst_pairwise(g) - worst_md = max(worst_md, w["max_diff"]) - worst_cos_min = min(worst_cos_min, w["cos_min"]) - all_cos_avg.append(w["cos_avg"]) - return CheckResult(name="logits", passed=worst_md == 0.0, - metrics={"n_tokens": copies[0].shape[0] if copies else 0, - "cos_avg": sum(all_cos_avg)/len(all_cos_avg) if all_cos_avg else 0.0, - "cos_min": worst_cos_min}) - - -# ── print wrappers ───────────────────────────────────────────── - -def _print_plain_baseline(r: CheckResult): - print(_SEP_SINGLE + f"\n [{r.name}] within-batch pairwise") + + multi_logits = torch.load(filepath, weights_only=True).float() + multi_logits = multi_logits.reshape(-1, multi_logits.size(-1)) + copies = [_slice_sequence(multi_logits, cu_seqlens, seq_index) + for seq_index in range(total_sequences)] + groups = _group_by_token_count(copies) + + worst_max_diff = 0.0 + worst_cos_min = 1.0 + all_cos_avgs: list[float] = [] + + for group in groups.values(): + if len(group) >= 2: + metrics = _pairwise_metrics(group) + worst_max_diff = max(worst_max_diff, metrics["max_diff"]) + worst_cos_min = min(worst_cos_min, metrics["cos_min"]) + all_cos_avgs.append(metrics["cos_avg"]) + + return CheckResult( + name="logits", passed=worst_max_diff == 0.0, + metrics={ + "n_tokens": copies[0].shape[0] if copies else 0, + "cos_avg": (sum(all_cos_avgs) / len(all_cos_avgs) + if all_cos_avgs else 0.0), + "cos_min": worst_cos_min, + }, + ) + + +# ══════════════════════════════════════════════════════════════════ +# Print wrappers (within-batch labels) +# ══════════════════════════════════════════════════════════════════ + +def _print_table_baseline(result: CheckResult): + """Print a per-layer comparison table for within-batch.""" + print(_SEP_SINGLE + f"\n [{result.name}] within-batch pairwise") print(_SEP_SINGLE) - m = r.metrics - if "error" in m: - print(f" {_CROSS} {m['error']}\n"); return - layers = m.get("layers", {}) - print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " - f"{'TOKENS':>8s} {'STATUS':>8s}") + metrics = result.metrics + if "error" in metrics: + print(f" {_CROSS} {metrics['error']}\n") + return + + layers = metrics.get("layers", {}) + header = (f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS_AVG':>10s} {'COS_MIN':>10s} " + f"{'TOKENS':>8s} {'STATUS':>8s}") + print(header) print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 10} {'─' * 8} {'─' * 8}") - for lyr in sorted(layers): - d = layers[lyr] - if "max_diff" not in d: - print(f" {lyr:>6d} {d.get('error','')}"); continue - md = d["max_diff"]; ca = d.get("cos_avg", 0); cm = d.get("cos_min", 1) - ok = md == 0.0 - print(f" {lyr:>6d} {md:>12.3e} {ca:>10.6f} {cm:>10.6f} " - f"{d.get('n_tokens','—'):>8} " - f"{'PASS' if ok else 'DIFF':>8s}") - print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " - f"{_CHECK if r.passed else _CROSS}") + + for layer_index in sorted(layers): + entry = layers[layer_index] + if "max_diff" not in entry: + print(f" {layer_index:>6d} {entry.get('error', '')}") + continue + max_diff = entry["max_diff"] + cos_avg = entry.get("cos_avg", 0.0) + cos_min = entry.get("cos_min", 1.0) + tokens = entry.get("n_tokens", "—") + ok = max_diff == 0.0 + print(f" {layer_index:>6d} {max_diff:>12.3e} {cos_avg:>10.6f} {cos_min:>10.6f} " + f"{tokens:>8} {'PASS' if ok else 'DIFF':>8s}") + + print(f"\n max_diff={metrics.get('max_diff')} cos_min={metrics.get('cos_min')} " + f"{_CHECK if result.passed else _CROSS}") print() -def _print_rope_baseline(r: CheckResult, label: str): - print(_SEP_SINGLE + f"\n [{r.name}] {label} within-batch pairwise") +def _print_kv_table_baseline(result: CheckResult, label: str): + """Print a per-layer Q/K or K/V comparison table for within-batch.""" + print(_SEP_SINGLE + f"\n [{result.name}] {label} within-batch pairwise") print(_SEP_SINGLE) - layers = r.metrics.get("layers") - if isinstance(layers, dict): - print(f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " + layers = result.metrics.get("layers") + if not isinstance(layers, dict): + return + + header = (f" {'LAYER':>6s} {'Q_MAXDIFF':>12s} {'Q_COS_AVG':>12s} {'Q_COS_MIN':>12s} " f"{'K_MAXDIFF':>12s} {'K_COS_AVG':>12s} {'K_COS_MIN':>12s} " f"{'TOKENS':>8s}") - print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " - f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") - for lyr in sorted(layers.keys()): - d = layers[lyr] - if "error" in d: - print(f" {lyr:>6d} {d['error']}"); continue - print(f" {lyr:>6d} {d.get('Q_max_diff',0.0):>12.3e} {d.get('Q_cos_avg',0.0):>12.6e} " - f"{d.get('Q_cos_min',0.0):>12.6e} " - f"{d.get('K_max_diff',0.0):>12.3e} {d.get('K_cos_avg',0.0):>12.6e} " - f"{d.get('K_cos_min',0.0):>12.6e} {d.get('n_tokens','—'):>8}") + print(header) + print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} " + f"{'─' * 12} {'─' * 12} {'─' * 12} {'─' * 8}") + + for layer_index in sorted(layers.keys()): + entry = layers[layer_index] + if "error" in entry: + print(f" {layer_index:>6d} {entry['error']}") + continue + print(f" {layer_index:>6d} " + f"{entry.get('Q_max_diff', 0.0):>12.3e} {entry.get('Q_cos_avg', 0.0):>12.6f} " + f"{entry.get('Q_cos_min', 0.0):>12.6f} " + f"{entry.get('K_max_diff', 0.0):>12.3e} {entry.get('K_cos_avg', 0.0):>12.6f} " + f"{entry.get('K_cos_min', 0.0):>12.6f} {entry.get('n_tokens', '—'):>8}") print() -# ── main ───────────────────────────────────────────────────────── +# ══════════════════════════════════════════════════════════════════ +# Main +# ══════════════════════════════════════════════════════════════════ def main(): - ap = argparse.ArgumentParser( + parser = argparse.ArgumentParser( description="GEMM precision baseline — within-batch pairwise comparison", epilog=__doc__, ) - ap.add_argument("--dir-multi", required=True, - help="Multi-copy dump directory (batch=[A x N])") - ap.add_argument("--num-copies", type=int, default=None, - help="Total sequences in batch (auto-detected from cu_seqlens)") - ap.add_argument("--num-seq", type=int, default=None, - help="Distinct sequences before stacking (required for 2D within-batch)") - ap.add_argument("--layer", type=int, default=None, - help="Compare specific layer 1-indexed (default: all)") - ap.add_argument("--tag", default="old", - help="2D file tag for logprobs/entropy (default: old)") - ap.add_argument("--atol", type=float, default=1e-5, - help="Absolute tolerance for 2D (default: 1e-5)") - ap.add_argument("--topk", type=int, default=0, - help="top-K worst dims (0=disabled)") - ap.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", - help="top-K sort: abs / rel / val") - ap.add_argument("--output", "-o", default=None, - help="Write JSON report to this path") - args = ap.parse_args() - - cu = _load_cu(args.dir_multi) - if cu is None: - print(f"{_CROSS} cu_seqlens_q.pt missing"); return 1 - - total_copies = cu.numel() - 1 - if args.num_copies is not None and args.num_copies != total_copies: - print(f"[warn] --num-copies={args.num_copies} but cu_seqlens has {total_copies} copies; using {total_copies}") - lengths = [int(cu[i + 1]) - int(cu[i]) for i in range(total_copies)] + parser.add_argument("--dir-multi", required=True, + help="Multi-copy dump directory") + parser.add_argument("--num-seq", type=int, default=None, + help="Number of distinct sequences (required for 2D within-batch)") + parser.add_argument("--layer", type=int, default=None, + help="Compare specific layer (1-indexed, default: all)") + parser.add_argument("--tag", default="old", + help="2D file tag for logprobs/entropy (default: old)") + parser.add_argument("--atol", type=float, default=1e-5, + help="Absolute tolerance for 2D (default: 1e-5)") + parser.add_argument("--topk", type=int, default=0, + help="top-K worst dims (0=disabled)") + parser.add_argument("--sort-err", choices=["abs", "rel", "val"], default="abs", + help="top-K sort order: abs / rel / val") + parser.add_argument("--output", "-o", default=None, + help="Write JSON report to this path") + args = parser.parse_args() + + cu_seqlens = _load_cu_seqlens(args.dir_multi) + if cu_seqlens is None: + print(f"{_CROSS} cu_seqlens_q.pt missing") + return 1 + + total_sequences = cu_seqlens.numel() - 1 + lengths = [int(cu_seqlens[i + 1]) - int(cu_seqlens[i]) + for i in range(total_sequences)] + print(_SEP_DOUBLE) print(" GEMM Precision Baseline — Within-Batch Pairwise Comparison") - print(f" Directory: {args.dir_multi} ({total_copies} sequences)") - print(f" Lengths: {lengths}") + print(f" Directory: {args.dir_multi} ({total_sequences} sequences)") + print(f" Sequence lengths: {lengths}") print(_SEP_DOUBLE) _print_shapes(args.dir_multi, args.dir_multi, args.tag) all_results: list[CheckResult] = [] - # ── hidden_states ── - r = _compare_plain_within(args.dir_multi, "hidden_states.pt", - cu, total_copies, args.layer, "hidden_states") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── build_kv_input_v ── - r = _compare_plain_within(args.dir_multi, "build_kv_input_v.pt", - cu, total_copies, args.layer, "build_kv_input_v") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── rope_preqk ── - r = _compare_rope_within(args.dir_multi, "rope_preqk.pt", - cu, total_copies, args.layer, "rope_preqk", "query", "key") - if r: all_results.append(r); _print_rope_baseline(r, r.name) - - # ── rope_freqs ── - r = _compare_plain_within(args.dir_multi, "rope_freqs.pt", - cu, total_copies, args.layer, "rope_freqs") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── rope_postqk ── - r = _compare_rope_within(args.dir_multi, "rope_postqk.pt", - cu, total_copies, args.layer, "rope_postqk", "query", "key") - if r: all_results.append(r); _print_rope_baseline(r, r.name) - - # ── attn_outputs ── - r = _compare_plain_within(args.dir_multi, "attn_outputs.pt", - cu, total_copies, args.layer, "attn_outputs") - if r: all_results.append(r); _print_plain_baseline(r) - - # ── full_kv ── - r = _compare_rope_within(args.dir_multi, "full_kv.pt", - cu, total_copies, args.layer, "full_kv", "key", "value") - if r: all_results.append(r); _print_rope_baseline(r, r.name) - - # ── logits ── - r = _compare_logits_within(args.dir_multi, cu, total_copies) - if r: all_results.append(r); _print_logits_packed(r) - - # ── 2D: logprobs + entropy ── - _2d_within: list[tuple[str, torch.Tensor]] = [] - _num_seq = args.num_seq if args.num_seq else (total_copies // (args.num_copies or total_copies) if args.num_copies else total_copies) - for fn, cn in [("logprobs", "logp"), ("entropy", "entropy")]: - fname = f"{fn}_{args.tag}.pt" - st = _load_tensor(args.dir_multi, fname) - if st is not None and st.dim() >= 2: - B = st.shape[0] - stack = B // _num_seq if _num_seq > 0 and B % _num_seq == 0 else 1 - all_abs, all_rel = [], [] - worst_md, worst_cos, worst_rel = 0.0, 1.0, 0.0 - # only compare stack copies of the SAME sequence (i vs i+k*num_seq) - for i in range(_num_seq): - if i >= B: break - a = st[i].float().reshape(-1) - for k in range(1, stack): - j = i + k * _num_seq - if j >= B: continue - b = st[j].float().reshape(-1) - diff = (a - b).abs() - rel = diff / a.abs().clamp(min=1e-8) - all_abs.extend(diff.tolist()); all_rel.extend(rel.tolist()) - md = float(diff.max()); rm = float(rel.max()) - worst_md = max(worst_md, md) - worst_rel = max(worst_rel, rm) - worst_cos = min(worst_cos, float(_cosine_sim(a, b, dim=-1))) - r = CheckResult(name=f"{cn}_{args.tag}", - passed=worst_md == 0.0, - metrics={"shape": tuple(st.shape), - "active": _num_seq * st.shape[1] if st.dim() >= 2 else st.numel(), - "abs_max": worst_md, - "abs_mean": sum(all_abs)/len(all_abs) if all_abs else 0.0, - "rel_max": worst_rel, - "rel_mean": sum(all_rel)/len(all_rel) if all_rel else 0.0, - "pearson_r": _pearson_r(st[:_num_seq].float().reshape(-1), - st[_num_seq:2*_num_seq].float().reshape(-1)) - if B >= 2*_num_seq else 1.0, - "atol": args.atol}) - all_results.append(r) - _print_2d_result(r) - _2d_within.append((f"{cn}_{args.tag}", st)) - - # ── top-K ── + # ── Per-layer plain dicts ── + for filename, label in [ + ("hidden_states.pt", "hidden_states"), + ("build_kv_input_v.pt", "build_kv_input_v"), + ("rope_freqs.pt", "rope_freqs"), + ("attn_outputs.pt", "attn_outputs"), + ]: + result = _compare_plain_within( + args.dir_multi, filename, cu_seqlens, total_sequences, + args.layer, label, + ) + if result: + all_results.append(result) + _print_table_baseline(result) + + # ── Per-layer KV dicts ── + for filename, label, field_a, field_b in [ + ("rope_preqk.pt", "rope_preqk", "query", "key"), + ("rope_postqk.pt", "rope_postqk", "query", "key"), + ("full_kv.pt", "full_kv", "key", "value"), + ]: + result = _compare_kv_within( + args.dir_multi, filename, cu_seqlens, total_sequences, + args.layer, label, field_a, field_b, + ) + if result: + all_results.append(result) + _print_kv_table_baseline(result, label) + + # ── Logits ── + result = _compare_logits_within(args.dir_multi, cu_seqlens, total_sequences) + if result: + all_results.append(result) + _print_logits_packed(result) + + # ── 2D ── + _2d_tensors: list[tuple[str, torch.Tensor]] = [] + num_sequences = args.num_seq or total_sequences + + for file_tag, compare_name in [("logprobs", "logp"), ("entropy", "entropy")]: + filename = f"{file_tag}_{args.tag}.pt" + tensor_2d = _load_tensor(args.dir_multi, filename) + if tensor_2d is None or tensor_2d.dim() < 2: + continue + tensor_2d = tensor_2d.float() + batch_size = tensor_2d.shape[0] + stack = (batch_size // num_sequences + if num_sequences > 0 and batch_size % num_sequences == 0 else 1) + + all_abs_diffs: list[float] = [] + all_rel_diffs: list[float] = [] + worst_max_diff = 0.0 + worst_cos_min = 1.0 + worst_rel_max = 0.0 + + # Only compare stack copies of the SAME logical sequence + for seq_index in range(num_sequences): + if seq_index >= batch_size: + break + row_a = tensor_2d[seq_index].reshape(-1) + for copy_index in range(1, stack): + row_b_offset = seq_index + copy_index * num_sequences + if row_b_offset >= batch_size: + continue + row_b = tensor_2d[row_b_offset].reshape(-1) + + abs_diff = (row_a - row_b).abs() + rel_diff = abs_diff / row_a.abs().clamp(min=1e-8) + all_abs_diffs.extend(abs_diff.tolist()) + all_rel_diffs.extend(rel_diff.tolist()) + worst_max_diff = max(worst_max_diff, float(abs_diff.max())) + worst_rel_max = max(worst_rel_max, float(rel_diff.max())) + worst_cos_min = min(worst_cos_min, + float(_cosine_sim(row_a, row_b, dim=-1))) + + # Pearson on first two stack copies + pearson_val = 1.0 + if batch_size >= 2 * num_sequences: + pearson_val = _pearson_r( + tensor_2d[:num_sequences].float().reshape(-1), + tensor_2d[num_sequences:2 * num_sequences].float().reshape(-1), + ) + + result = CheckResult( + name=f"{compare_name}_{args.tag}", + passed=worst_max_diff == 0.0, + metrics={ + "shape": tuple(tensor_2d.shape), + "active": num_sequences * tensor_2d.shape[1], + "abs_max": worst_max_diff, + "abs_mean": (sum(all_abs_diffs) / len(all_abs_diffs) + if all_abs_diffs else 0.0), + "rel_max": worst_rel_max, + "rel_mean": (sum(all_rel_diffs) / len(all_rel_diffs) + if all_rel_diffs else 0.0), + "pearson_r": pearson_val, + "atol": args.atol, + }, + ) + all_results.append(result) + _print_2d_result(result) + _2d_tensors.append((f"{compare_name}_{args.tag}", tensor_2d)) + + # ── Top-K ── if args.topk > 0: - # 2D top-K: compare row 0 vs row _num_seq (stack copies of same seq) - for label, st2d in _2d_within: - ns = _num_seq if _num_seq > 0 and (st2d.shape[0] % _num_seq == 0) else st2d.shape[0] - if st2d.shape[0] >= ns + 1: - _print_topk_2d(st2d[:1].cpu(), st2d[ns:ns+1].cpu(), None, - args.topk, args.sort_err, label) + # 2D top-K: compare row 0 vs row num_seq (stack copies of same seq) + for label, tensor_2d in _2d_tensors: + ns = (num_sequences if num_sequences > 0 + and tensor_2d.shape[0] % num_sequences == 0 + else tensor_2d.shape[0]) + if tensor_2d.shape[0] >= ns + 1: + _print_topk_2d(tensor_2d[:1].cpu(), tensor_2d[ns:ns + 1].cpu(), + None, args.topk, args.sort_err, label) + # build_kv_input_v top-K - d = _load_dict(args.dir_multi, "build_kv_input_v.pt") - if d: - lyr = max(int(k) for k in d.keys()) - mt = d[lyr].float() - copies = [mt[int(cu[i]):int(cu[i+1])].reshape(-1) for i in range(total_copies)] - groups = _group_by_len(copies) - worst_md = 0.0; worst_a = worst_b = None - for g in groups.values(): - if len(g) < 2: continue - for i in range(len(g)): - for j in range(i + 1, len(g)): - md = float((g[i] - g[j]).abs().max()) - if md > worst_md: - worst_md = md; worst_a = g[i]; worst_b = g[j] + data = _load_per_layer_dict(args.dir_multi, "build_kv_input_v.pt") + if data: + last_layer = max(int(k) for k in data.keys()) + multi_tensor = data[last_layer].float() + copies = [_slice_sequence(multi_tensor, cu_seqlens, seq_index) + for seq_index in range(total_sequences)] + groups = _group_by_token_count(copies) + worst_max_diff = 0.0 + worst_a = worst_b = None + for group in groups.values(): + if len(group) < 2: + continue + for i in range(len(group)): + for j in range(i + 1, len(group)): + max_diff = float((group[i].reshape(-1) - + group[j].reshape(-1)).abs().max()) + if max_diff > worst_max_diff: + worst_max_diff = max_diff + worst_a = group[i].reshape(-1) + worst_b = group[j].reshape(-1) if worst_a is not None: _print_topk_vec(worst_a.cpu(), worst_b.cpu(), args.topk, args.sort_err, - f"build_kv_input_v_L{lyr}_token0") + f"build_kv_input_v_L{last_layer}_token0") _print_summary(all_results) if args.output: _dump_json(all_results, args.output, "—", args.dir_multi, - tag=f"within_batch_N{total_copies}", dir_off2=None) + tag=f"within_batch_N{total_sequences}", dir_off2=None) if __name__ == "__main__": From ca37506956174d287e721997215bd5f8667a25cf Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 30 Jun 2026 17:50:22 +0800 Subject: [PATCH 54/61] [refactor] standardize all PS-diag dump block markers to ##### [PS-diag] ... ##### MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - attention.py (patch ON): added missing END marker - megatron_runtime.py: unified ######### → ##### [PS-diag] format - Megatron attention.py (invasive): added ##### wrappers to both blocks - All 12 dump blocks now consistently wrapped for easy identification Co-Authored-By: Claude Fable 5 --- .../megatron/core/transformer/attention.py | 9 ++++----- .../integrations/megatron_runtime.py | 16 ++++++++-------- .../verl080_mcore0161_ms0160/attention.py | 1 + 3 files changed, 13 insertions(+), 13 deletions(-) diff --git a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py index 100cbfba..2e6db54a 100644 --- a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py +++ b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py @@ -1069,8 +1069,7 @@ def forward( value = value.squeeze(1) nvtx_range_pop(suffix="adjust_key_value") - # [PS-diag] OFF pre-RoPE Q/K/V dump — 侵入式，post-squeeze、pre-RoPE， - # 与 ON 侧（attention.py ON 分支 get_qkv+squeeze 之后 dump）完全对称。 + # ##### [PS-diag] OFF pre-RoPE Q/K/V dump（侵入式，post-squeeze、pre-RoPE）##### import os as _ps_os if _ps_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: try: @@ -1085,6 +1084,7 @@ def forward( self.config.num_layers) except Exception as _ps_e: print(f"[PS-diag] OFF preqkv dump failed: {_ps_e}", flush=True) + # ##### [PS-diag] OFF pre-RoPE Q/K/V dump end ##### # ================================================ # relative positional embedding (rotary embedding) @@ -1143,9 +1143,7 @@ def forward( # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb) nvtx_range_pop(suffix="rotary_pos_emb") - # [PS-diag] OFF post-RoPE Q/K + full_kv dump — 侵入式，rotary block 之后、core attention 之前， - # 与 ON 侧（_apply_positioned_rope 返回后 dump rope_postqk；build_kv 后 dump expanded_kv）对称。 - # 此处 query/key 是 post-RoPE，value 是 raw（未旋转），都在 scope。 + # ##### [PS-diag] OFF post-RoPE Q/K + full_kv dump（侵入式，rotary block 之后）##### if _ps_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: try: from prefix_sharing.tools.diagnostic_dump_verl080 import ( @@ -1157,6 +1155,7 @@ def forward( self.config.num_layers) except Exception as _ps_e2: print(f"[PS-diag] OFF postqk/full_kv dump failed: {_ps_e2}", flush=True) + # ##### [PS-diag] OFF post-RoPE Q/K + full_kv dump end ##### # ================================== # core attention computation diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index 90f88019..ac51e4d0 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -106,14 +106,14 @@ def prefix_attention( f"built expanded kv: expanded_key_shape={tuple(expanded_key.shape)}, expanded_value_shape={tuple(expanded_value.shape)}" ) - ######### prefix-sharing diag: ON expanded K/V dump（build_kv 输出，attention 实际用的完整 KV）######### + ##### [PS-diag] ON expanded K/V dump（build_kv 输出）##### try: from prefix_sharing.tools.diagnostic_dump_verl080 import dump_expanded_kv_on dump_expanded_kv_on(layer_id, expanded_key, expanded_value, attention_module.config.num_layers) except Exception as _e: print(f"expanded_kv dump failed: {_e}", flush=True) - ######### prefix-sharing diag: ON expanded K/V dump end ######### + ##### [PS-diag] ON expanded K/V dump end ##### # 注意力计算 core_attn_out = attention_backend.attention( @@ -128,7 +128,7 @@ def prefix_attention( core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) output = attention_module.linear_proj(core_attn_out) # (tensor, bias) tuple - ######### prefix-sharing diag: ON attention_output (per-layer) ######### + ##### [PS-diag] ON attention_output (per-layer) ##### try: from prefix_sharing.tools.diagnostic_dump import dump_attn_on dump_attn_on(output[0], packed_seq_params, prefix_sharing_context.prefix_sharing_plan, @@ -136,7 +136,7 @@ def prefix_attention( attention_module.config.num_layers) except Exception as e: print(f"last-attn dump (ON) failed: {e}") - ######### prefix-sharing diag: ON attention_output (per-layer) ######### + ##### [PS-diag] ON attention_output (per-layer) end ##### # --- return output @@ -315,14 +315,14 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: if q_pos_emb is not None: q_freqs = q_pos_emb.index_select(0, positions) - ######### prefix-sharing diag: ON rope_freqs (per-layer) ######### + ##### [PS-diag] ON rope_freqs (per-layer) ##### try: from prefix_sharing.tools.diagnostic_dump import dump_rope_freqs dump_rope_freqs(q_freqs, attention_module.layer_number, attention_module.config.num_layers) except Exception as e: print(f"rope_freqs dump failed: {e}") - ######### prefix-sharing diag: ON rope_freqs (per-layer) ######### + ##### [PS-diag] ON rope_freqs (per-layer) end ##### query = apply_rotary_pos_emb( query.unsqueeze(1), q_freqs, @@ -336,7 +336,7 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: **_rope_kwargs(cu_seqlens_kv), ).squeeze(1) - ######### prefix-sharing diag: ON post-RoPE Q/K dump (per-layer) ######### + ##### [PS-diag] ON post-RoPE Q/K dump (per-layer) ##### try: from prefix_sharing.tools.diagnostic_dump_verl080 import dump_rope_postqk_verl080 dump_rope_postqk_verl080( @@ -347,7 +347,7 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: ) except Exception as e: print(f"rope_postqk_layer dump failed: {e}") - ######### prefix-sharing diag: ON post-RoPE Q/K dump end ######### + ##### [PS-diag] ON post-RoPE Q/K dump end ##### return query, key diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index d0d06865..6cebf20a 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -137,6 +137,7 @@ def patched_forward( self.config.num_layers) except Exception as _e: print(f"hidden_states dump failed: {_e}", flush=True) + # ##### [PS-diag] ON pre-RoPE Q/K/V dump end ##### # delegate to verified integrations code from prefix_sharing.integrations.megatron_runtime import ( From 14f6fb510365563a7fbaa6256018c448e5111a26 Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 1 Jul 2026 10:48:01 +0800 Subject: [PATCH 55/61] [feat] multi-rank (TP/PP) precision verification: dump routing + assemble tool + variable readability refactor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## New features - dump 侧 PP-aware routing: 替换 _rank0_only() 为 scope 模型 (global/tp_vocab/pp_last/pp_stage) - 新增 _stage_last_layer() 解决 PP 下 buffer flush 永不触发的问题 - 新增 assemble_dump.py: 将多卡分散 dump 文件合并为单卡兼容格式 - 修复 forward_step.py PP 崩溃: restore_via_2d_unfold_verl080 加 isinstance guard ## Refactored - 所有 dump 代码段变量名可读性优化 (_lp→log_probs_nested, _ol→original_lengths 等) - cmp_diag.py / cmp_diag_verl080.py 变量名规范化 (lo/lf→logits_on/logits_off, lyr→layer_idx 等) - exception 变量统一为 exc ## Files: 10 files (9 modified + 1 new) --- .../megatron/core/transformer/attention.py | 8 +- .../integrations/megatron_runtime.py | 16 +- .../verl080_mcore0161_ms0160/attention.py | 55 ++-- .../verl080_mcore0161_ms0160/forward_step.py | 119 +++++---- .../vocab_logprobs.py | 26 +- .../prefix_sharing/tools/assemble_dump.py | 196 ++++++++++++++ .../prefix_sharing/tools/cmp_diag.py | 124 ++++----- .../prefix_sharing/tools/cmp_diag_verl080.py | 244 +++++++++--------- .../prefix_sharing/tools/diagnostic_dump.py | 139 ++++++++-- .../tools/diagnostic_dump_verl080.py | 56 ++-- 10 files changed, 653 insertions(+), 330 deletions(-) create mode 100644 prefix-sharing/prefix_sharing/tools/assemble_dump.py diff --git a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py index 2e6db54a..a10ff8e3 100644 --- a/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py +++ b/dependency/Megatron-LM-core_v0.16.1/megatron/core/transformer/attention.py @@ -1082,8 +1082,8 @@ def forward( self.config.num_layers) dump_hidden_states_on(self.layer_number, hidden_states, self.config.num_layers) - except Exception as _ps_e: - print(f"[PS-diag] OFF preqkv dump failed: {_ps_e}", flush=True) + except Exception as exc: + print(f"[PS-diag] OFF preqkv dump failed: {exc}", flush=True) # ##### [PS-diag] OFF pre-RoPE Q/K/V dump end ##### # ================================================ @@ -1153,8 +1153,8 @@ def forward( self.config.num_layers) dump_full_kv_off(self.layer_number, key, value, self.config.num_layers) - except Exception as _ps_e2: - print(f"[PS-diag] OFF postqk/full_kv dump failed: {_ps_e2}", flush=True) + except Exception as exc: + print(f"[PS-diag] OFF postqk/full_kv dump failed: {exc}", flush=True) # ##### [PS-diag] OFF post-RoPE Q/K + full_kv dump end ##### # ================================== diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py index ac51e4d0..bbbf06ca 100644 --- a/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py +++ b/prefix-sharing/prefix_sharing/integrations/megatron_runtime.py @@ -111,8 +111,8 @@ def prefix_attention( from prefix_sharing.tools.diagnostic_dump_verl080 import dump_expanded_kv_on dump_expanded_kv_on(layer_id, expanded_key, expanded_value, attention_module.config.num_layers) - except Exception as _e: - print(f"expanded_kv dump failed: {_e}", flush=True) + except Exception as exc: + print(f"expanded_kv dump failed: {exc}", flush=True) ##### [PS-diag] ON expanded K/V dump end ##### # 注意力计算 @@ -134,8 +134,8 @@ def prefix_attention( dump_attn_on(output[0], packed_seq_params, prefix_sharing_context.prefix_sharing_plan, attention_module.layer_number, attention_module.config.num_layers) - except Exception as e: - print(f"last-attn dump (ON) failed: {e}") + except Exception as exc: + print(f"last-attn dump (ON) failed: {exc}") ##### [PS-diag] ON attention_output (per-layer) end ##### # --- @@ -320,8 +320,8 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: from prefix_sharing.tools.diagnostic_dump import dump_rope_freqs dump_rope_freqs(q_freqs, attention_module.layer_number, attention_module.config.num_layers) - except Exception as e: - print(f"rope_freqs dump failed: {e}") + except Exception as exc: + print(f"rope_freqs dump failed: {exc}") ##### [PS-diag] ON rope_freqs (per-layer) end ##### query = apply_rotary_pos_emb( query.unsqueeze(1), @@ -345,8 +345,8 @@ def _rope_kwargs(_unused_cu_seqlens: Any | None) -> dict[str, Any]: attention_module.config.num_layers, positions=packed_position_ids, ) - except Exception as e: - print(f"rope_postqk_layer dump failed: {e}") + except Exception as exc: + print(f"rope_postqk_layer dump failed: {exc}") ##### [PS-diag] ON post-RoPE Q/K dump end ##### return query, key diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py index 6cebf20a..2fb7fbb0 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/attention.py @@ -38,9 +38,9 @@ def patched_forward( # ── normal path: 调用原始 forward ── # post-RoPE Q/K / full_kv / preqk 由 Megatron attention.py 侵入式 dump 写入； # patch 层只负责 attn_outputs + rope_freqs（侵入式未覆盖的）。 - import os as _os - _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None - _result = original_forward( + import os as _diag_os + diag_enabled = _diag_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None + forward_result = original_forward( self, hidden_states, attention_mask, @@ -56,46 +56,47 @@ def patched_forward( inference_params=inference_params, ) # ##### [PS-diag] OFF attn_outputs + rope_freqs dump ##### - if _diag_on: + if diag_enabled: import torch # 仅用于构造 positions（freqs 切 per-token + Q/K debug） from prefix_sharing.tools.diagnostic_dump import ( dump_attn_off, dump_rope_freqs, ) # rope_postqk / preqk / full_kv 已由 Megatron 侵入式 dump 覆盖 from prefix_sharing.integrations.megatron_runtime import _unpack_rotary_pos_emb - _attn_out = _result[0] if isinstance(_result, tuple) else _result - _bs = ( + attn_output = forward_result[0] if isinstance(forward_result, tuple) else forward_result + batch_size = ( len(packed_seq_params.cu_seqlens_q_padded) - 1 if (packed_seq_params is not None and hasattr(packed_seq_params, "cu_seqlens_q_padded")) else 0 ) - dump_attn_off(_attn_out, packed_seq_params, - self.layer_number, _bs, self.config.num_layers) + dump_attn_off(attn_output, packed_seq_params, + self.layer_number, batch_size, self.config.num_layers) if rotary_pos_emb is not None: - _q_pos_emb, _k_pos_emb = _unpack_rotary_pos_emb(rotary_pos_emb) + q_pos_emb, k_pos_emb = _unpack_rotary_pos_emb(rotary_pos_emb) # OFF 标准 positions（每 segment 内 0..seg-1）：切 per-token freqs + Q/K debug - _off_positions = None + per_token_positions = None if (packed_seq_params is not None and hasattr(packed_seq_params, "cu_seqlens_q_padded")): - _cu = packed_seq_params.cu_seqlens_q_padded - # device 必须显式到 _cu.device(GPU)：arange 默认 CPU，否则后面 - # _q_pos_emb.index_select(0, _off_positions) 会 device 不匹配崩 forward。 - _off_positions = torch.cat([ - torch.arange(int(_cu[i + 1] - _cu[i]), device=_cu.device) - for i in range(len(_cu) - 1) + cu_seqlens_tensor = packed_seq_params.cu_seqlens_q_padded + # device 必须显式到 cu_seqlens_tensor.device(GPU)：arange 默认 CPU，否则后面 + # q_pos_emb.index_select(0, per_token_positions) 会 device 不匹配崩 forward。 + per_token_positions = torch.cat([ + torch.arange(int(cu_seqlens_tensor[i + 1] - cu_seqlens_tensor[i]), + device=cu_seqlens_tensor.device) + for i in range(len(cu_seqlens_tensor) - 1) ]).long() # rope_freqs：存 per-token 角度（与 ON 同款），统一 rope_freqs.pt - if _off_positions is not None: + if per_token_positions is not None: dump_rope_freqs( - _q_pos_emb.index_select(0, _off_positions), + q_pos_emb.index_select(0, per_token_positions), self.layer_number, self.config.num_layers, ) # rope_postqk + full_kv 已由 megatron attention.py 侵入式 dump # （rotary block 之后），不在 patch 层重复——避免 hook 二次 flush # 覆盖侵入式已写好的完整 24 层文件。 # ##### [PS-diag] OFF attn_outputs + rope_freqs dump end ##### - return _result + return forward_result # ── prefix-sharing path ── # phase 1: training, THD, no fusion, no output gate @@ -115,28 +116,28 @@ def patched_forward( # ##### [PS-diag] ON pre-RoPE Q/K/V 统一 dump（get_qkv 之后、RoPE 之前）##### # 全部在此点 dump（squeeze 后、_apply_positioned_rope / build_kv 之前）， # 与 OFF baseline（hook 在 get_qkv 输出处截）同口径，集中对比，避免分散。 - import os as _os - if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + import os as _diag_os + if _diag_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump_verl080 import ( dump_rope_preqk_verl080, dump_build_kv_input_v_on, ) try: dump_rope_preqk_verl080(self.layer_number, query, key, self.config.num_layers) - except Exception as _e: - print(f"rope_preqk (pre-RoPE Q/K) dump failed: {_e}", flush=True) + except Exception as exc: + print(f"rope_preqk (pre-RoPE Q/K) dump failed: {exc}", flush=True) try: dump_build_kv_input_v_on(self.layer_number, value, self.config.num_layers) - except Exception as _e: - print(f"build_kv_input_v (pre-RoPE V) dump failed: {_e}", flush=True) + except Exception as exc: + print(f"build_kv_input_v (pre-RoPE V) dump failed: {exc}", flush=True) # [PS-diag] dump hidden_states for input-level comparison try: from prefix_sharing.tools.diagnostic_dump_verl080 import dump_hidden_states_on dump_hidden_states_on(self.layer_number, hidden_states, self.config.num_layers) - except Exception as _e: - print(f"hidden_states dump failed: {_e}", flush=True) + except Exception as exc: + print(f"hidden_states dump failed: {exc}", flush=True) # ##### [PS-diag] ON pre-RoPE Q/K/V dump end ##### # delegate to verified integrations code diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py index f6243a7c..863d847f 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py @@ -122,8 +122,8 @@ def patched_forward_step( # ON: prefix_lens / original_lengths 取自 plan； # OFF: prefix_lens 全0、original_lengths 从 input_ids NestedTensor offsets diff 推。 # cu_seqlens 取送进 forward 的 input_ids NestedTensor offsets（ON=裁剪后 packed 边界, OFF=完整）。 - import os as _os - if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + import os as _diag_os + if _diag_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump_verl080 import ( dump_meta_verl080, dump_attention_mask_verl080, dump_label_mask_verl080, @@ -131,48 +131,49 @@ def patched_forward_step( nested_offsets_to_cu, ) from prefix_sharing.integrations.verl_mcore import _is_nested_tensor - _ids_nested = batch_for_forward["input_ids"] - if _is_nested_tensor(_ids_nested): + input_ids_nested = batch_for_forward["input_ids"] + if _is_nested_tensor(input_ids_nested): if ps_state is not None: - _plan = ps_state.prefix_sharing_plan - _prefix_lens = list(_plan.prefix_lens) - _orig_lens = list(_plan.original_lengths) + sharing_plan = ps_state.prefix_sharing_plan + prefix_lens_list = list(sharing_plan.prefix_lens) + original_lengths = list(sharing_plan.original_lengths) else: - _diffs = _ids_nested.offsets().diff().tolist() - _orig_lens = [int(d) for d in _diffs] - _prefix_lens = [0] * len(_orig_lens) - dump_meta_verl080(_prefix_lens, nested_offsets_to_cu(_ids_nested)) + seq_lengths = input_ids_nested.offsets().diff().tolist() + original_lengths = [int(length) for length in seq_lengths] + prefix_lens_list = [0] * len(original_lengths) + dump_meta_verl080(prefix_lens_list, + nested_offsets_to_cu(input_ids_nested)) # attention_mask + label_mask：两种 log_probs 对比范围，都不含越界预测位 # （POS L_i-1，其 logp 预测不存在的 token[L_i]）。对齐 restore 后 log_probs # 的 [B, L_max] 紧凑坐标系。 # attention_mask：[0:L_i-1) prompt 区+prompt-last+response 区（整体 restore 验证） # label_mask：[prompt-last:L_i-1) prompt-last+response 区（PPO loss 范围） - _Lmax_lm = max(_orig_lens) if _orig_lens else 0 + max_seq_len = max(original_lengths) if original_lengths else 0 # tag 与 logprobs 一致（接入点2 用 model.training 区分 old/train）， # 保证 mask 和 logprobs_{tag} 来自同一 forward（同 batch、同 L_max）。 - _tag_lm = "train" if model.training else "old" - # attention_mask 仅依赖 _orig_lens，不需要 loss_mask。 + diag_tag = "train" if model.training else "old" + # attention_mask 仅依赖 original_lengths，不需要 loss_mask。 dump_attention_mask_verl080( - build_attention_mask_2d(_orig_lens, _Lmax_lm), _tag_lm) - # label_mask 用 response_lens（每行 response token 数）。verl080 padding 后 + build_attention_mask_2d(original_lengths, max_seq_len), diag_tag) + # label_mask 用 response_lengths（每行 response token 数）。verl080 padding 后 # loss_mask = response_mask 是 2D left-right padded（非 NestedTensor，见 # verl padding.py:71），不能走 nested_to_2d_full；但 response token 数 = # loss_mask 行 sum，与坐标系无关，据此推 prompt_len 最稳（2D/NestedTensor 均适用）。 - _lm = original_batch.get("loss_mask") - if _lm is not None: - if _is_nested_tensor(_lm): - _lm_off = _lm.offsets() - _lm_val = _lm.values() - _response_lens = [ - int(_lm_val[_lm_off[i]:_lm_off[i + 1]].sum()) - for i in range(len(_orig_lens))] + loss_mask_tensor = original_batch.get("loss_mask") + if loss_mask_tensor is not None: + if _is_nested_tensor(loss_mask_tensor): + loss_offsets = loss_mask_tensor.offsets() + loss_values = loss_mask_tensor.values() + response_lengths = [ + int(loss_values[loss_offsets[i]:loss_offsets[i + 1]].sum()) + for i in range(len(original_lengths))] else: # .long() 免 import torch（本文件顶部未导入 torch）； # .cpu() 防御 on-device tensor 的 tolist() - _response_lens = _lm.sum(dim=-1).long().cpu().tolist() + response_lengths = loss_mask_tensor.sum(dim=-1).long().cpu().tolist() dump_label_mask_verl080( - build_label_mask_2d(_response_lens, _orig_lens, _Lmax_lm), - _tag_lm) + build_label_mask_2d(response_lengths, original_lengths, max_seq_len), + diag_tag) # ##### [PS-diag] dump 元数据 + masks end ##### # ── 构造修改后的 iterator 喂回原始 forward_step ── @@ -201,6 +202,9 @@ def patched_forward_step( # forward_step 返回 (output_dict, partial(postprocess_func))， # 解包处理 output_dict 再重包。restore_via_2d_unfold_verl080 内部 # 会检查 context / restore_indices，无 restore 需求时 early return。 + # + # PP guard: 非末 stage 返回的是 tensor（hidden_states），不是 dict； + # restore 只在末 stage（output 是含 log_probs 的 dict）才有意义。 if ps_state is not None: from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080 from prefix_sharing.integrations.context import current_prefix_sharing_context @@ -209,22 +213,25 @@ def patched_forward_step( vocab_parallel_log_probs_from_logits, ) output_dict, postprocess_fn = output - output_dict = restore_via_2d_unfold_verl080( - output_dict, - vocab_parallel_log_probs_from_logits, - vocab_parallel_entropy, - ) - # 释放 vocab 维 logits（占用大，只在 context 生命周期内持有， - # restore 已消费完毕）。clear 职责在此，不在包装函数内。 - ctx = current_prefix_sharing_context() - if ctx is not None: - ctx.prefix_last_logits_saved.clear() + if isinstance(output_dict, dict): + output_dict = restore_via_2d_unfold_verl080( + output_dict, + vocab_parallel_log_probs_from_logits, + vocab_parallel_entropy, + ) + # 释放 vocab 维 logits（占用大，只在 context 生命周期内持有， + # restore 已消费完毕）。clear 职责在此，不在包装函数内。 + ctx = current_prefix_sharing_context() + if ctx is not None: + ctx.prefix_last_logits_saved.clear() output = (output_dict, postprocess_fn) # ##### [PS-diag] dump 2D logprobs/entropy（ON=restore后, OFF=原始） ##### # restore 后（ON）或原始 forward（OFF）的 log_probs/entropy 都是 NestedTensor， # 每行长度 = original_lengths[i]，展开到统一 [B, L_max] 供 cmp_diag.cmp_2d 逐元素对比。 - import os as _os2 - if _os2.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: + # + # PP guard: 非末 stage 返回 tensor 而非 dict；跳过 get() 避免 AttributeError。 + import os as _diag_os + if _diag_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None: from prefix_sharing.tools.diagnostic_dump_verl080 import ( nested_to_2d_full, dump_logprobs_2d_verl080, dump_entropy_2d_verl080, ) @@ -234,19 +241,27 @@ def patched_forward_step( # （eval_mode→training=False→"old" 对应 old_logp 阶段； # train_mode→training=True→"train" 对应 update_actor 阶段）。 # 这样一次 run 自动产出 logprobs_old + logprobs_train 两份，不互相覆盖。 - _tag = "train" if model.training else "old" - _out_dict, _ = output - _lp = _out_dict.get("log_probs") - if _is_nested_tensor(_lp): - if ps_state is not None: - _ol = list(ps_state.prefix_sharing_plan.original_lengths) - else: - _ol = [int(d) for d in _lp.offsets().diff().tolist()] - _Lmax = max(_ol) if _ol else 0 - dump_logprobs_2d_verl080(nested_to_2d_full(_lp, _ol, _Lmax), _tag) - _ent = _out_dict.get("entropy") - if _is_nested_tensor(_ent): - dump_entropy_2d_verl080(nested_to_2d_full(_ent, _ol, _Lmax), _tag) + diag_tag = "train" if model.training else "old" + output_dict_raw, _ = output + if isinstance(output_dict_raw, dict): + log_probs_nested = output_dict_raw.get("log_probs") + if _is_nested_tensor(log_probs_nested): + if ps_state is not None: + original_lengths = list( + ps_state.prefix_sharing_plan.original_lengths) + else: + original_lengths = [ + int(length) for length + in log_probs_nested.offsets().diff().tolist()] + max_seq_len = max(original_lengths) if original_lengths else 0 + dump_logprobs_2d_verl080( + nested_to_2d_full(log_probs_nested, original_lengths, max_seq_len), + diag_tag, scope="pp_last") + entropy_nested = output_dict_raw.get("entropy") + if _is_nested_tensor(entropy_nested): + dump_entropy_2d_verl080( + nested_to_2d_full(entropy_nested, original_lengths, max_seq_len), + diag_tag, scope="pp_last") # ##### [PS-diag] dump 2D logprobs/entropy end ##### _ps_forward_step_probe("after_original_forward_step") return output diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py index 65d7b746..5ebcfef8 100644 --- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py +++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/vocab_logprobs.py @@ -26,9 +26,9 @@ def patched_fn(logits, labels): # ##### [PS-diag] dump logits（ON/OFF 都 dump，必须在 original_fn 之前） ##### # logits 形态 [N, V//tp]（或 [N,1,V//tp]），cmp_diag.cmp_logits_packed 会 # reshape 成 token-major [N,V] 再用 cu_seqlens+prefix_lens 对齐。 - import os as _os - _diag_on = _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None - if _diag_on: + import os as _diag_os + diag_enabled = _diag_os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None + if diag_enabled: from prefix_sharing.tools.diagnostic_dump_verl080 import dump_logits_verl080 dump_logits_verl080(logits) # ##### [PS-diag] dump logits end ##### @@ -50,21 +50,21 @@ def patched_fn(logits, labels): logits_2d = logits.view(-1, logits.size(-1)) # ##### [PS-diag] 验证 packed 坐标对齐（logits N 是 valid 还是 padded） ##### - if _diag_on: - _layout = ctx.packed_batch_layout + if diag_enabled: + packed_layout = ctx.packed_batch_layout print( f"[PS-diag][packed-align] logits_N={logits_2d.shape[0]} " - f"total_padded={_layout.total_padded_length} " - f"total_valid={_layout.total_valid_length} " - f"has_padding={_layout.has_padding}", + f"total_padded={packed_layout.total_padded_length} " + f"total_valid={packed_layout.total_valid_length} " + f"has_padding={packed_layout.has_padding}", flush=True, ) - for _idx in ctx.prefix_last_restore_indices: + for restore_index in ctx.prefix_last_restore_indices: print( - f"[PS-diag][packed-align] reuser={_idx.reuse_idx_in_batch} " - f"provider={_idx.provider_idx_in_batch} " - f"provider_1d_pos={_idx.provider_1d_pos} " - f"target_2d_pos={_idx.target_2d_pos}", + f"[PS-diag][packed-align] reuser={restore_index.reuse_idx_in_batch} " + f"provider={restore_index.provider_idx_in_batch} " + f"provider_1d_pos={restore_index.provider_1d_pos} " + f"target_2d_pos={restore_index.target_2d_pos}", flush=True, ) # ##### [PS-diag] 验证 packed 坐标对齐 end ##### diff --git a/prefix-sharing/prefix_sharing/tools/assemble_dump.py b/prefix-sharing/prefix_sharing/tools/assemble_dump.py new file mode 100644 index 00000000..d0a23c66 --- /dev/null +++ b/prefix-sharing/prefix_sharing/tools/assemble_dump.py @@ -0,0 +1,196 @@ +"""Multi-rank diagnostic dump assembler — merges PP/TP-sharded files into a +single-card-compatible flat directory. + +Reads a raw dump directory produced by the diagnostic dump infrastructure under +TP/PP parallelism (with ``_pp{p}`` / ``_tp{r}`` file suffixes) and assembles a +clean flat directory that looks exactly like a single-card dump. The output can +then be fed directly to the **unmodified** ``cmp_diag_verl080.py``. + +Usage:: + + python assemble_dump.py --input-dir /path/to/raw_dump --output-dir /path/to/assembled + +Single-card dumps (tp==1, pp==1) are a fast path: all files are copied verbatim. +""" + +from __future__ import annotations + +import argparse +import json +import os +import shutil +from typing import Any + +import torch + +# ── Per-layer dict files that may be PP-sharded ────────────────── +# These glob patterns match the stem (without _pp suffix or .pt extension). +_PP_STAGE_STEMS: list[str] = [ + "attn_outputs", + "rope_preqk", + "rope_postqk", + "rope_freqs", + "expanded_kv", + "full_kv", + "build_kv_input_v", + "hidden_states", +] + +# ── Global (non-sharded) files — copied verbatim ───────────────── +_GLOBAL_FILES: list[str] = [ + "parallel_info.json", + "cu_seqlens_q.pt", + "cu_seqlens_q_logits.pt", + "prefix_lens.pt", +] + +# ── Optional tag-suffixed files (glob-matched) ─────────────────── +_GLOB_TAG_FILES: list[str] = [ + "attention_mask_", + "label_mask_", + "logprobs_", + "entropy_", +] + + +def _strip_ext(fname: str) -> tuple[str, str]: + """Split ``stem.ext`` → ``(stem, ext)``. ``ext`` includes the dot.""" + idx = fname.rfind(".") + if idx == -1: + return fname, "" + return fname[:idx], fname[idx:] + + +def assemble(input_dir: str, output_dir: str) -> None: + """Assemble a multi-rank dump into a single-card-compatible directory. + + Reads ``parallel_info.json`` to determine topology, then: + - copies global metadata files verbatim + - merges per-PP-stage layer dicts into single dicts + - concatenates TP-sharded logits along the vocab dimension + - copies 2D files verbatim + """ + os.makedirs(output_dir, exist_ok=True) + + # ── Load topology ───────────────────────────────────────────── + manifest_path = os.path.join(input_dir, "parallel_info.json") + if not os.path.exists(manifest_path): + print(f"[assemble] ERROR: {manifest_path} not found — is this a diagnostic dump?") + return + + with open(manifest_path, encoding="utf-8") as f: + manifest = json.load(f) + + tp_size: int = manifest.get("tp_size", 1) + pp_size: int = manifest.get("pp_size", 1) + scopes: dict[str, str] = manifest.get("scopes", {}) + + is_multi_rank = tp_size > 1 or pp_size > 1 + if not is_multi_rank: + print("[assemble] tp=1 pp=1 — fast path: copying all files") + _copy_tree(input_dir, output_dir) + return + + print(f"[assemble] tp_size={tp_size} pp_size={pp_size}") + + # ── Copy global metadata files ───────────────────────────────── + for fname in _GLOBAL_FILES: + src = os.path.join(input_dir, fname) + if os.path.exists(src): + shutil.copy2(src, os.path.join(output_dir, fname)) + print(f" [copy] {fname}") + + # ── Copy glob-tagged files (attention_mask_old.pt, etc.) ──────── + for prefix in _GLOB_TAG_FILES: + if not os.path.isdir(input_dir): + continue + for fname in os.listdir(input_dir): + if fname.startswith(prefix) and fname.endswith(".pt"): + src = os.path.join(input_dir, fname) + shutil.copy2(src, os.path.join(output_dir, fname)) + print(f" [copy] {fname}") + + # ── Merge per-PP-stage layer dicts ───────────────────────────── + for stem in _PP_STAGE_STEMS: + fname = f"{stem}.pt" + # Determine if this file is PP-sharded from manifest scopes + scope = scopes.get(stem, "") + if scope != "pp_stage" or pp_size <= 1: + # Single file — copy verbatim + src = os.path.join(input_dir, fname) + if os.path.exists(src): + shutil.copy2(src, os.path.join(output_dir, fname)) + print(f" [copy] {fname}") + continue + + # PP-sharded: load and merge + merged: dict[int, Any] = {} + found_any = False + for p in range(pp_size): + stem_ext = f"{stem}_pp{p}.pt" + src = os.path.join(input_dir, stem_ext) + if os.path.exists(src): + d = torch.load(src, weights_only=True) + if isinstance(d, dict): + merged.update(d) + found_any = True + if found_any: + torch.save(merged, os.path.join(output_dir, fname)) + print(f" [merge] {fname} ← {pp_size} stage(s), {len(merged)} layers") + + # ── Concat TP-sharded logits ────────────────────────────────── + if scopes.get("logits", "") == "tp_vocab" and tp_size > 1: + shards = [] + for t in range(tp_size): + src = os.path.join(input_dir, f"logits_tp{t}.pt") + if os.path.exists(src): + shards.append(torch.load(src, weights_only=True)) + if shards: + full = torch.cat(shards, dim=-1) + torch.save(full, os.path.join(output_dir, "logits.pt")) + print(f" [concat] logits.pt ← {len(shards)} tp shards, shape {list(full.shape)}") + else: + # tp_size==1: copy logits.pt verbatim + src = os.path.join(input_dir, "logits.pt") + if os.path.exists(src): + shutil.copy2(src, os.path.join(output_dir, "logits.pt")) + print(" [copy] logits.pt") + + # ── Summary ─────────────────────────────────────────────────── + n_files = len(os.listdir(output_dir)) + print(f"\n[assemble] done — {n_files} files written to {output_dir}") + + +def _copy_tree(src_dir: str, dst_dir: str) -> None: + """Copy all .pt and .json files from src_dir to dst_dir (fast path for single-card).""" + if not os.path.isdir(src_dir): + return + for fname in os.listdir(src_dir): + if fname.endswith(".pt") or fname.endswith(".json"): + src = os.path.join(src_dir, fname) + if os.path.isfile(src): + shutil.copy2(src, os.path.join(dst_dir, fname)) + + +# ── CLI ─────────────────────────────────────────────────────────── + +def main() -> None: + ap = argparse.ArgumentParser( + description="Assemble multi-rank (TP/PP) diagnostic dump into single-card format", + ) + ap.add_argument( + "--input-dir", "-i", + required=True, + help="Raw multi-rank dump directory (with _pp{p}/_tp{r} suffixes)", + ) + ap.add_argument( + "--output-dir", "-o", + required=True, + help="Assembled single-card-compatible directory", + ) + args = ap.parse_args() + assemble(args.input_dir, args.output_dir) + + +if __name__ == "__main__": + main() diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag.py b/prefix-sharing/prefix_sharing/tools/cmp_diag.py index 77300bbd..46fce34a 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag.py @@ -201,7 +201,7 @@ def _first_token_metrics(a_vec: torch.Tensor, b_vec: torch.Tensor) -> dict: "cos": cos, "pearson": pr} -def _logits_first_token(lo: torch.Tensor, lf: torch.Tensor +def _logits_first_token(logits_on: torch.Tensor, logits_off: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor] | None: """Extract first token's full-vocab vector from ON/OFF logits. @@ -210,21 +210,21 @@ def _logits_first_token(lo: torch.Tensor, lf: torch.Tensor output_layer → [S, B, V//tp] → model returns [B, S, V//tp]). """ # Flatten all leading batch/token dims into N, keep V as last dim - lo_2d = lo.reshape(-1, lo.size(-1)) - lf_2d = lf.reshape(-1, lf.size(-1)) + on_flat = logits_on.reshape(-1, logits_on.size(-1)) + off_flat = logits_off.reshape(-1, logits_off.size(-1)) # First token = first row → full-vocab vector [V] - return lo_2d[0, :].contiguous(), lf_2d[0, :].contiguous() + return on_flat[0, :].contiguous(), off_flat[0, :].contiguous() -def _logits_ensure_token_major(lo: torch.Tensor, lf: torch.Tensor +def _logits_ensure_token_major(logits_on: torch.Tensor, logits_off: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """Ensure logits are 2D [N, V] for packed alignment. Vocab is always the last dim (verified from Megatron model forward). Batch dim on leading axes is flattened into N. """ - return (lo.reshape(-1, lo.size(-1)).contiguous(), - lf.reshape(-1, lf.size(-1)).contiguous()) + return (logits_on.reshape(-1, logits_on.size(-1)).contiguous(), + logits_off.reshape(-1, logits_off.size(-1)).contiguous()) # ══════════════════════════════════════════════════════════════════ @@ -421,9 +421,9 @@ def cmp_rope_postqk(dir_on: str, dir_off: str) -> CheckResult | None: first_row_len = 0 num_layers = len(la) - for lyr in sorted(la): - on_entry = a[lyr] - off_entry = b[lyr] + for layer_idx in sorted(la): + on_entry = a[layer_idx] + off_entry = b[layer_idx] on_q, on_k = on_entry["query"], on_entry["key"] off_q, off_k = off_entry["query"], off_entry["key"] on_pos = on_entry.get("positions") @@ -549,18 +549,18 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: max_diff = 0.0 mismatches: list[dict] = [] # [{layer, token_idx, dim, on_val, off_val, diff}] - for lyr in sorted(la): - on_freqs = on_dict[lyr] # [T_on, 1, 1, D] + for layer_idx in sorted(la): + on_freqs = on_dict[layer_idx] # [T_on, 1, 1, D] # Reconstruct OFF per-token for this layer off_freqs = torch.cat( - [off_dict[lyr][:s, :, :, :] for s in seqlens], dim=0) # [T_off, 1, 1, D] + [off_dict[layer_idx][:s, :, :, :] for s in seqlens], dim=0) # [T_off, 1, 1, D] try: on_aligned, off_aligned = _align_packed( on_freqs, off_freqs, align_mask) except ValueError as e: return CheckResult(name="rope_freqs", passed=False, - metrics={"error": f"align failed L{lyr}: {e}"}) + metrics={"error": f"align failed L{layer_idx}: {e}"}) diff = (on_aligned - off_aligned).abs() # [N, 1, 1, D] md = float(diff.max()) @@ -574,7 +574,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: t = int(t) d = int(token_diff.indices[t]) mismatches.append({ - "layer": lyr, + "layer": layer_idx, "token_idx": t, "dim": d, "on_val": float(on_aligned[t, 0, 0, d]), @@ -602,7 +602,7 @@ def _per_layer_cos(dir_on: str, dir_off: str, layer: int | None) -> dict | None: extract the matching suffix region from OFF before comparison. """ - def _cos_for_layer(a, b, lyr, need_align, align_mask): + def _cos_for_layer(a, b, layer_idx, need_align, align_mask): if a.dim() == 3: a, b = a.squeeze(1), b.squeeze(1) if need_align and a.shape[0] != b.shape[0]: @@ -659,8 +659,8 @@ def _cos_for_layer(a, b, lyr, need_align, align_mask): T = int(mb["cu_seqlens"][-1]) if mb and mb["cu_seqlens"].numel() > 0 else 0 if T > 0: # Check if any layer has shape mismatch - for lyr in da: - if lyr in db and da[lyr].shape != db[lyr].shape: + for layer_idx in da: + if layer_idx in db and da[layer_idx].shape != db[layer_idx].shape: need_align = True break if need_align: @@ -675,8 +675,8 @@ def _cos_for_layer(a, b, lyr, need_align, align_mask): mb["cu_seqlens"], ma["prefix_lens"], T) results = {} - for lyr in sorted(set(da.keys()) & set(db.keys())): - results[lyr] = _cos_for_layer(da[lyr], db[lyr], lyr, need_align, align_mask) + for layer_idx in sorted(set(da.keys()) & set(db.keys())): + results[layer_idx] = _cos_for_layer(da[layer_idx], db[layer_idx], layer_idx, need_align, align_mask) return results @@ -686,8 +686,8 @@ def cmp_attn_layer(dir_on: str, dir_off: str, if r is None: return None if "cos_avg" in r: - lyr = r["layer"] - return CheckResult(name=f"attn_L{lyr}", + layer_idx = r["layer"] + return CheckResult(name=f"attn_L{layer_idx}", passed=r["cos_avg"] > 0.9999 and r["cos_min"] > 0.999, metrics=r) return CheckResult(name="attn_per_layer", passed=True, @@ -712,22 +712,22 @@ def cmp_first_token(dir_on: str, dir_off: str) -> list[CheckResult]: a = _load_attn_output(dir_on, last) b = _load_attn_output(dir_off, last) if a is not None and b is not None: - a0 = a.squeeze(1) if a.dim() == 3 else a - b0 = b.squeeze(1) if b.dim() == 3 else b - ft = _first_token_metrics(a0[0], b0[0]) - results.append(CheckResult(name="first_token_attn", metrics=ft)) + on_token0 = a.squeeze(1) if a.dim() == 3 else a + off_token0 = b.squeeze(1) if b.dim() == 3 else b + metrics = _first_token_metrics(on_token0[0], off_token0[0]) + results.append(CheckResult(name="first_token_attn", metrics=metrics)) # logits — packed[0], auto-detect [N,V] vs [V,N] format - lo = _load_tensor(dir_on, "logits.pt") - lf = _load_tensor(dir_off, "logits.pt") - if lo is not None and lf is not None: - lo_first, lf_first = _logits_first_token(lo, lf) - if lo_first is not None: - ft = _first_token_metrics(lo_first, lf_first) - results.append(CheckResult(name="first_token_logits", metrics=ft)) + logits_on = _load_tensor(dir_on, "logits.pt") + logits_off = _load_tensor(dir_off, "logits.pt") + if logits_on is not None and logits_off is not None: + logits_on_first, logits_off_first = _logits_first_token(logits_on, logits_off) + if logits_on_first is not None: + metrics = _first_token_metrics(logits_on_first, logits_off_first) + results.append(CheckResult(name="first_token_logits", metrics=metrics)) else: _log.warning("first_token_logits skipped: cannot determine token dim " - "(ON %s, OFF %s)", _fmt_shape(lo.shape), _fmt_shape(lf.shape)) + "(ON %s, OFF %s)", _fmt_shape(logits_on.shape), _fmt_shape(logits_off.shape)) return results @@ -743,22 +743,22 @@ def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: with OFF (full-sequence) logits, same alignment logic as attn_output per-layer comparison. """ - lo = _load_tensor(dir_on, "logits.pt") - lf = _load_tensor(dir_off, "logits.pt") - if lo is None or lf is None: + logits_on = _load_tensor(dir_on, "logits.pt") + logits_off = _load_tensor(dir_off, "logits.pt") + if logits_on is None or logits_off is None: return None # Logits may be [N,V] or [V,N] — ensure token-major [N,V] for alignment - lo, lf = _logits_ensure_token_major(lo, lf) + logits_on, logits_off = _logits_ensure_token_major(logits_on, logits_off) # Metadata for logits uses cu_seqlens_q_logits.pt - ma = _load_packed_meta(dir_on, "cu_seqlens_q_logits.pt") - mb = _load_packed_meta(dir_off, "cu_seqlens_q_logits.pt") - if ma is None or mb is None: + meta_on = _load_packed_meta(dir_on, "cu_seqlens_q_logits.pt") + meta_off = _load_packed_meta(dir_off, "cu_seqlens_q_logits.pt") + if meta_on is None or meta_off is None: return None - T_off = int(mb["cu_seqlens"][-1]) if mb["cu_seqlens"].numel() > 0 else 0 - if T_off == 0 or lo.shape[0] == 0 or lf.shape[0] == 0: + total_off_tokens = int(meta_off["cu_seqlens"][-1]) if meta_off["cu_seqlens"].numel() > 0 else 0 + if total_off_tokens == 0 or logits_on.shape[0] == 0 or logits_off.shape[0] == 0: return None # Build alignment mask (same logic as attn_output all-layers) @@ -766,19 +766,19 @@ def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: mask_off_2d = _load_attention_mask_2d(dir_off) if mask_on_2d is not None and mask_off_2d is not None: align_mask = _build_alignment_mask_from_2d( - mask_on_2d, mask_off_2d, mb["cu_seqlens"], T_off) + mask_on_2d, mask_off_2d, meta_off["cu_seqlens"], total_off_tokens) else: align_mask = _build_alignment_mask( - mb["cu_seqlens"], ma["prefix_lens"], T_off) + meta_off["cu_seqlens"], meta_on["prefix_lens"], total_off_tokens) # Align: extract suffix-only region from OFF try: - on_aligned, off_aligned = _align_packed(lo, lf, align_mask) + on_aligned, off_aligned = _align_packed(logits_on, logits_off, align_mask) except ValueError as e: return CheckResult(name="logits", passed=False, metrics={"error": str(e), - "n_on": lo.shape[0], - "n_off": lf.shape[0]}) + "n_on": logits_on.shape[0], + "n_off": logits_off.shape[0]}) n_tokens = on_aligned.shape[0] @@ -905,13 +905,13 @@ def _print_per_layer(r: CheckResult): print(f" {'LAYER':>6s} {'COS_AVG':>14s} {'COS_MIN':>14s} {'TOKENS':>8s} {'STATUS':>8s}") print(f" {'─'*6} {'─'*14} {'─'*14} {'─'*8} {'─'*8}") bad = [] - for lyr in sorted(layers.keys()): - d = layers[lyr] + for layer_idx in sorted(layers.keys()): + d = layers[layer_idx] ok = d["cos_avg"] > 0.9999 and d["cos_min"] > 0.999 - print(f" {lyr:>6d} {d['cos_avg']:>14.6e} {d['cos_min']:>14.6e} " + print(f" {layer_idx:>6d} {d['cos_avg']:>14.6e} {d['cos_min']:>14.6e} " f"{d['n_tokens']:>8d} {'PASS' if ok else 'WARN':>8s}") if not ok: - bad.append(lyr) + bad.append(layer_idx) if bad: print(f"\n ⚠ First deviating layer: {bad[0]}") elif "cos_avg" in r.metrics: @@ -1139,20 +1139,20 @@ def main(): a = _load_attn_output(args.dir_on, last) b = _load_attn_output(args.dir_off, last) if a is not None and b is not None: - a0 = (a.squeeze(1) if a.dim() == 3 else a)[0].cpu() - b0 = (b.squeeze(1) if b.dim() == 3 else b)[0].cpu() - _print_topk_vec(a0, b0, args.topk, "val", + on_token0 = (a.squeeze(1) if a.dim() == 3 else a)[0].cpu() + off_token0 = (b.squeeze(1) if b.dim() == 3 else b)[0].cpu() + _print_topk_vec(on_token0, off_token0, args.topk, "val", "first_token_attn") # first_token_logits — per-dim top-K (real = ON signed value, # so positive logits — the tokens actually selectable by # sampling — surface first, not the large-magnitude negatives) - lo = _load_tensor(args.dir_on, "logits.pt") - lf = _load_tensor(args.dir_off, "logits.pt") - if lo is not None and lf is not None: - ft = _logits_first_token(lo, lf) - if ft is not None: - _print_topk_vec(ft[0].cpu(), ft[1].cpu(), args.topk, - "real", "first_token_logits") + logits_on = _load_tensor(args.dir_on, "logits.pt") + logits_off = _load_tensor(args.dir_off, "logits.pt") + if logits_on is not None and logits_off is not None: + first_token_pair = _logits_first_token(logits_on, logits_off) + if first_token_pair is not None: + _print_topk_vec(first_token_pair[0].cpu(), first_token_pair[1].cpu(), + args.topk, "real", "first_token_logits") # ── ④b Logits (full packed alignment via attention_mask + dual pointer) ── if not stop: diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index acbff5fa..b3452e47 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -306,11 +306,11 @@ def _aligned_vec_at_pos( return on[pos].contiguous(), off[pos].contiguous() -def _logits_ensure_token_major(lo: torch.Tensor, lf: torch.Tensor +def _logits_ensure_token_major(logits_on: torch.Tensor, logits_off: torch.Tensor ) -> tuple[torch.Tensor, torch.Tensor]: """确保 logits 为 2D [N, V]（token-major），vocab 在最后一维。""" - return (lo.reshape(-1, lo.size(-1)).contiguous(), - lf.reshape(-1, lf.size(-1)).contiguous()) + return (logits_on.reshape(-1, logits_on.size(-1)).contiguous(), + logits_off.reshape(-1, logits_off.size(-1)).contiguous()) # ════════════════════════════════════════════════════════════════ @@ -376,13 +376,13 @@ def cmp_attn_layer(dir_on: str, dir_off: str, return None results = {} - for lyr in sorted(set(da.keys()) & set(db.keys())): - a, b = da[lyr], db[lyr] + for layer_idx in sorted(set(da.keys()) & set(db.keys())): + a, b = da[layer_idx], db[layer_idx] need = align_mask is not None and a.shape[0] != b.shape[0] try: - results[lyr] = _cos_for_layer(a, b, align_mask if need else None) + results[layer_idx] = _cos_for_layer(a, b, align_mask if need else None) except ValueError as e: - results[lyr] = {"error": str(e)} + results[layer_idx] = {"error": str(e)} return CheckResult(name="attn_per_layer", passed=True, metrics={"layers": results}) @@ -421,9 +421,9 @@ def cmp_packed_token(dir_on: str, dir_off: str, name=f"attn_L{attn_layer}_pos{pos}", metrics=_vec_metrics(vecs[0], vecs[1]))) - lo = _load_logits(dir_on) - lf = _load_logits(dir_off) - vecs = _aligned_vec_at_pos(lo, lf, False, pos, align_mask) + logits_on = _load_logits(dir_on) + logits_off = _load_logits(dir_off) + vecs = _aligned_vec_at_pos(logits_on, logits_off, False, pos, align_mask) if vecs is None: results.append(CheckResult( name=f"logits_pos{pos}", @@ -523,18 +523,18 @@ def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, return None results = {} - for lyr in sorted(set(da.keys()) & set(db.keys())): - ea, eb = da[lyr], db[lyr] + for layer_idx in sorted(set(da.keys()) & set(db.keys())): + ea, eb = da[layer_idx], db[layer_idx] qa, ka = ea.get("query"), ea.get("key") qb, kb = eb.get("query"), eb.get("key") if qa is None or qb is None: continue need = align_mask is not None and qa.shape[0] != qb.shape[0] try: - results[lyr] = _rope_postqk_cos_for_layer(qa, ka, qb, kb, + results[layer_idx] = _rope_postqk_cos_for_layer(qa, ka, qb, kb, align_mask if need else None) except ValueError as e: - results[lyr] = {"error": str(e)} + results[layer_idx] = {"error": str(e)} return CheckResult(name=f"{label}_per_layer", passed=True, metrics={"layers": results}) @@ -636,7 +636,7 @@ def cmp_rope_postqk_token(dir_on: str, dir_off: str, """Q/K packed[pos] 对比（**suffix 对齐后**），单 stage。 stage="pre" → rope_preqk.pt（旋转前），stage="post" → rope_postqk.pt（旋转后）。 - 对 Q、K 分别输出 {label}_L{lyr}_Q_pos{pos} / {label}_L{lyr}_K_pos{pos}。 + 对 Q、K 分别输出 {label}_L{layer_idx}_Q_pos{pos} / {label}_L{layer_idx}_K_pos{pos}。 调用方按 pre → rope_freqs → post 顺序分别调用。 """ if stage == "pre": @@ -648,27 +648,27 @@ def cmp_rope_postqk_token(dir_on: str, dir_off: str, def cmp_logits_packed(dir_on: str, dir_off: str) -> CheckResult | None: """全 packed logits suffix 对齐 + per-token cosine。""" - lo = _load_logits(dir_on) - lf = _load_logits(dir_off) - if lo is None or lf is None: + logits_on = _load_logits(dir_on) + logits_off = _load_logits(dir_off) + if logits_on is None or logits_off is None: return None - lo, lf = _logits_ensure_token_major(lo, lf) + logits_on, logits_off = _logits_ensure_token_major(logits_on, logits_off) - ma = _load_packed_meta(dir_on, "cu_seqlens_q_logits.pt") - mb = _load_packed_meta(dir_off, "cu_seqlens_q_logits.pt") - if ma is None or mb is None: + meta_on = _load_packed_meta(dir_on, "cu_seqlens_q_logits.pt") + meta_off = _load_packed_meta(dir_off, "cu_seqlens_q_logits.pt") + if meta_on is None or meta_off is None: return None - T_off = int(mb["cu_seqlens"][-1]) if mb["cu_seqlens"].numel() > 0 else 0 - if T_off == 0 or lo.shape[0] == 0 or lf.shape[0] == 0: + total_off_tokens = int(meta_off["cu_seqlens"][-1]) if meta_off["cu_seqlens"].numel() > 0 else 0 + if total_off_tokens == 0 or logits_on.shape[0] == 0 or logits_off.shape[0] == 0: return None - align_mask = _build_alignment_mask(mb["cu_seqlens"], ma["prefix_lens"], T_off) + align_mask = _build_alignment_mask(meta_off["cu_seqlens"], meta_on["prefix_lens"], total_off_tokens) try: - on_aligned, off_aligned = _align_packed(lo, lf, align_mask) + on_aligned, off_aligned = _align_packed(logits_on, logits_off, align_mask) except ValueError as e: return CheckResult(name="logits", passed=False, metrics={"error": str(e), - "n_on": lo.shape[0], "n_off": lf.shape[0]}) + "n_on": logits_on.shape[0], "n_off": logits_off.shape[0]}) cos = _cosine_sim(on_aligned, off_aligned, dim=-1) cos_avg, cos_min = float(cos.mean()), float(cos.min()) @@ -714,8 +714,8 @@ def _load_rope_freqs_vec_at_pos(dir_on: str, dir_off: str, layer: int, pos: int, align_mask = _build_attn_align_mask(dir_on, dir_off) if align_mask is None: return None, None - _aligned = _align_rope_freqs_layer(on_dict[layer], off_dict[layer], align_mask) - if _aligned is None: + aligned_result = _align_rope_freqs_layer(on_dict[layer], off_dict[layer], align_mask) + if aligned_result is None: return None, None on_a, off_a = _aligned if pos < 0 or pos >= on_a.shape[0]: @@ -743,21 +743,21 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) if layer is not None: layers = [l for l in layers if l == layer] - _name = f"rope_freqs_L{layer}" if layer is not None else "rope_freqs" + result_name = f"rope_freqs_L{layer}" if layer is not None else "rope_freqs" if not layers: - return CheckResult(name=_name, passed=False, + return CheckResult(name=result_name, passed=False, metrics={"error": f"layer {layer} 不在双方 rope_freqs 中"}) align_mask = _build_attn_align_mask(dir_on, dir_off) if align_mask is None: - return CheckResult(name=_name, passed=False, + return CheckResult(name=result_name, passed=False, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) max_diff = 0.0 mismatches: list[dict] = [] - for lyr in layers: - _aligned = _align_rope_freqs_layer(on_dict[lyr], off_dict[lyr], align_mask) - if _aligned is None: + for layer_idx in layers: + aligned_result = _align_rope_freqs_layer(on_dict[layer_idx], off_dict[layer_idx], align_mask) + if aligned_result is None: continue on_a, off_a = _aligned diff = (on_a - off_a).abs() # [N,1,1,D] @@ -770,7 +770,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, t = int(t) d = int(token_diff.indices[t]) mismatches.append({ - "layer": lyr, "token_idx": t, "dim": d, + "layer": layer_idx, "token_idx": t, "dim": d, "on_val": float(on_a[t, 0, 0, d]), "off_val": float(off_a[t, 0, 0, d]), "diff": float(token_diff.values[t]), @@ -780,7 +780,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, if mismatches: metrics["mismatches"] = mismatches[:20] metrics["total_mismatches"] = len(mismatches) - return CheckResult(name=_name, passed=max_diff == 0.0, metrics=metrics) + return CheckResult(name=result_name, passed=max_diff == 0.0, metrics=metrics) def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, @@ -801,27 +801,27 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, return None common = set(on_dict.keys()) & set(off_dict.keys()) rf_layer = layer if layer is not None else (max(common) if common else 0) - _name = f"rope_freqs_L{rf_layer}_pos{pos}" + result_name = f"rope_freqs_L{rf_layer}_pos{pos}" if rf_layer not in on_dict or rf_layer not in off_dict: - return CheckResult(name=_name, metrics={"error": f"layer {rf_layer} 缺失"}) + return CheckResult(name=result_name, metrics={"error": f"layer {rf_layer} 缺失"}) if align_mask is None: align_mask = _build_attn_align_mask(dir_on, dir_off) if align_mask is None: - return CheckResult(name=_name, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) + return CheckResult(name=result_name, metrics={"error": "cu_seqlens/prefix_lens 缺失"}) - _aligned = _align_rope_freqs_layer(on_dict[rf_layer], off_dict[rf_layer], align_mask) - if _aligned is None: - return CheckResult(name=_name, metrics={"error": "对齐失败"}) + aligned_result = _align_rope_freqs_layer(on_dict[rf_layer], off_dict[rf_layer], align_mask) + if aligned_result is None: + return CheckResult(name=result_name, metrics={"error": "对齐失败"}) on_a, off_a = _aligned n = on_a.shape[0] if pos < 0 or pos >= n: - return CheckResult(name=_name, + return CheckResult(name=result_name, metrics={"error": f"pos {pos} 越界: 对齐后 token 数={n}"}) on_vec = on_a[pos].reshape(-1) off_vec = off_a[pos].reshape(-1) m = _vec_metrics(on_vec, off_vec) - return CheckResult(name=_name, passed=m["max_abs"] == 0.0, metrics=m) + return CheckResult(name=result_name, passed=m["max_abs"] == 0.0, metrics=m) # ════════════════════════════════════════════════════════════════ @@ -851,49 +851,53 @@ def cmp_attn_kv(dir_on: str, dir_off: str, 完整 KV）。相同 → attention 输入一致，attention_output 差异必来自 attention 计算/mask； 不同 → bug 在 build_kv 的 prefix 复用（store/expand）。 """ - fa = os.path.join(dir_on, "expanded_kv.pt") - fb = os.path.join(dir_off, "full_kv.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "expanded_kv.pt") + filepath_off = os.path.join(dir_off, "full_kv.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) if layer is not None: layers = [l for l in layers if l == layer] - _name = f"attn_kv_L{layer}" if layer is not None else "attn_kv" + resultresult_name = f"attn_kv_L{layer}" if layer is not None else "attn_kv" if not layers: - return CheckResult(name=_name, passed=False, + return CheckResult(name=result_name, passed=False, metrics={"error": f"layer {layer} 不在双方 attn_kv 中"}) per_layer: dict = {} worst = {"max_diff": 0.0, "cos_min": 1.0} - for lyr in layers: - ek, ev = on_dict[lyr].get("key"), on_dict[lyr].get("value") - fk, fv = off_dict[lyr].get("key"), off_dict[lyr].get("value") - d: dict = {} - for _tag, (_a, _b) in [("K", (ek, fk)), ("V", (ev, fv))]: - if _a is None or _b is None: - d[_tag] = {"error": "缺失"} + for layer_idx in layers: + on_key = on_dict[layer_idx].get("key") + on_value = on_dict[layer_idx].get("value") + off_key = off_dict[layer_idx].get("key") + off_value = off_dict[layer_idx].get("value") + entry_result: dict = {} + for kv_type, (on_kv, off_kv) in [("K", (on_key, off_key)), ("V", (on_value, off_value))]: + if on_kv is None or off_kv is None: + entry_result[kv_type] = {"error": "缺失"} continue - if _a.shape != _b.shape: - d[_tag] = {"error": f"shape mismatch ON{tuple(_a.shape)} vs OFF{tuple(_b.shape)}"} + if on_kv.shape != off_kv.shape: + entry_result[kv_type] = { + "error": f"shape mismatch ON{tuple(on_kv.shape)} vs OFF{tuple(off_kv.shape)}"} continue - _af = _a.reshape(_a.shape[0], -1).float() - _bf = _b.reshape(_b.shape[0], -1).float() - _diff = (_af - _bf).abs() - _cos = _cosine_sim(_af, _bf, dim=-1) - _md = float(_diff.max()) - d[_tag] = {"max_diff": _md, - "cos_avg": float(_cos.mean()), "cos_min": float(_cos.min()), - "n_tokens": _af.shape[0]} - worst["max_diff"] = max(worst["max_diff"], _md) - worst["cos_min"] = min(worst["cos_min"], float(_cos.min())) - per_layer[lyr] = d + on_flat = on_kv.reshape(on_kv.shape[0], -1).float() + off_flat = off_kv.reshape(off_kv.shape[0], -1).float() + element_diff = (on_flat - off_flat).abs() + token_cos = _cosine_sim(on_flat, off_flat, dim=-1) + max_elem_diff = float(element_diff.max()) + entry_result[kv_type] = { + "max_diff": max_elem_diff, + "cos_avg": float(token_cos.mean()), "cos_min": float(token_cos.min()), + "n_tokens": on_flat.shape[0]} + worst["max_diff"] = max(worst["max_diff"], max_elem_diff) + worst["cos_min"] = min(worst["cos_min"], float(token_cos.min())) + per_layer[layer_idx] = entry_result # expanded 应精确等于 full → 阈值极严 passed = worst["max_diff"] < 1e-5 and worst["cos_min"] > 0.9999 - return CheckResult(name=_name, passed=passed, + return CheckResult(name=result_name, passed=passed, metrics={"layers": per_layer, "max_diff": worst["max_diff"], "cos_min": worst["cos_min"], "num_layers": len(layers)}) @@ -909,18 +913,18 @@ def _print_attn_kv(r: CheckResult): f"{'V_MAXDIFF':>12s} {'V_COS':>10s} {'STATUS':>8s}") print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 12} {'─' * 10} {'─' * 8}") bad = [] - for lyr in sorted(layers): - d = layers[lyr] + for layer_idx in sorted(layers): + d = layers[layer_idx] kd, vd = d.get("K", {}), d.get("V", {}) if "error" in kd or "error" in vd: - print(f" {lyr:>6d} K:{kd.get('error','')} V:{vd.get('error','')}") - bad.append(lyr); continue + print(f" {layer_idx:>6d} K:{kd.get('error','')} V:{vd.get('error','')}") + bad.append(layer_idx); continue kmd, kcos = kd["max_diff"], kd["cos_avg"] vmd, vcos = vd["max_diff"], vd["cos_avg"] ok = kmd < 1e-5 and vmd < 1e-5 if not ok: - bad.append(lyr) - print(f" {lyr:>6d} {kmd:>12.3e} {kcos:>10.6f} " + bad.append(layer_idx) + print(f" {layer_idx:>6d} {kmd:>12.3e} {kcos:>10.6f} " f"{vmd:>12.3e} {vcos:>10.6f} {'OK' if ok else 'DIFF':>8s}") print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " f"{_CHECK if r.passed else _CROSS} " @@ -949,19 +953,19 @@ def cmp_build_kv_input_v(dir_on: str, dir_off: str, layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) if layer is not None: layers = [l for l in layers if l == layer] - _name = f"build_kv_input_v_L{layer}" if layer is not None else "build_kv_input_v" + result_name = f"build_kv_input_v_L{layer}" if layer is not None else "build_kv_input_v" if not layers: - return CheckResult(name=_name, passed=False, metrics={"error": "no layers"}) + return CheckResult(name=result_name, passed=False, metrics={"error": "no layers"}) align_mask = _build_attn_align_mask(dir_on, dir_off) per_layer: dict = {} worst_md = 0.0 worst_cos = 1.0 - for lyr in layers: - on_v = on_dict[lyr] - off_v = off_dict[lyr] + for layer_idx in layers: + on_v = on_dict[layer_idx] + off_v = off_dict[layer_idx] if on_v is None or off_v is None: - per_layer[lyr] = {"error": "缺失"}; continue + per_layer[layer_idx] = {"error": "缺失"}; continue on_f = on_v.reshape(on_v.shape[0], -1).float() off_f = off_v.reshape(off_v.shape[0], -1).float() on_T, off_T = int(on_v.shape[0]), int(off_v.shape[0]) @@ -969,18 +973,18 @@ def cmp_build_kv_input_v(dir_on: str, dir_off: str, try: on_f, off_f = _align_packed(on_f, off_f, align_mask) except ValueError as e: - per_layer[lyr] = {"error": str(e), "on_T": on_T, "off_T": off_T} + per_layer[layer_idx] = {"error": str(e), "on_T": on_T, "off_T": off_T} continue diff = (on_f - off_f).abs() cos = _cosine_sim(on_f, off_f, dim=-1) md = float(diff.max()) - per_layer[lyr] = {"max_diff": md, "cos_avg": float(cos.mean()), + per_layer[layer_idx] = {"max_diff": md, "cos_avg": float(cos.mean()), "cos_min": float(cos.min()), "n_tokens": on_f.shape[0], "on_T": on_T, "off_T": off_T} worst_md = max(worst_md, md) worst_cos = min(worst_cos, float(cos.min())) passed = worst_md < 1e-5 - return CheckResult(name=_name, passed=passed, + return CheckResult(name=result_name, passed=passed, metrics={"layers": per_layer, "max_diff": worst_md, "cos_min": worst_cos}) @@ -994,15 +998,15 @@ def _print_build_kv_input_v(r: CheckResult): print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS':>10s} " f"{'ON_T':>8s} {'OFF_T':>8s} {'STATUS':>8s}") print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") - for lyr in sorted(layers): - d = layers[lyr] + for layer_idx in sorted(layers): + d = layers[layer_idx] if "max_diff" not in d: - print(f" {lyr:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") + print(f" {layer_idx:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") continue md, cos = d["max_diff"], d["cos_avg"] ok = md < 1e-5 _crop = " (cropped)" if d.get("on_T") != d.get("off_T") else "" - print(f" {lyr:>6d} {md:>12.3e} {cos:>10.6f} " + print(f" {layer_idx:>6d} {md:>12.3e} {cos:>10.6f} " f"{d.get('on_T', '—'):>8} {d.get('off_T', '—'):>8} " f"{'OK' if ok else 'DIFF':>8s}{_crop}") print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " @@ -1029,19 +1033,19 @@ def cmp_hidden_states(dir_on: str, dir_off: str, layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) if layer is not None: layers = [l for l in layers if l == layer] - _name = f"hidden_states_L{layer}" if layer is not None else "hidden_states" + result_name = f"hidden_states_L{layer}" if layer is not None else "hidden_states" if not layers: - return CheckResult(name=_name, passed=False, metrics={"error": "no layers"}) + return CheckResult(name=result_name, passed=False, metrics={"error": "no layers"}) align_mask = _build_attn_align_mask(dir_on, dir_off) per_layer: dict = {} worst_md = 0.0 worst_cos = 1.0 - for lyr in layers: - on_hs = on_dict[lyr] - off_hs = off_dict[lyr] + for layer_idx in layers: + on_hs = on_dict[layer_idx] + off_hs = off_dict[layer_idx] if on_hs is None or off_hs is None: - per_layer[lyr] = {"error": "缺失"}; continue + per_layer[layer_idx] = {"error": "缺失"}; continue on_f = on_hs.reshape(on_hs.shape[0], -1).float() off_f = off_hs.reshape(off_hs.shape[0], -1).float() on_T, off_T = int(on_hs.shape[0]), int(off_hs.shape[0]) @@ -1049,18 +1053,18 @@ def cmp_hidden_states(dir_on: str, dir_off: str, try: on_f, off_f = _align_packed(on_f, off_f, align_mask) except ValueError as e: - per_layer[lyr] = {"error": str(e), "on_T": on_T, "off_T": off_T} + per_layer[layer_idx] = {"error": str(e), "on_T": on_T, "off_T": off_T} continue diff = (on_f - off_f).abs() cos = _cosine_sim(on_f, off_f, dim=-1) md = float(diff.max()) - per_layer[lyr] = {"max_diff": md, "cos_avg": float(cos.mean()), + per_layer[layer_idx] = {"max_diff": md, "cos_avg": float(cos.mean()), "cos_min": float(cos.min()), "n_tokens": on_f.shape[0], "on_T": on_T, "off_T": off_T} worst_md = max(worst_md, md) worst_cos = min(worst_cos, float(cos.min())) passed = worst_md < 1e-5 - return CheckResult(name=_name, passed=passed, + return CheckResult(name=result_name, passed=passed, metrics={"layers": per_layer, "max_diff": worst_md, "cos_min": worst_cos}) @@ -1074,14 +1078,14 @@ def _print_hidden_states(r: CheckResult): print(f" {'LAYER':>6s} {'MAXDIFF':>12s} {'COS':>10s} " f"{'ON_T':>8s} {'OFF_T':>8s} {'STATUS':>8s}") print(f" {'─' * 6} {'─' * 12} {'─' * 10} {'─' * 8} {'─' * 8} {'─' * 8}") - for lyr in sorted(layers): - d = layers[lyr] + for layer_idx in sorted(layers): + d = layers[layer_idx] if "max_diff" not in d: - print(f" {lyr:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") + print(f" {layer_idx:>6d} {d.get('error', '')} ON_T={d.get('on_T')} OFF_T={d.get('off_T')}") continue md, cos = d["max_diff"], d["cos_avg"] ok = md < 1e-5 - print(f" {lyr:>6d} {md:>12.3e} {cos:>10.6f} " + print(f" {layer_idx:>6d} {md:>12.3e} {cos:>10.6f} " f"{d.get('on_T', '—'):>8} {d.get('off_T', '—'):>8} " f"{'OK' if ok else 'DIFF':>8s}") print(f"\n max_diff={m.get('max_diff')} cos_min={m.get('cos_min')} " @@ -1317,17 +1321,17 @@ def _print_per_layer(r: CheckResult): f"{'TOKENS':>8s} {'STATUS':>8s}") print(f" {'─' * 6} {'─' * 14} {'─' * 14} {'─' * 8} {'─' * 8}") bad = [] - for lyr in sorted(layers.keys()): - d = layers[lyr] + for layer_idx in sorted(layers.keys()): + d = layers[layer_idx] if "error" in d: - print(f" {lyr:>6d} {d['error']}") - bad.append(lyr) + print(f" {layer_idx:>6d} {d['error']}") + bad.append(layer_idx) continue ok = d["cos_avg"] > _COS_AVG_PASS and d["cos_min"] > _COS_MIN_PASS - print(f" {lyr:>6d} {d['cos_avg']:>14.6e} {d['cos_min']:>14.6e} " + print(f" {layer_idx:>6d} {d['cos_avg']:>14.6e} {d['cos_min']:>14.6e} " f"{d['n_tokens']:>8d} {'PASS' if ok else 'WARN':>8s}") if not ok: - bad.append(lyr) + bad.append(layer_idx) if bad: print(f"\n ⚠ First deviating layer: {bad[0]}") elif "cos_avg" in r.metrics: @@ -1353,20 +1357,20 @@ def _print_rope_postqk_per_layer(r: CheckResult): print(f" {'─' * 6} {'─' * 12} {'─' * 12} {'─' * 12} {'─' * 12} " f"{'─' * 8} {'─' * 8}") bad = [] - for lyr in sorted(layers.keys()): - d = layers[lyr] + for layer_idx in sorted(layers.keys()): + d = layers[layer_idx] if "error" in d: - print(f" {lyr:>6d} {d['error']}") - bad.append(lyr) + print(f" {layer_idx:>6d} {d['error']}") + bad.append(layer_idx) continue ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS and d["K_cos_avg"] > _COS_AVG_PASS and d["K_cos_min"] > _COS_MIN_PASS) - print(f" {lyr:>6d} {d.get('Q_max_diff', 0.0):>12.3e} " + print(f" {layer_idx:>6d} {d.get('Q_max_diff', 0.0):>12.3e} " f"{d['Q_cos_avg']:>12.6e} " f"{d.get('K_max_diff', 0.0):>12.3e} {d['K_cos_avg']:>12.6e} " f"{d['n_tokens']:>8d} {'PASS' if ok else 'WARN':>8s}") if not ok: - bad.append(lyr) + bad.append(layer_idx) if bad: print(f"\n ⚠ First deviating layer: {bad[0]}") print(f" （Q/K max_diff 与 build_kv_input_v 的 V max_diff 同口径，可直接对比）") @@ -1659,9 +1663,9 @@ def main(): if vecs is not None: _print_topk_vec(vecs[0].cpu(), vecs[1].cpu(), args.topk, "val", f"attn_L{attn_layer}_pos{pos}") - lo = _load_logits(args.dir_on) - lf = _load_logits(args.dir_off) - vecs = _aligned_vec_at_pos(lo, lf, False, pos, align_mask) + logits_on = _load_logits(args.dir_on) + logits_off = _load_logits(args.dir_off) + vecs = _aligned_vec_at_pos(logits_on, logits_off, False, pos, align_mask) if vecs is not None: _print_topk_vec(vecs[0].cpu(), vecs[1].cpu(), args.topk, "val", f"logits_pos{pos}", show_rel=False) diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py index 3af88d19..5f57e04a 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump.py @@ -86,7 +86,15 @@ def _rank0_only() -> bool: # cmp_diag knows how to gather each tensor without guessing. _TENSOR_SCOPES: dict[str, str] = { - "logits": "tp_vocab", # packed logits [N, V//tp] — vocab-sharded under TP + "logits": "tp_vocab", # packed logits [N, V//tp] — vocab-sharded under TP + "attn_outputs": "pp_stage", # per-layer attn output dict — PP-sharded + "rope_postqk": "pp_stage", # per-layer post-RoPE Q/K dict — PP-sharded + "rope_preqk": "pp_stage", # per-layer pre-RoPE Q/K dict — PP-sharded + "rope_freqs": "pp_stage", # per-layer RoPE freqs dict — PP-sharded + "expanded_kv": "pp_stage", # per-layer expanded KV dict — PP-sharded (ON only) + "full_kv": "pp_stage", # per-layer full KV dict — PP-sharded (OFF only) + "build_kv_input_v": "pp_stage", # per-layer V dict — PP-sharded + "hidden_states": "pp_stage", # per-layer hidden_states dict — PP-sharded } _PARALLEL_INFO_CACHE: Any = None @@ -112,6 +120,63 @@ def _cached_parallel_info() -> Any: return _PARALLEL_INFO_CACHE +def _stage_last_layer(num_layers_global: int) -> int: + """Return the last global layer number owned by the current PP stage. + + Under PP, ``num_layers_global`` is split across stages. + Non-last stages never reach ``layer_number == num_layers_global``, + so the standard auto-flush condition silently drops their data. + This function computes the stage-local last layer so every stage + flushes independently. + + Uniform split (standard Megatron): + layers_per_stage = num_layers // pp_size + remainder stages get 1 extra layer. + """ + pi = _cached_parallel_info() + if pi is None or pi.pp_size <= 1: + return num_layers_global + base = num_layers_global // pi.pp_size + rem = num_layers_global % pi.pp_size + if pi.pp_rank < rem: + return (pi.pp_rank + 1) * (base + 1) + else: + return rem * (base + 1) + (pi.pp_rank - rem + 1) * base + + +def _pp_suffix() -> str: + """Return ``'_pp{r}'`` when ``pp_size > 1``, else ``''``.""" + pi = _cached_parallel_info() + if pi is not None and pi.pp_size > 1: + return f"_pp{pi.pp_rank}" + return "" + + +def _should_write_for_scope(scope: str) -> bool: + """Return True if this rank should dump data for the given scope. + + Gate logic: + - ``"global"`` → rank 0 only (all ranks have identical data) + - ``"tp_vocab"`` → every tp rank dumps (each has different shard) + - ``"pp_last"`` → tp_rank==0 on pp_last stage only + - ``"pp_stage"`` → tp_rank==0 within each PP stage + """ + pi = _cached_parallel_info() + if scope == "global": + return _rank0_only() + if scope == "tp_vocab": + return True + if scope == "pp_last": + if pi is not None and not pi.is_pipeline_last_stage: + return False + return pi is None or pi.tp_rank == 0 + if scope == "pp_stage": + if pi is None or pi.pp_size <= 1: + return _rank0_only() + return pi.tp_rank == 0 + return _rank0_only() + + def _with_suffix(name: str, suffix: str) -> str: """Insert a rank suffix before the extension: logits.pt → logits_tp0.pt.""" if not suffix: @@ -130,7 +195,13 @@ def _ensure_manifest(dump_dir: str, pi: Any) -> None: if dump_dir in _MANIFEST_WRITTEN: return _MANIFEST_WRITTEN.add(dump_dir) - if not _rank0_only(): + # Under PP, each stage is a separate process; allow tp_rank==0 on any + # PP stage to write the manifest (content is identical across stages). + # This prevents data loss if pp0's dump path is never reached. + if pi is not None and pi.pp_size > 1: + if pi.tp_rank != 0: + return + elif not _rank0_only(): return import json manifest = { @@ -158,21 +229,28 @@ def _save_tensor(name: str, tensor: torch.Tensor, dump_dir: str, See the ``_TENSOR_SCOPES`` block above for the scope semantics. ``scope`` defaults to ``"global"`` (rank-0-only, plain filename) so existing callers - are unchanged; per-rank tensors opt in via ``scope="tp_vocab"`` (etc.). + are unchanged; per-rank tensors opt in via ``scope="tp_vocab"`` / ``"pp_last"`` + (etc.). + + Gate logic (delegated to ``_should_write_for_scope``): + - ``"global"`` → rank 0 only + - ``"tp_vocab"`` → every tp rank dumps + - ``"pp_last"`` → tp_rank==0 on pp_last only """ pi = _cached_parallel_info() _ensure_manifest(dump_dir, pi) + if not _should_write_for_scope(scope): + return False + + # Determine filename suffix from scope if scope == "tp_vocab" and pi is not None and pi.tp_size > 1: # Every tp rank dumps its own vocab shard; no rank-0 gate, no comm. fname = _with_suffix(name, f"_tp{pi.tp_rank}") else: - # global scope, OR tp_vocab with tp_size==1 (single shard == full): - # rank-0 dumps once under the plain name (single-card compatible). + # global / pp_last / tp_vocab with tp_size==1: plain filename. if scope == "tp_vocab": scope = "global" # tp==1 → behaves as global for logging - if not _rank0_only(): - return False fname = name try: path = os.path.join(dump_dir, fname) @@ -231,18 +309,23 @@ def _add_to_attn_buffer(layer_number: int, tensor: torch.Tensor) -> None: def _flush_attn_buffer(dump_dir: str) -> None: - """Write accumulated attn_outputs dict to disk and clear buffer.""" + """Write accumulated attn_outputs dict to disk and clear buffer. + + Under PP, each stage's tp_rank==0 writes with ``_pp{r}`` suffix; + assembly later merges all stage files. + """ global _ATTN_BUFFER if _ATTN_BUFFER is None: return - if not _rank0_only(): + if not _should_write_for_scope("pp_stage"): _ATTN_BUFFER = None return try: # Move every tensor to CPU (dict has no .detach()/.cpu()/.clone()) moved = {k: v.detach().cpu().clone() for k, v in _ATTN_BUFFER.items()} - torch.save(moved, os.path.join(dump_dir, "attn_outputs.pt")) - _log.warning("attn_outputs.pt saved (%d layers)", len(moved)) + fname = f"attn_outputs{_pp_suffix()}.pt" + torch.save(moved, os.path.join(dump_dir, fname)) + _log.warning("%s saved (%d layers)", fname, len(moved)) _ATTN_BUFFER = None except Exception as e: _log.warning("attn_outputs.pt save failed: %s", e) @@ -263,16 +346,21 @@ def _add_to_rope_buffer(layer_number: int, rotated_query: torch.Tensor, def _flush_rope_buffer(dump_dir: str) -> None: - """Write accumulated rope_postqk dict to disk and clear buffer.""" + """Write accumulated rope_postqk dict to disk and clear buffer. + + Under PP, each stage's tp_rank==0 writes with ``_pp{r}`` suffix; + assembly later merges all stage files. + """ global _ROPE_BUFFER if _ROPE_BUFFER is None: return - if not _rank0_only(): + if not _should_write_for_scope("pp_stage"): _ROPE_BUFFER = None return try: - torch.save(_ROPE_BUFFER, os.path.join(dump_dir, "rope_postqk.pt")) - _log.warning("rope_postqk.pt saved (%d layers)", len(_ROPE_BUFFER)) + fname = f"rope_postqk{_pp_suffix()}.pt" + torch.save(_ROPE_BUFFER, os.path.join(dump_dir, fname)) + _log.warning("%s saved (%d layers)", fname, len(_ROPE_BUFFER)) _ROPE_BUFFER = None except Exception as e: _log.warning("rope_postqk.pt save failed: %s", e) @@ -296,13 +384,14 @@ def dump_rope_freqs(q_freqs: torch.Tensor, layer_number: int, if _ROPE_FREQS_BUFFER is None: _ROPE_FREQS_BUFFER = {} _ROPE_FREQS_BUFFER[layer_number] = q_freqs.detach().cpu().clone() - if layer_number == num_layers: - if _rank0_only(): + if layer_number == _stage_last_layer(num_layers): + if _should_write_for_scope("pp_stage"): try: + fname = f"rope_freqs{_pp_suffix()}.pt" torch.save(_ROPE_FREQS_BUFFER, - os.path.join(dump_dir, "rope_freqs.pt")) - _log.warning("rope_freqs.pt saved (%d layers)", - len(_ROPE_FREQS_BUFFER)) + os.path.join(dump_dir, fname)) + _log.warning("%s saved (%d layers)", + fname, len(_ROPE_FREQS_BUFFER)) except Exception as e: _log.warning("rope_freqs.pt save failed: %s", e) _ROPE_FREQS_BUFFER = None @@ -335,7 +424,7 @@ def dump_rope_postqk_layer(layer_number: int, rotated_query: torch.Tensor, if dump_dir is None: return _add_to_rope_buffer(layer_number, rotated_query, rotated_key, positions) - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_rope_buffer(dump_dir) @@ -365,8 +454,8 @@ def dump_attn_on( _save_meta(packed_seq_params, list(prefix_sharing_plan.prefix_lens), dump_dir, meta_key="attn") - # Flush on last layer - if layer_number == num_layers: + # Flush on last layer of this PP stage + if layer_number == _stage_last_layer(num_layers): _flush_attn_buffer(dump_dir) @@ -395,8 +484,8 @@ def dump_attn_off( if layer_number == 1: _save_meta(packed_seq_params, [0] * batch_size, dump_dir, meta_key="attn") - # Flush on last layer - if layer_number == num_layers: + # Flush on last layer of this PP stage + if layer_number == _stage_last_layer(num_layers): _flush_attn_buffer(dump_dir) diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index d2eff422..6a5f8104 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -31,6 +31,9 @@ from prefix_sharing.tools.diagnostic_dump import ( _get_dump_dir, _save_tensor, + _stage_last_layer, + _pp_suffix, + _cached_parallel_info, ) @@ -159,12 +162,16 @@ def dump_logits_verl080(logits: torch.Tensor) -> None: _save_tensor("logits.pt", logits, dump_dir, scope="tp_vocab") -def dump_logprobs_2d_verl080(logp_2d: torch.Tensor, tag: str) -> None: - """存 2D log_probs ``[B, L_max]``，文件名 ``logprobs_{tag}.pt``。""" +def dump_logprobs_2d_verl080(logp_2d: torch.Tensor, tag: str, + scope: str = "global") -> None: + """存 2D log_probs ``[B, L_max]``，文件名 ``logprobs_{tag}.pt``。 + + ``scope="pp_last"`` 时仅最后一个 PP stage 落盘（多卡下 logprobs 只在末 stage 产生）。 + """ dump_dir = _get_dump_dir() if dump_dir is None: return - _save_tensor(f"logprobs_{tag}.pt", logp_2d, dump_dir) + _save_tensor(f"logprobs_{tag}.pt", logp_2d, dump_dir, scope=scope) def dump_attention_mask_verl080(mask_2d: torch.Tensor, tag: str) -> None: @@ -198,12 +205,16 @@ def dump_label_mask_verl080(mask_2d: torch.Tensor, tag: str) -> None: _save_tensor(f"label_mask_{tag}.pt", mask_2d.to(torch.bool), dump_dir) -def dump_entropy_2d_verl080(ent_2d: torch.Tensor | None, tag: str) -> None: - """存 2D entropy ``[B, L_max]``，文件名 ``entropy_{tag}.pt``。``None`` 时跳过。""" +def dump_entropy_2d_verl080(ent_2d: torch.Tensor | None, tag: str, + scope: str = "global") -> None: + """存 2D entropy ``[B, L_max]``，文件名 ``entropy_{tag}.pt``。``None`` 时跳过。 + + ``scope="pp_last"`` 时仅最后一个 PP stage 落盘（多卡下 entropy 只在末 stage 产生）。 + """ dump_dir = _get_dump_dir() if dump_dir is None or ent_2d is None: return - _save_tensor(f"entropy_{tag}.pt", ent_2d, dump_dir) + _save_tensor(f"entropy_{tag}.pt", ent_2d, dump_dir, scope=scope) # ════════════════════════════════════════════════════════════════ @@ -238,25 +249,32 @@ def dump_rope_postqk_verl080(layer_number: int, if positions is not None: entry["positions"] = positions.detach().cpu().clone() _ROPE_POSTQK_BUFFER[layer_number] = entry - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("rope_postqk.pt", _ROPE_POSTQK_BUFFER, dump_dir) _ROPE_POSTQK_BUFFER = None def _flush_dict_buffer(fname: str, buffer: dict, dump_dir: str) -> None: - """rank0 直接 torch.save 一个 dict buffer。 + """rank0 直接 torch.save 一个 dict buffer。PP-aware gating + suffix。 不能用 _save_tensor：它对入参做 .detach().cpu().clone()，dict 没 .detach() → AttributeError 被其 except 吞掉，文件永不写盘（rope_postqk.pt 曾因此丢失）。 entries 应在插入时已 detach().cpu().clone()。仿 _flush_attn_buffer。 + + Under PP, each stage's tp_rank==0 writes with ``_pp{r}`` suffix; + assembly later merges all stage files. """ import os as _os - from prefix_sharing.tools.diagnostic_dump import _rank0_only - if _rank0_only(): - try: - torch.save(buffer, _os.path.join(dump_dir, fname)) - except Exception as _e: - print(f"[PS-diag] {fname} save failed: {_e}", flush=True) + from prefix_sharing.tools.diagnostic_dump import _should_write_for_scope + if not _should_write_for_scope("pp_stage"): + return + try: + stem, sep, ext = fname.rpartition(".") + pp_sfx = _pp_suffix() + fname_pp = f"{stem}{pp_sfx}{sep}{ext}" if sep else f"{fname}{pp_sfx}" + torch.save(buffer, _os.path.join(dump_dir, fname_pp)) + except Exception as _e: + print(f"[PS-diag] {fname} save failed: {_e}", flush=True) _ROPE_PREQK_BUFFER: dict[int, dict] | None = None @@ -283,7 +301,7 @@ def dump_rope_preqk_verl080(layer_number: int, "query": query.detach().cpu().clone(), "key": key.detach().cpu().clone(), } - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("rope_preqk.pt", _ROPE_PREQK_BUFFER, dump_dir) _ROPE_PREQK_BUFFER = None @@ -309,7 +327,7 @@ def dump_expanded_kv_on(layer_number: int, expanded_key: torch.Tensor, "key": expanded_key.detach().cpu().clone(), "value": expanded_value.detach().cpu().clone(), } - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("expanded_kv.pt", _EXPANDED_KV_BUFFER, dump_dir) _EXPANDED_KV_BUFFER = None @@ -334,7 +352,7 @@ def dump_full_kv_off(layer_number: int, key: torch.Tensor, value: torch.Tensor, "key": key.detach().cpu().clone(), "value": value.detach().cpu().clone(), } - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("full_kv.pt", _FULL_KV_BUFFER, dump_dir) _FULL_KV_BUFFER = None @@ -357,7 +375,7 @@ def dump_build_kv_input_v_on(layer_number: int, value: torch.Tensor, if _BUILD_KV_INPUT_V_BUFFER is None: _BUILD_KV_INPUT_V_BUFFER = {} _BUILD_KV_INPUT_V_BUFFER[layer_number] = value.detach().cpu().clone() - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("build_kv_input_v.pt", _BUILD_KV_INPUT_V_BUFFER, dump_dir) _BUILD_KV_INPUT_V_BUFFER = None @@ -379,7 +397,7 @@ def dump_hidden_states_on(layer_number: int, hidden_states: torch.Tensor, if _HIDDEN_STATES_BUFFER is None: _HIDDEN_STATES_BUFFER = {} _HIDDEN_STATES_BUFFER[layer_number] = hidden_states.detach().cpu().clone() - if layer_number == num_layers: + if layer_number == _stage_last_layer(num_layers): _flush_dict_buffer("hidden_states.pt", _HIDDEN_STATES_BUFFER, dump_dir) _HIDDEN_STATES_BUFFER = None From ae48d978679aa7a005e5a0782118747b98ae906e Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 1 Jul 2026 11:25:42 +0800 Subject: [PATCH 56/61] [refactor] comprehensive variable rename in cmp_diag.py and cmp_diag_verl080.py MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Key renames: - fp→filepath, fa/fb→filepath_on/off, da/db→attn_dict_on/off - ma/mb→meta_on/off, la/lb→layers_on/off - d→attn_dict/rope_dict/kv_dict (context-specific) - s→shard_tensor, t→tp_rank, pl_fp→prefix_lens_filepath - _q→query_tensor, _name→result_name, _D→head_dim - _sec→section_label, _stage→stage_label, _crop→crop_note - _on_row1_end→on_row1_end, _cmp_len→compare_length, etc. - Various underscore-prefixed locals→descriptive names --- .../prefix_sharing/tools/cmp_diag.py | 180 ++++++++--------- .../prefix_sharing/tools/cmp_diag_verl080.py | 182 +++++++++--------- 2 files changed, 181 insertions(+), 181 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag.py b/prefix-sharing/prefix_sharing/tools/cmp_diag.py index 46fce34a..6dcc1df7 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag.py @@ -232,30 +232,30 @@ def _logits_ensure_token_major(logits_on: torch.Tensor, logits_off: torch.Tensor # ══════════════════════════════════════════════════════════════════ def _load_tensor(dir_path: str, filename: str) -> torch.Tensor | None: - fp = os.path.join(dir_path, filename) - return torch.load(fp, weights_only=True).float() if os.path.exists(fp) else None + filepath = os.path.join(dir_path, filename) + return torch.load(filepath, weights_only=True).float() if os.path.exists(filepath) else None def _load_packed_meta(dir_path: str, cu_fname: str = "cu_seqlens_q.pt") -> dict | None: - fp = os.path.join(dir_path, cu_fname) - if not os.path.exists(fp): - fp = os.path.join(dir_path, "cu_seqlens_q.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, cu_fname) + if not os.path.exists(filepath): + filepath = os.path.join(dir_path, "cu_seqlens_q.pt") + if not os.path.exists(filepath): return None - pl_fp = os.path.join(dir_path, "prefix_lens.pt") - if not os.path.exists(pl_fp): + prefix_lens_filepath = os.path.join(dir_path, "prefix_lens.pt") + if not os.path.exists(prefix_lens_filepath): return None - return {"cu_seqlens": torch.load(fp, weights_only=True), - "prefix_lens": torch.load(pl_fp, weights_only=True)} + return {"cu_seqlens": torch.load(filepath, weights_only=True), + "prefix_lens": torch.load(prefix_lens_filepath, weights_only=True)} def _load_attn_output(dir_path: str, layer: int) -> torch.Tensor | None: """Load a single layer's attn_output from attn_outputs.pt dict.""" - fp = os.path.join(dir_path, "attn_outputs.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, "attn_outputs.pt") + if not os.path.exists(filepath): return None - d = torch.load(fp, weights_only=True) + d = torch.load(filepath, weights_only=True) return d.get(layer) if isinstance(d, dict) else None @@ -271,17 +271,17 @@ def _load_attention_mask_2d(dir_path: str) -> torch.Tensor | None: runs). ON mask is a strict subset of OFF mask per row, which ``_build_alignment_mask_from_2d`` relies on. """ - fp = os.path.join(dir_path, "attention_mask.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, "attention_mask.pt") + if not os.path.exists(filepath): return None - return torch.load(fp, weights_only=True).to(torch.bool) + return torch.load(filepath, weights_only=True).to(torch.bool) def _get_num_layers(dir_path: str) -> int: - fp = os.path.join(dir_path, "attn_outputs.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, "attn_outputs.pt") + if not os.path.exists(filepath): return 0 - d = torch.load(fp, weights_only=True) + d = torch.load(filepath, weights_only=True) return max(d.keys()) if isinstance(d, dict) and d else 0 @@ -399,16 +399,16 @@ def cmp_rope_postqk(dir_on: str, dir_off: str) -> CheckResult | None: 3. For reuser rows: ON uses absolute positions (prefix_len..), OFF uses relative (0..). Positions differ by design — skip direct comparison. """ - fa = os.path.join(dir_on, "rope_postqk.pt") - fb = os.path.join(dir_off, "rope_postqk.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "rope_postqk.pt") + filepath_off = os.path.join(dir_off, "rope_postqk.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - a = torch.load(fa, weights_only=True) - b = torch.load(fb, weights_only=True) + a = torch.load(filepath_on, weights_only=True) + b = torch.load(filepath_off, weights_only=True) if not isinstance(a, dict) or not isinstance(b, dict): return None - la, lb = set(a.keys()), set(b.keys()) - if la != lb: + layers_on, layers_off =set(a.keys()), set(b.keys()) + if layers_on != layers_off: return CheckResult(name="rope_postqk", passed=False, metrics={"error": "layer set mismatch"}) @@ -419,9 +419,9 @@ def cmp_rope_postqk(dir_on: str, dir_off: str) -> CheckResult | None: first_row_q_diff = 0.0 first_row_k_diff = 0.0 first_row_len = 0 - num_layers = len(la) + num_layers = len(layers_on) - for layer_idx in sorted(la): + for layer_idx in sorted(layers_on): on_entry = a[layer_idx] off_entry = b[layer_idx] on_q, on_k = on_entry["query"], on_entry["key"] @@ -441,40 +441,40 @@ def cmp_rope_postqk(dir_on: str, dir_off: str) -> CheckResult | None: off_pos_t = off_pos.long() # Find first-row extent in ON: tokens before the first position reset # (position decreases or jumps to prefix_start) - _on_row1_end = 1 - for _i in range(1, len(on_pos_t)): - if on_pos_t[_i] <= on_pos_t[_i - 1]: + on_row1_end = 1 + for j in range(1, len(on_pos_t)): + if on_pos_t[j] <= on_pos_t[j - 1]: break - _on_row1_end = _i + 1 - _on_row1_len = _on_row1_end # positions 0..L-1 + on_row1_end = j + 1 + on_row1_len = on_row1_end # positions 0..L-1 # First row in OFF: positions go 0..L-1 (find matching extent) - _off_row1_end = 1 - for _i in range(1, len(off_pos_t)): - if off_pos_t[_i] <= off_pos_t[_i - 1]: + off_row1_end = 1 + for j in range(1, len(off_pos_t)): + if off_pos_t[j] <= off_pos_t[j - 1]: break - _off_row1_end = _i + 1 - _off_row1_len = _off_row1_end + off_row1_end = j + 1 + off_row1_len = off_row1_end - _cmp_len = min(_on_row1_len, _off_row1_len) - if _cmp_len > 0: - first_row_len = max(first_row_len, _cmp_len) + compare_length = min(on_row1_len, off_row1_len) + if compare_length > 0: + first_row_len = max(first_row_len, compare_length) # Match by position ID within the first row - for _pos_id in range(_cmp_len): - _on_idx = _pos_id # ON row 1 starts at packed index 0 - _off_idx = _pos_id # OFF row 1 starts at packed index 0 + for position_id in range(compare_length): + on_index = position_id # ON row 1 starts at packed index 0 + off_index = position_id # OFF row 1 starts at packed index 0 first_row_q_diff = max(first_row_q_diff, - float((on_q[_on_idx] - off_q[_off_idx]).abs().max())) + float((on_q[on_index] - off_q[off_index]).abs().max())) first_row_k_diff = max(first_row_k_diff, - float((on_k[_on_idx] - off_k[_off_idx]).abs().max())) + float((on_k[on_index] - off_k[off_index]).abs().max())) max_diff_q = max(max_diff_q, first_row_q_diff) max_diff_k = max(max_diff_k, first_row_k_diff) else: # Fallback: no positions available, compare first row by direct offset - _cmp_len = min(on_q.shape[0], off_q.shape[0], 128) - first_row_len = _cmp_len - first_row_q_diff = float((on_q[:_cmp_len] - off_q[:_cmp_len]).abs().max()) - first_row_k_diff = float((on_k[:_cmp_len] - off_k[:_cmp_len]).abs().max()) + compare_length = min(on_q.shape[0], off_q.shape[0], 128) + first_row_len = compare_length + first_row_q_diff = float((on_q[:compare_length] - off_q[:compare_length]).abs().max()) + first_row_k_diff = float((on_k[:compare_length] - off_k[:compare_length]).abs().max()) max_diff_q = first_row_q_diff max_diff_k = first_row_k_diff @@ -506,40 +506,40 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: The two are aligned to suffix-only via the same attention_mask dual-pointer logic used for attn_outputs / logits. """ - fa = os.path.join(dir_on, "rope_freqs_on.pt") - fb = os.path.join(dir_off, "rope_freqs_off.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "rope_freqs_on.pt") + filepath_off = os.path.join(dir_off, "rope_freqs_off.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None - la, lb = set(on_dict.keys()), set(off_dict.keys()) - if la != lb: + layers_on, layers_off =set(on_dict.keys()), set(off_dict.keys()) + if layers_on != layers_off: return CheckResult(name="rope_freqs", passed=False, metrics={"error": "layer set mismatch", - "on_layers": sorted(la), - "off_layers": sorted(lb)}) + "on_layers": sorted(layers_on), + "off_layers": sorted(layers_off)}) # Load OFF metadata for per-token reconstruction + alignment - mb = _load_packed_meta(dir_off) - if mb is None: + meta_off = _load_packed_meta(dir_off) + if meta_off is None: return CheckResult(name="rope_freqs", passed=False, metrics={"error": "OFF cu_seqlens missing"}) - cu_off = mb["cu_seqlens"] + cu_off = meta_off["cu_seqlens"] T_off = int(cu_off[-1]) if cu_off.numel() > 0 else 0 # Build alignment mask (prefer 2D attention_mask) mask_on_2d = _load_attention_mask_2d(dir_on) mask_off_2d = _load_attention_mask_2d(dir_off) - ma = _load_packed_meta(dir_on) + meta_on = _load_packed_meta(dir_on) if mask_on_2d is not None and mask_off_2d is not None: align_mask = _build_alignment_mask_from_2d( mask_on_2d, mask_off_2d, cu_off, T_off) - elif ma is not None: + elif meta_on is not None: align_mask = _build_alignment_mask( - cu_off, ma["prefix_lens"], T_off) + cu_off, meta_on["prefix_lens"], T_off) else: return CheckResult(name="rope_freqs", passed=False, metrics={"error": "cannot build alignment mask"}) @@ -549,7 +549,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: max_diff = 0.0 mismatches: list[dict] = [] # [{layer, token_idx, dim, on_val, off_val, diff}] - for layer_idx in sorted(la): + for layer_idx in sorted(layers_on): on_freqs = on_dict[layer_idx] # [T_on, 1, 1, D] # Reconstruct OFF per-token for this layer off_freqs = torch.cat( @@ -582,7 +582,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str) -> CheckResult | None: "diff": float(token_diff.values[t]), }) - metrics: dict = {"max_diff": max_diff, "num_layers": len(la)} + metrics: dict = {"max_diff": max_diff, "num_layers": len(layers_on)} if mismatches: metrics["mismatches"] = mismatches[:20] # cap to top 20 metrics["total_mismatches"] = len(mismatches) @@ -621,46 +621,46 @@ def _cos_for_layer(a, b, layer_idx, need_align, align_mask): b = _load_attn_output(dir_off, layer) if a is None or b is None: return None - ma = _load_packed_meta(dir_on) + meta_on = _load_packed_meta(dir_on) align_mask = None - need_align = (ma is not None and a.shape[0] != b.shape[0]) + need_align = (meta_on is not None and a.shape[0] != b.shape[0]) if need_align: - mb = _load_packed_meta(dir_off) - T = int(mb["cu_seqlens"][-1]) if mb and mb["cu_seqlens"].numel() > 0 else b.shape[0] + meta_off = _load_packed_meta(dir_off) + T = int(meta_off["cu_seqlens"][-1]) if meta_off and meta_off["cu_seqlens"].numel() > 0 else b.shape[0] # Prefer 2D attention_mask.pt for exact suffix alignment (dp aligned) mask_on_2d = _load_attention_mask_2d(dir_on) mask_off_2d = _load_attention_mask_2d(dir_off) - if mask_on_2d is not None and mask_off_2d is not None and mb is not None: + if mask_on_2d is not None and mask_off_2d is not None and meta_off is not None: align_mask = _build_alignment_mask_from_2d( - mask_on_2d, mask_off_2d, mb["cu_seqlens"], T) + mask_on_2d, mask_off_2d, meta_off["cu_seqlens"], T) else: align_mask = _build_alignment_mask( - mb["cu_seqlens"], ma["prefix_lens"], T) + meta_off["cu_seqlens"], meta_on["prefix_lens"], T) d = _cos_for_layer(a, b, layer, need_align, align_mask) d["layer"] = layer return d # All-layers mode - fa = os.path.join(dir_on, "attn_outputs.pt") - fb = os.path.join(dir_off, "attn_outputs.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "attn_outputs.pt") + filepath_off = os.path.join(dir_off, "attn_outputs.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - da = torch.load(fa, weights_only=True) - db = torch.load(fb, weights_only=True) - if not isinstance(da, dict) or not isinstance(db, dict): + attn_dict_on = torch.load(filepath_on, weights_only=True) + attn_dict_off = torch.load(filepath_off, weights_only=True) + if not isinstance(attn_dict_on, dict) or not isinstance(attn_dict_off, dict): return None # Build alignment mask once from ON metadata - ma = _load_packed_meta(dir_on) + meta_on = _load_packed_meta(dir_on) align_mask = None need_align = False - if ma is not None: - mb = _load_packed_meta(dir_off) - T = int(mb["cu_seqlens"][-1]) if mb and mb["cu_seqlens"].numel() > 0 else 0 + if meta_on is not None: + meta_off = _load_packed_meta(dir_off) + T = int(meta_off["cu_seqlens"][-1]) if meta_off and meta_off["cu_seqlens"].numel() > 0 else 0 if T > 0: # Check if any layer has shape mismatch - for layer_idx in da: - if layer_idx in db and da[layer_idx].shape != db[layer_idx].shape: + for layer_idx in attn_dict_on: + if layer_idx in attn_dict_off and attn_dict_on[layer_idx].shape != attn_dict_off[layer_idx].shape: need_align = True break if need_align: @@ -669,14 +669,14 @@ def _cos_for_layer(a, b, layer_idx, need_align, align_mask): mask_off_2d = _load_attention_mask_2d(dir_off) if mask_on_2d is not None and mask_off_2d is not None: align_mask = _build_alignment_mask_from_2d( - mask_on_2d, mask_off_2d, mb["cu_seqlens"], T) + mask_on_2d, mask_off_2d, meta_off["cu_seqlens"], T) else: align_mask = _build_alignment_mask( - mb["cu_seqlens"], ma["prefix_lens"], T) + meta_off["cu_seqlens"], meta_on["prefix_lens"], T) results = {} - for layer_idx in sorted(set(da.keys()) & set(db.keys())): - results[layer_idx] = _cos_for_layer(da[layer_idx], db[layer_idx], layer_idx, need_align, align_mask) + for layer_idx in sorted(set(attn_dict_on.keys()) & set(attn_dict_off.keys())): + results[layer_idx] = _cos_for_layer(attn_dict_on[layer_idx], attn_dict_off[layer_idx], layer_idx, need_align, align_mask) return results diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index b3452e47..685916f3 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -147,8 +147,8 @@ def _vec_metrics(a_vec: torch.Tensor, b_vec: torch.Tensor) -> dict: # ════════════════════════════════════════════════════════════════ def _load_tensor(dir_path: str, filename: str) -> torch.Tensor | None: - fp = os.path.join(dir_path, filename) - return torch.load(fp, weights_only=True).float() if os.path.exists(fp) else None + filepath = os.path.join(dir_path, filename) + return torch.load(filepath, weights_only=True).float() if os.path.exists(filepath) else None def _load_manifest(dir_path: str) -> dict | None: @@ -157,11 +157,11 @@ def _load_manifest(dir_path: str) -> dict | None: Returns None when absent (single-card or pre-manifest dumps) → callers fall back to tp_size==1 behavior (plain filenames, single-card compatible). """ - fp = os.path.join(dir_path, "parallel_info.json") - if not os.path.exists(fp): + manifest_filepath = os.path.join(dir_path, "parallel_info.json") + if not os.path.exists(manifest_filepath): return None try: - with open(fp, encoding="utf-8") as f: + with open(manifest_filepath, encoding="utf-8") as f: return json.load(f) except Exception: return None @@ -183,44 +183,44 @@ def _load_logits(dir_path: str, manifest: dict | None = None) -> torch.Tensor | if tp_size <= 1: return _load_tensor(dir_path, "logits.pt") shards = [] - for t in range(tp_size): - s = _load_tensor(dir_path, f"logits_tp{t}.pt") - if s is None: + for tp_rank in range(tp_size): + shard_tensor = _load_tensor(dir_path, f"logits_tp{tp_rank}.pt") + if shard_tensor is None: return None - shards.append(s) + shards.append(shard_tensor) return torch.cat(shards, dim=-1) def _load_packed_meta(dir_path: str, cu_fname: str = "cu_seqlens_q.pt") -> dict | None: """加载 cu_seqlens + prefix_lens（suffix 对齐所需）。""" - fp = os.path.join(dir_path, cu_fname) - if not os.path.exists(fp): - fp = os.path.join(dir_path, "cu_seqlens_q.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, cu_fname) + if not os.path.exists(filepath): + filepath = os.path.join(dir_path, "cu_seqlens_q.pt") + if not os.path.exists(filepath): return None - pl_fp = os.path.join(dir_path, "prefix_lens.pt") - if not os.path.exists(pl_fp): + prefix_lens_filepath = os.path.join(dir_path, "prefix_lens.pt") + if not os.path.exists(prefix_lens_filepath): return None - return {"cu_seqlens": torch.load(fp, weights_only=True), - "prefix_lens": torch.load(pl_fp, weights_only=True)} + return {"cu_seqlens": torch.load(filepath, weights_only=True), + "prefix_lens": torch.load(prefix_lens_filepath, weights_only=True)} def _load_attn_output(dir_path: str, layer: int) -> torch.Tensor | None: """加载单层 attn_output（attn_outputs.pt = dict {layer: tensor}）。""" - fp = os.path.join(dir_path, "attn_outputs.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, "attn_outputs.pt") + if not os.path.exists(filepath): return None - d = torch.load(fp, weights_only=True) - return d.get(layer) if isinstance(d, dict) else None + attn_dict = torch.load(filepath, weights_only=True) + return attn_dict.get(layer) if isinstance(attn_dict, dict) else None def _get_num_layers(dir_path: str) -> int: - fp = os.path.join(dir_path, "attn_outputs.pt") - if not os.path.exists(fp): + filepath = os.path.join(dir_path, "attn_outputs.pt") + if not os.path.exists(filepath): return 0 - d = torch.load(fp, weights_only=True) - return max(d.keys()) if isinstance(d, dict) and d else 0 + attn_dict = torch.load(filepath, weights_only=True) + return max(attn_dict.keys()) if isinstance(attn_dict, dict) and attn_dict else 0 # ════════════════════════════════════════════════════════════════ @@ -331,15 +331,15 @@ def _cos_for_layer(a: torch.Tensor, b: torch.Tensor, def _build_attn_align_mask(dir_on: str, dir_off: str) -> torch.Tensor | None: """从 OFF cu_seqlens + ON prefix_lens 构建 suffix 对齐 mask（None=无法构建）。""" - ma = _load_packed_meta(dir_on) - mb = _load_packed_meta(dir_off) - if ma is None or mb is None: + meta_on = _load_packed_meta(dir_on) + meta_off = _load_packed_meta(dir_off) + if meta_on is None or meta_off is None: return None - cu_off = mb["cu_seqlens"] + cu_off = meta_off["cu_seqlens"] T = int(cu_off[-1]) if cu_off.numel() > 0 else 0 if T == 0: return None - return _build_alignment_mask(cu_off, ma["prefix_lens"], T) + return _build_alignment_mask(cu_off, meta_on["prefix_lens"], T) def cmp_attn_layer(dir_on: str, dir_off: str, @@ -366,18 +366,18 @@ def cmp_attn_layer(dir_on: str, dir_off: str, passed=d["cos_avg"] > _COS_AVG_PASS and d["cos_min"] > _COS_MIN_PASS, metrics=d) - fa = os.path.join(dir_on, "attn_outputs.pt") - fb = os.path.join(dir_off, "attn_outputs.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "attn_outputs.pt") + filepath_off = os.path.join(dir_off, "attn_outputs.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - da = torch.load(fa, weights_only=True) - db = torch.load(fb, weights_only=True) - if not isinstance(da, dict) or not isinstance(db, dict): + attn_dict_on = torch.load(filepath_on, weights_only=True) + attn_dict_off = torch.load(filepath_off, weights_only=True) + if not isinstance(attn_dict_on, dict) or not isinstance(attn_dict_off, dict): return None results = {} - for layer_idx in sorted(set(da.keys()) & set(db.keys())): - a, b = da[layer_idx], db[layer_idx] + for layer_idx in sorted(set(attn_dict_on.keys()) & set(attn_dict_off.keys())): + a, b = attn_dict_on[layer_idx], attn_dict_off[layer_idx] need = align_mask is not None and a.shape[0] != b.shape[0] try: results[layer_idx] = _cos_for_layer(a, b, align_mask if need else None) @@ -453,10 +453,10 @@ def _load_rope_postqk(dir_path: str, layer: int, fname: str = "rope_postqk.pt" Returns ``(query, key)`` or ``(None, None)``. """ - fp = os.path.join(dir_path, fname) - if not os.path.exists(fp): + filepath = os.path.join(dir_path, fname) + if not os.path.exists(filepath): return None, None - d = torch.load(fp, weights_only=True) + d = torch.load(filepath, weights_only=True) if not isinstance(d, dict): return None, None entry = d.get(layer) @@ -513,18 +513,18 @@ def _cmp_rope_stage_layer(dir_on: str, dir_off: str, layer: int | None, return CheckResult(name=f"{label}_L{layer}", passed=ok, metrics=d) # All layers - fa = os.path.join(dir_on, fname) - fb = os.path.join(dir_off, fname) - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, fname) + filepath_off = os.path.join(dir_off, fname) + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - da = torch.load(fa, weights_only=True) - db = torch.load(fb, weights_only=True) - if not isinstance(da, dict) or not isinstance(db, dict): + attn_dict_on = torch.load(filepath_on, weights_only=True) + attn_dict_off = torch.load(filepath_off, weights_only=True) + if not isinstance(attn_dict_on, dict) or not isinstance(attn_dict_off, dict): return None results = {} - for layer_idx in sorted(set(da.keys()) & set(db.keys())): - ea, eb = da[layer_idx], db[layer_idx] + for layer_idx in sorted(set(attn_dict_on.keys()) & set(attn_dict_off.keys())): + ea, eb = attn_dict_on[layer_idx], attn_dict_off[layer_idx] qa, ka = ea.get("query"), ea.get("key") qb, kb = eb.get("query"), eb.get("key") if qa is None or qb is None: @@ -700,12 +700,12 @@ def _load_rope_freqs_vec_at_pos(dir_on: str, dir_off: str, layer: int, pos: int, 供 top-K 跨 stage 对齐用（freqs dim = Q/K dim % D，角度按 head_dim 共享）。 """ - fa = os.path.join(dir_on, "rope_freqs.pt") - fb = os.path.join(dir_off, "rope_freqs.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "rope_freqs.pt") + filepath_off = os.path.join(dir_off, "rope_freqs.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None, None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None, None if layer not in on_dict or layer not in off_dict: @@ -717,7 +717,7 @@ def _load_rope_freqs_vec_at_pos(dir_on: str, dir_off: str, layer: int, pos: int, aligned_result = _align_rope_freqs_layer(on_dict[layer], off_dict[layer], align_mask) if aligned_result is None: return None, None - on_a, off_a = _aligned + on_a, off_a = aligned_result if pos < 0 or pos >= on_a.shape[0]: return None, None return on_a[pos].reshape(-1), off_a[pos].reshape(-1) @@ -731,12 +731,12 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, suffix 对齐后逐元素比。角度是 RoPE 输入，应精确相等（max_diff==0）。 ``layer`` 给定则只比该层。 """ - fa = os.path.join(dir_on, "rope_freqs.pt") - fb = os.path.join(dir_off, "rope_freqs.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "rope_freqs.pt") + filepath_off = os.path.join(dir_off, "rope_freqs.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None @@ -759,7 +759,7 @@ def cmp_rope_freqs(dir_on: str, dir_off: str, aligned_result = _align_rope_freqs_layer(on_dict[layer_idx], off_dict[layer_idx], align_mask) if aligned_result is None: continue - on_a, off_a = _aligned + on_a, off_a = aligned_result diff = (on_a - off_a).abs() # [N,1,1,D] md = float(diff.max()) max_diff = max(max_diff, md) @@ -791,12 +791,12 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, 取 ``layer``（默认最后一层）对齐后第 ``pos`` 个 token 的角度向量 [D]，比 ON/OFF。 角度是 RoPE 输入，应逐元素相等 → max_abs 应为 0。 """ - fa = os.path.join(dir_on, "rope_freqs.pt") - fb = os.path.join(dir_off, "rope_freqs.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "rope_freqs.pt") + filepath_off = os.path.join(dir_off, "rope_freqs.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None common = set(on_dict.keys()) & set(off_dict.keys()) @@ -813,7 +813,7 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, aligned_result = _align_rope_freqs_layer(on_dict[rf_layer], off_dict[rf_layer], align_mask) if aligned_result is None: return CheckResult(name=result_name, metrics={"error": "对齐失败"}) - on_a, off_a = _aligned + on_a, off_a = aligned_result n = on_a.shape[0] if pos < 0 or pos >= n: return CheckResult(name=result_name, @@ -831,13 +831,13 @@ def cmp_rope_freqs_token(dir_on: str, dir_off: str, pos: int, def _load_attn_kv(dir_path: str, layer: int, fname: str) -> tuple[torch.Tensor | None, torch.Tensor | None]: """load {key, value} for a layer from fname. Returns (key, value) or (None, None).""" - fp = os.path.join(dir_path, fname) - if not os.path.exists(fp): + filepath = os.path.join(dir_path, fname) + if not os.path.exists(filepath): return None, None - d = torch.load(fp, weights_only=True) - if not isinstance(d, dict): + kv_dict = torch.load(filepath, weights_only=True) + if not isinstance(kv_dict, dict): return None, None - entry = d.get(layer) + entry = kv_dict.get(layer) if entry is None: return None, None return entry.get("key"), entry.get("value") @@ -862,7 +862,7 @@ def cmp_attn_kv(dir_on: str, dir_off: str, layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) if layer is not None: layers = [l for l in layers if l == layer] - resultresult_name = f"attn_kv_L{layer}" if layer is not None else "attn_kv" + result_name = f"attn_kv_L{layer}" if layer is not None else "attn_kv" if not layers: return CheckResult(name=result_name, passed=False, metrics={"error": f"layer {layer} 不在双方 attn_kv 中"}) @@ -942,12 +942,12 @@ def cmp_build_kv_input_v(dir_on: str, dir_off: str, 应逐元素相同——若不同则问题在 QKV 投影阶段（hidden_states / QKV 权重）。 ON_T vs OFF_T 还能看出 ON 有没有把 hidden_states 裁成 suffix-only。 """ - fa = os.path.join(dir_on, "build_kv_input_v.pt") - fb = os.path.join(dir_off, "build_kv_input_v.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "build_kv_input_v.pt") + filepath_off = os.path.join(dir_off, "build_kv_input_v.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) @@ -1022,12 +1022,12 @@ def cmp_hidden_states(dir_on: str, dir_off: str, 这是 QKV 投影的 INPUT。如果 hidden_states 一致但 V 不一致 → GEMM 精度差异； 如果 hidden_states 就不一致 → 根因在上游（embedding / input_layernorm）。 """ - fa = os.path.join(dir_on, "hidden_states.pt") - fb = os.path.join(dir_off, "hidden_states.pt") - if not os.path.exists(fa) or not os.path.exists(fb): + filepath_on = os.path.join(dir_on, "hidden_states.pt") + filepath_off = os.path.join(dir_off, "hidden_states.pt") + if not os.path.exists(filepath_on) or not os.path.exists(filepath_off): return None - on_dict = torch.load(fa, weights_only=True) - off_dict = torch.load(fb, weights_only=True) + on_dict = torch.load(filepath_on, weights_only=True) + off_dict = torch.load(filepath_off, weights_only=True) if not isinstance(on_dict, dict) or not isinstance(off_dict, dict): return None layers = sorted(set(on_dict.keys()) & set(off_dict.keys())) @@ -1103,10 +1103,10 @@ def _load_mask_2d(dir_path: str, mask_kind: str, tag: str) -> torch.Tensor | Non if mask_kind == "none": return None fname = f"{mask_kind}_mask_{tag}.pt" # label_mask_{tag} / attention_mask_{tag} - fp = os.path.join(dir_path, fname) - if not os.path.exists(fp): + filepath = os.path.join(dir_path, fname) + if not os.path.exists(filepath): return None - return torch.load(fp, weights_only=True).to(torch.bool) + return torch.load(filepath, weights_only=True).to(torch.bool) def _resolve_mask(dir_off: str, mask_kind: str, tag: str, @@ -1172,18 +1172,18 @@ def cmp_2d(dir_on: str, dir_off: str, filename: str, name: str, # ════════════════════════════════════════════════════════════════ def _shape_of(dir_path: str, filename: str) -> str: - fp = os.path.join(dir_path, filename) - if not os.path.exists(fp): + filepath = os.path.join(dir_path, filename) + if not os.path.exists(filepath): return "(missing)" try: - obj = torch.load(fp, weights_only=True) + obj = torch.load(filepath, weights_only=True) if isinstance(obj, dict): # per-layer dict（attn_outputs / rope_freqs_*）：显示层数 + 首层 shape sample = next(iter(obj.values())) if obj else None # rope_postqk.pt：每层值是 {"query","key"[,"positions"]} dict，取 query 的 shape 代表 if isinstance(sample, dict): - _q = sample.get("query") - sample_shape = f",Q{tuple(_q.shape)}" if _q is not None else "" + query_tensor = sample.get("query") + sample_shape = f",Q{tuple(query_tensor.shape)}" if query_tensor is not None else "" elif sample is not None: sample_shape = f",{tuple(sample.shape)}" else: From 716b4163ab02bbc4216ddb6e130f4d0ec96ad662 Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 1 Jul 2026 12:03:04 +0800 Subject: [PATCH 57/61] [fix] assemble_dump: rewrite parallel_info.json with tp=1 pp=1 after assembly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Root cause: assemble_dump.py copied parallel_info.json verbatim from raw dump, leaving tp_size/pp_size > 1 and scopes like 'tp_vocab'/'pp_stage'. cmp_diag then tried to read logits_tp0.pt instead of logits.pt → showed (missing). Fix: after merging all shards, update parallel_info.json to tp_size=1, pp_size=1, all scopes='global', matching the single-card assembled directory. --- .../prefix_sharing/tools/assemble_dump.py | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/prefix-sharing/prefix_sharing/tools/assemble_dump.py b/prefix-sharing/prefix_sharing/tools/assemble_dump.py index d0a23c66..ce25e8d4 100644 --- a/prefix-sharing/prefix_sharing/tools/assemble_dump.py +++ b/prefix-sharing/prefix_sharing/tools/assemble_dump.py @@ -156,6 +156,26 @@ def assemble(input_dir: str, output_dir: str) -> None: shutil.copy2(src, os.path.join(output_dir, "logits.pt")) print(" [copy] logits.pt") + # ── Rewrite manifest: assembled output is single-card ────────── + # parallel_info.json was copied verbatim from raw dump and still + # records tp_size/pp_size > 1. The assembled directory has merged + # all shards → update to tp_size=1 / pp_size=1 so cmp_diag reads + # plain ``logits.pt`` (not ``logits_tp0.pt``), etc. + if is_multi_rank: + manifest_out_path = os.path.join(output_dir, "parallel_info.json") + if os.path.exists(manifest_out_path): + with open(manifest_out_path, encoding="utf-8") as f: + manifest_out = json.load(f) + manifest_out["tp_size"] = 1 + manifest_out["pp_size"] = 1 + manifest_out["cp_size"] = 1 + # all scopes are now "global" — no sharding in assembled output + for stem in manifest_out.get("scopes", {}): + manifest_out["scopes"][stem] = "global" + with open(manifest_out_path, "w", encoding="utf-8") as f: + json.dump(manifest_out, f, indent=2, ensure_ascii=False) + print(" [update] parallel_info.json → tp=1 pp=1") + # ── Summary ─────────────────────────────────────────────────── n_files = len(os.listdir(output_dir)) print(f"\n[assemble] done — {n_files} files written to {output_dir}") From 8b8a042bef821f7dd38adc324480c29a0b5bbad9 Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 1 Jul 2026 12:43:40 +0800 Subject: [PATCH 58/61] [refactor] remove TP/PP awareness from cmp_diag_verl080.py; revert assemble manifest rewrite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit cmp_diag_verl080.py: - _load_logits: remove TP vocab concat logic → just load logits.pt - _logits_shape: deleted - _load_manifest: deleted - _print_topology: deleted - _print_shapes: remove manifest params, treat logits like any other file - main(): remove manifest loading and topology printing Rationale: multi-rank assembly is done by assemble_dump.py before cmp runs; cmp should work on flat single-card data only — no parallel awareness. assemble_dump.py: revert the manifest rewrite (no longer needed since cmp doesn't read manifest). --- .../prefix_sharing/tools/assemble_dump.py | 20 ---- .../prefix_sharing/tools/cmp_diag_verl080.py | 91 ++----------------- 2 files changed, 8 insertions(+), 103 deletions(-) diff --git a/prefix-sharing/prefix_sharing/tools/assemble_dump.py b/prefix-sharing/prefix_sharing/tools/assemble_dump.py index ce25e8d4..d0a23c66 100644 --- a/prefix-sharing/prefix_sharing/tools/assemble_dump.py +++ b/prefix-sharing/prefix_sharing/tools/assemble_dump.py @@ -156,26 +156,6 @@ def assemble(input_dir: str, output_dir: str) -> None: shutil.copy2(src, os.path.join(output_dir, "logits.pt")) print(" [copy] logits.pt") - # ── Rewrite manifest: assembled output is single-card ────────── - # parallel_info.json was copied verbatim from raw dump and still - # records tp_size/pp_size > 1. The assembled directory has merged - # all shards → update to tp_size=1 / pp_size=1 so cmp_diag reads - # plain ``logits.pt`` (not ``logits_tp0.pt``), etc. - if is_multi_rank: - manifest_out_path = os.path.join(output_dir, "parallel_info.json") - if os.path.exists(manifest_out_path): - with open(manifest_out_path, encoding="utf-8") as f: - manifest_out = json.load(f) - manifest_out["tp_size"] = 1 - manifest_out["pp_size"] = 1 - manifest_out["cp_size"] = 1 - # all scopes are now "global" — no sharding in assembled output - for stem in manifest_out.get("scopes", {}): - manifest_out["scopes"][stem] = "global" - with open(manifest_out_path, "w", encoding="utf-8") as f: - json.dump(manifest_out, f, indent=2, ensure_ascii=False) - print(" [update] parallel_info.json → tp=1 pp=1") - # ── Summary ─────────────────────────────────────────────────── n_files = len(os.listdir(output_dir)) print(f"\n[assemble] done — {n_files} files written to {output_dir}") diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index 685916f3..f1aab81c 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -151,44 +151,13 @@ def _load_tensor(dir_path: str, filename: str) -> torch.Tensor | None: return torch.load(filepath, weights_only=True).float() if os.path.exists(filepath) else None -def _load_manifest(dir_path: str) -> dict | None: - """Load ``parallel_info.json`` written by the dump layer (topology + scopes). +def _load_logits(dir_path: str) -> torch.Tensor | None: + """Load packed logits from ``logits.pt``. - Returns None when absent (single-card or pre-manifest dumps) → callers fall - back to tp_size==1 behavior (plain filenames, single-card compatible). + Multi-rank assembly (tp vocab concat) is done by ``assemble_dump.py`` + before cmp is called; cmp works on flat single-card data only. """ - manifest_filepath = os.path.join(dir_path, "parallel_info.json") - if not os.path.exists(manifest_filepath): - return None - try: - with open(manifest_filepath, encoding="utf-8") as f: - return json.load(f) - except Exception: - return None - - -def _load_logits(dir_path: str, manifest: dict | None = None) -> torch.Tensor | None: - """Load packed logits, gathering tp vocab shards to full vocab when tp>1. - - tp_size==1 (or no manifest) → single ``logits.pt`` (single-card compatible). - tp_size>1 → concat ``logits_tp{0..tp-1}.pt`` on the vocab (last) dim, - reconstructing ``[N, V]`` so ON-vs-OFF compares on the same - full-vocab coordinate system as single-card. A missing shard - aborts the reconstruction (returns None) rather than silently - comparing partial vocab. - """ - if manifest is None: - manifest = _load_manifest(dir_path) - tp_size = (manifest or {}).get("tp_size", 1) - if tp_size <= 1: - return _load_tensor(dir_path, "logits.pt") - shards = [] - for tp_rank in range(tp_size): - shard_tensor = _load_tensor(dir_path, f"logits_tp{tp_rank}.pt") - if shard_tensor is None: - return None - shards.append(shard_tensor) - return torch.cat(shards, dim=-1) + return _load_tensor(dir_path, "logits.pt") def _load_packed_meta(dir_path: str, @@ -1194,43 +1163,9 @@ def _shape_of(dir_path: str, filename: str) -> str: return "(error)" -def _logits_shape(dir_path: str, manifest: dict | None) -> str: - """Shape string for logits, manifest-aware: tp>1 → show shard shape tagged. - Under TP the file is sharded (``logits_tp{r}.pt``); report one shard's shape - prefixed with ``tp{N}×`` so the shapes table still flags mismatches without - pretending a plain ``logits.pt`` exists. - """ - tp_size = (manifest or {}).get("tp_size", 1) - if tp_size <= 1: - return _shape_of(dir_path, "logits.pt") - s0 = _shape_of(dir_path, "logits_tp0.pt") - if s0 in ("(missing)", "(error)"): - return s0 - return f"tp{tp_size}×{s0}" - - -def _print_topology(manifest_on: dict | None, manifest_off: dict | None) -> None: - """Print ON/OFF parallel topology from manifests; warn on mismatch.""" - def _topo(m): - if not m: - return "single-card (no manifest)" - return f"tp={m.get('tp_size', 1)} pp={m.get('pp_size', 1)} cp={m.get('cp_size', 1)}" - print(_SEP_SINGLE + "\n [topology] ON vs OFF parallel config") - print(_SEP_SINGLE) - print(f" ON : {_topo(manifest_on)}") - print(f" OFF: {_topo(manifest_off)}") - if manifest_on and manifest_off: - for key in ("tp_size", "pp_size", "cp_size"): - if manifest_on.get(key) != manifest_off.get(key): - print(f" {_CROSS} MISMATCH on {key}: ON={manifest_on.get(key)} " - f"OFF={manifest_off.get(key)} — comparison may be invalid") - print() - -def _print_shapes(dir_on: str, dir_off: str, tag: str, - manifest_on: dict | None = None, - manifest_off: dict | None = None): +def _print_shapes(dir_on: str, dir_off: str, tag: str): """打印 ON/OFF 各 .pt 文件 shape —— 定位 shape mismatch 根因的第一手信息。""" print(_SEP_SINGLE + "\n [shapes] ON vs OFF dump shapes") print(_SEP_SINGLE) @@ -1254,12 +1189,7 @@ def _print_shapes(dir_on: str, dir_off: str, tag: str, print(f" {'FILE':<28s} {'ON':<16s} {'OFF':<16s} {'STATUS'}") print(f" {'─' * 28} {'─' * 16} {'─' * 16} {'─' * 10}") for fname in files: - if fname == "logits.pt": - # TP-sharded: per-rank logits_tp{r}.pt, not a plain logits.pt - s_on = _logits_shape(dir_on, manifest_on) - s_off = _logits_shape(dir_off, manifest_off) - else: - s_on, s_off = _shape_of(dir_on, fname), _shape_of(dir_off, fname) + s_on, s_off = _shape_of(dir_on, fname), _shape_of(dir_off, fname) if s_on == "(missing)" or s_off == "(missing)": status = "—" elif s_on == s_off: @@ -1581,13 +1511,8 @@ def main(): _print_header(args.dir_on, args.dir_off, args.dir_off2, args.tag, args.mask, args.layer) - # ── parallel topology (manifest-driven: TP shards, future SP/PP) ── - manifest_on = _load_manifest(args.dir_on) - manifest_off = _load_manifest(args.dir_off) - _print_topology(manifest_on, manifest_off) - # ── shape diagnostics ── - _print_shapes(args.dir_on, args.dir_off, args.tag, manifest_on, manifest_off) + _print_shapes(args.dir_on, args.dir_off, args.tag) # ── resolve 2D mask ── ref = _load_tensor(args.dir_off, f"logprobs_{args.tag}.pt") From b7fe4009981b35ddd29853a4496814d6dac77d7f Mon Sep 17 00:00:00 2001 From: Boundless Date: Tue, 14 Jul 2026 20:32:56 +0800 Subject: [PATCH 59/61] [fix] replace torch.nested.nested_tensor with as_nested_tensor for NPU autograd compat On Ascend NPU, torch.nested.nested_tensor is intercepted by torch_npu and throws. verl/__init__.py unwraps it (via __wrapped__), but the unwrapped version does not support autograd backward on NPU tensors -- gradients are silently detached. In PP training, this causes the last stage's loss backward to never reach the model, so no send_backward P2P is issued, and all prior PP ranks hang forever at recv_backward. TP doesn't hang because there is no P2P dependency, even though gradients are also broken. Fix: replace all torch.nested.nested_tensor(rows, layout=torch.jagged) with torch.nested.as_nested_tensor(rows, layout=torch.jagged) in verl_mcore.py. The latter preserves the autograd graph correctly on NPU after unwrapping. Affected functions: - _fold_2d_to_nested (root cause of PP hang) - _trim_nested_batch (forward batch trimming) - _trim_plain_batch_thd (forward batch trimming) Co-Authored-By: Claude --- .../prefix_sharing/integrations/verl_mcore.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index 2743e1c7..fa12c38e 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -532,7 +532,7 @@ def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any: import torch rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))] - return torch.nested.nested_tensor(rows, layout=torch.jagged) + return torch.nested.as_nested_tensor(rows, layout=torch.jagged) def _clone_batch(batch: Any) -> Any: @@ -728,13 +728,13 @@ def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any: # 裁剪 input_ids NestedTensor trimmed_ids_seqs = _slice_nested_sequences(input_ids, plan) - new_input_ids = torch.nested.nested_tensor(trimmed_ids_seqs, layout=torch.jagged) + new_input_ids = torch.nested.as_nested_tensor(trimmed_ids_seqs, layout=torch.jagged) trimmed_batch["input_ids"] = new_input_ids # 裁剪 position_ids NestedTensor if _is_nested_tensor(position_ids): trimmed_pos_seqs = _slice_nested_sequences(position_ids, plan) - new_position_ids = torch.nested.nested_tensor(trimmed_pos_seqs, layout=torch.jagged) + new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged) else: # position_ids 是 2D tensor → 需要用 attention_mask 的 # valid_indices 切片（keep_range 是序列偏移，不是列索引） @@ -754,14 +754,14 @@ def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any: trimmed_pos_seqs = _slice_2d_position_rows( position_ids, plan, attention_mask_bool, ) - new_position_ids = torch.nested.nested_tensor(trimmed_pos_seqs, layout=torch.jagged) + new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged) trimmed_batch["position_ids"] = new_position_ids # loss_mask 也需要裁剪（如果存在） loss_mask = batch.get("loss_mask") if loss_mask is not None and _is_nested_tensor(loss_mask): trimmed_loss_seqs = _slice_nested_sequences(loss_mask, plan) - trimmed_batch["loss_mask"] = torch.nested.nested_tensor( + trimmed_batch["loss_mask"] = torch.nested.as_nested_tensor( trimmed_loss_seqs, layout=torch.jagged ) @@ -815,8 +815,8 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan) -> Any: # 用裁剪后的序列构建 NestedTensor（jagged layout） # 这样 preprocess_thd_engine 会从 offsets 正确计算 cu_seqlens trimmed_batch = _clone_batch(batch) - trimmed_batch["input_ids"] = torch.nested.nested_tensor(kept_id_rows, layout=torch.jagged) - trimmed_batch["position_ids"] = torch.nested.nested_tensor(kept_pos_rows, layout=torch.jagged) + trimmed_batch["input_ids"] = torch.nested.as_nested_tensor(kept_id_rows, layout=torch.jagged) + trimmed_batch["position_ids"] = torch.nested.as_nested_tensor(kept_pos_rows, layout=torch.jagged) # loss_mask loss_mask = batch.get("loss_mask") @@ -827,7 +827,7 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan) -> Any: keep_start, keep_end = plan.input_keep_ranges[row] kept_indices = indices[keep_start:keep_end] kept_loss_rows.append(loss_mask[row, kept_indices]) - trimmed_batch["loss_mask"] = torch.nested.nested_tensor( + trimmed_batch["loss_mask"] = torch.nested.as_nested_tensor( kept_loss_rows, layout=torch.jagged ) From 861be1358ad21b8b33753ab9e90a02c73d999a1a Mon Sep 17 00:00:00 2001 From: Boundless Date: Wed, 15 Jul 2026 21:01:26 +0800 Subject: [PATCH 60/61] [fix] resolve CI lint failures (F401/F541/F821/F841) - Add missing _load_manifest function to cmp_diag_verl080.py (F821) - Remove unused imports: Sequence, ensure_global_packed_token_lengths, inspect, torch, _cached_parallel_info, json, os (F401) - Fix f-strings without placeholders to plain strings (F541) - Remove unused variable device assignment (F841) All flagged by ruff --select F,E9 (pyflakes + syntax errors). --- .../prefix_sharing/integrations/verl_mcore.py | 30 ++++++++----------- .../prefix_sharing/tools/cmp_diag_verl080.py | 18 ++++++++++- .../tools/diagnostic_dump_verl080.py | 1 - .../tools/inject_baseline_synthetic.py | 2 -- 4 files changed, 29 insertions(+), 22 deletions(-) diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index fa12c38e..4b4cb228 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -22,7 +22,7 @@ import importlib from contextlib import contextmanager from dataclasses import dataclass -from typing import Any, Iterator, Mapping, Sequence +from typing import Any, Iterator, Mapping from prefix_sharing.backends.factory import get_backend_instance from prefix_sharing.backends.packed_layout import PackedBatchLayout @@ -34,8 +34,6 @@ from prefix_sharing.integrations.parallel_info import MegatronParallelInfo from prefix_sharing.integrations.parallel_info import get_megatron_parallel_info from prefix_sharing.integrations.patch_manager import PatchHandle -from prefix_sharing.utils import ensure_global_packed_token_lengths - @dataclass(frozen=True) class PrefixSharingRuntimeState: @@ -118,26 +116,25 @@ def build_prefix_sharing_micro_batch_verl070( # --- Path 1: prefix sharing disabled by config --- if not config.enable_prefix_sharing: - print(f"[PS][prepare] PATH 1: prefix sharing disabled (config.enable_prefix_sharing=False), returning (batch, None)") + print("[PS][prepare] PATH 1: prefix sharing disabled (config.enable_prefix_sharing=False), returning (batch, None)") return batch, None - print(f"[PS][prepare] config.enable_prefix_sharing=True, validating config...") + print("[PS][prepare] config.enable_prefix_sharing=True, validating config...") config.validate(model_config=model_config, integrate_mode="verl_megatron_actor") - print(f"[PS][prepare] config.validate() returned OK") + print("[PS][prepare] config.validate() returned OK") # --- Path 2: missing use_remove_padding --- - print(f"[PS][prepare] checking megatron.use_remove_padding...") + print("[PS][prepare] checking megatron.use_remove_padding...") if not _read_actor_bool(actor_config, "megatron.use_remove_padding", False): - print(f"[PS][prepare] PATH 2: megatron.use_remove_padding=False, raising RuntimeError") + print("[PS][prepare] PATH 2: megatron.use_remove_padding=False, raising RuntimeError") raise RuntimeError("prefix sharing phase 1 requires verl megatron.use_remove_padding=True") # --- Path 3: multi_modal check --- - print(f"[PS][prepare] use_remove_padding=True, about to batch.get(multi_modal_inputs)...") + print("[PS][prepare] use_remove_padding=True, about to batch.get(multi_modal_inputs)...") multi_modal_inputs = batch.get("multi_modal_inputs") if multi_modal_inputs is not None: # tensorclass 无法遍历（触发 CUDA 同步），改用底层 td 检查字段数 - import inspect is_tensorclass = hasattr(multi_modal_inputs, 'batch_size') print(f"[PS][prepare] multi_modal_inputs type: tensorclass={is_tensorclass}, type={type(multi_modal_inputs).__name__}") if is_tensorclass: @@ -148,9 +145,9 @@ def build_prefix_sharing_micro_batch_verl070( else: has_mm = any(mmi is not None and len(mmi.keys()) > 0 for mmi in multi_modal_inputs) if has_mm: - print(f"[PS][prepare] PATH 3: multi_modal_inputs has content, raising RuntimeError") + print("[PS][prepare] PATH 3: multi_modal_inputs has content, raising RuntimeError") raise RuntimeError("prefix sharing phase 1 supports only text-only actor micro-batches") - print(f"[PS][prepare] multi_modal check PASSED (no real multi-modal content)") + print("[PS][prepare] multi_modal check PASSED (no real multi-modal content)") # --- Read tensors --- attention_mask = batch["attention_mask"].to(bool) @@ -160,7 +157,7 @@ def build_prefix_sharing_micro_batch_verl070( # --- Path 4: wrong tensor dims --- if attention_mask.dim() != 2 or input_ids.dim() != 2 or position_ids.dim() != 2: - print(f"[PS][prepare] PATH 4: non-2D tensors detected, raising RuntimeError") + print("[PS][prepare] PATH 4: non-2D tensors detected, raising RuntimeError") raise RuntimeError("prefix sharing phase 1 expects 2D input_ids/attention_mask/position_ids") # --- Planning --- @@ -178,11 +175,11 @@ def build_prefix_sharing_micro_batch_verl070( # --- Path 5: no sharing found --- if not prefix_sharing_plan.has_sharing: - print(f"[PS][prepare] PATH 5: no sharing detected, returning (batch, None)") + print("[PS][prepare] PATH 5: no sharing detected, returning (batch, None)") return batch, None # --- Path 6: sharing found, trim the original micro-batch --- - print(f"[PS][prepare] PATH 6: sharing detected, preparing trimmed batch...") + print("[PS][prepare] PATH 6: sharing detected, preparing trimmed batch...") trimmed_micro_batch = _clone_batch(batch) new_attention_mask = attention_mask.clone() new_attention_mask[:] = False @@ -387,7 +384,6 @@ def restore_via_2d_unfold_verl080( Returns: ``output``（``log_probs``/``entropy`` 被替换为重组后的 NestedTensor）。 """ - import torch ctx = current_prefix_sharing_context() if ctx is None: @@ -422,8 +418,6 @@ def restore_via_2d_unfold_verl080( return output L_max = max(original_lengths) - device = log_probs_nested.values().device - # --- Step 1: 展开裁剪后 NestedTensor → 完整 2D [B, L_max] --- log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d( log_probs_nested, diff --git a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py index ded68a53..c30f0bfd 100644 --- a/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/cmp_diag_verl080.py @@ -160,6 +160,22 @@ def _load_logits(dir_path: str) -> torch.Tensor | None: return _load_tensor(dir_path, "logits.pt") +def _load_manifest(dir_path: str) -> dict | None: + """Load ``parallel_info.json`` written by the dump layer (topology + scopes). + + Returns None when absent (single-card or pre-manifest dumps) -> callers fall + back to tp_size==1 behavior (plain filenames, single-card compatible). + """ + fp = os.path.join(dir_path, "parallel_info.json") + if not os.path.exists(fp): + return None + try: + with open(fp, encoding="utf-8") as f: + return json.load(f) + except Exception: + return None + + def _load_packed_meta(dir_path: str, cu_fname: str = "cu_seqlens_q.pt") -> dict | None: """加载 cu_seqlens + prefix_lens（suffix 对齐所需）。""" @@ -1343,7 +1359,7 @@ def _print_rope_postqk_per_layer(r: CheckResult): bad.append(layer_idx) if bad: print(f"\n ⚠ First deviating layer: {bad[0]}") - print(f" （Q/K max_diff 与 build_kv_input_v 的 V max_diff 同口径，可直接对比）") + print(" （Q/K max_diff 与 build_kv_input_v 的 V max_diff 同口径，可直接对比）") elif "Q_cos_avg" in r.metrics: d = r.metrics ok = (d["Q_cos_avg"] > _COS_AVG_PASS and d["Q_cos_min"] > _COS_MIN_PASS diff --git a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py index 6a5f8104..2fdf0243 100644 --- a/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py +++ b/prefix-sharing/prefix_sharing/tools/diagnostic_dump_verl080.py @@ -33,7 +33,6 @@ _save_tensor, _stage_last_layer, _pp_suffix, - _cached_parallel_info, ) diff --git a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py index f6bb22e6..9c9f335a 100644 --- a/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py +++ b/prefix-sharing/prefix_sharing/tools/inject_baseline_synthetic.py @@ -12,8 +12,6 @@ # Multi-copy: stack=3 → 12 sequences (4 shuffled × 3 stacked) """ -import json -import os import random import torch From c6b0ecdd4fc1cc8d2d27d827a7a4bc0f92e28118 Mon Sep 17 00:00:00 2001 From: Boundless Date: Thu, 16 Jul 2026 16:22:19 +0800 Subject: [PATCH 61/61] [cleanup] remove deprecated flash_atten_npu_tnd and flash_atten_npu_test backends Delete flash_atten_npu_tnd.py and flash_atten_npu_test.py, and clean up all references in __init__.py, factory.py, config.py, and tests. Co-Authored-By: Claude --- .../prefix_sharing/backends/__init__.py | 2 - .../prefix_sharing/backends/factory.py | 7 +- .../backends/flash_atten_npu_test.py | 468 ------------------ .../backends/flash_atten_npu_tnd.py | 257 ---------- prefix-sharing/prefix_sharing/core/config.py | 6 +- .../tests/unit_test/test_backend_factory.py | 11 +- 6 files changed, 5 insertions(+), 746 deletions(-) delete mode 100644 prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py delete mode 100644 prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py diff --git a/prefix-sharing/prefix_sharing/backends/__init__.py b/prefix-sharing/prefix_sharing/backends/__init__.py index 63cc533c..0996d446 100644 --- a/prefix-sharing/prefix_sharing/backends/__init__.py +++ b/prefix-sharing/prefix_sharing/backends/__init__.py @@ -6,7 +6,6 @@ from prefix_sharing.backends.flash_atten_base import FlashAttentionMixin, FlashBackendValidationError from prefix_sharing.backends.flash_atten_gpu import GpuFlashAttentionBackend from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend -from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend from prefix_sharing.backends.torch_ref import TorchReferenceBackend __all__ = [ @@ -15,7 +14,6 @@ "FlashBackendValidationError", "GpuFlashAttentionBackend", "NpuFlashAttentionBackend", - "NpuFlashAttentionTndBackend", "PrefixAttentionBackend", "PrefixDeltanetBackend", "TorchReferenceBackend", diff --git a/prefix-sharing/prefix_sharing/backends/factory.py b/prefix-sharing/prefix_sharing/backends/factory.py index 2359a73b..c761d3c2 100644 --- a/prefix-sharing/prefix_sharing/backends/factory.py +++ b/prefix-sharing/prefix_sharing/backends/factory.py @@ -17,7 +17,6 @@ def get_backend_instance( * ``"torch_ref"`` -> :class:`~prefix_sharing.backends.torch_ref.TorchReferenceBackend` * ``"flash_atten_gpu"`` -> :class:`~prefix_sharing.backends.flash_atten_gpu.GpuFlashAttentionBackend` * ``"flash_atten_npu"`` -> :class:`~prefix_sharing.backends.flash_atten_npu.NpuFlashAttentionBackend` - * ``"flash_atten_npu_tnd"`` -> :class:`~prefix_sharing.backends.flash_atten_npu_tnd.NpuFlashAttentionTndBackend` (recommended for NPU) """ if backend is not None: return backend @@ -34,11 +33,7 @@ def get_backend_instance( from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend return NpuFlashAttentionBackend() - if config.backend == "flash_atten_npu_tnd": - from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend - return NpuFlashAttentionTndBackend() - raise ValueError( f"Unknown backend '{config.backend}'. " - f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu, flash_atten_npu_tnd" + f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu" ) diff --git a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py deleted file mode 100644 index 1c09d708..00000000 --- a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_test.py +++ /dev/null @@ -1,468 +0,0 @@ -"""TND (varlen) + sparse_mode=3 (rightDownCausal) 可行性实验（NPU only）。 - -背景 ----- -当前 ON 路径 ([flash_atten_npu.py](prefix_sharing/backends/flash_atten_npu.py)) 用 -BSH + per-sample B1SS mask + sparse_mode=1，在 reuser 上结果错；改 sparse_mode=3 又崩 -（"attenmask compression requires [2048,2048]"，因为 mode 3 要的是压缩 [2048,2048]， -不是 B1SS）。 - -调研 Ascend 60RC2 官方文档后发现：**sparse_mode=3 (rightDownCausal) 原生支持 Q Any: - """构造一个结构可控的 PrefixSharingPlan。 - - - batch_sizes[i] = 样本 i 的原始长度 - - prefix_lens[i] = 样本 i 的共享前缀长度（0 表示 provider / 无共享） - """ - config = PrefixSharingConfig(enable_prefix_sharing=True, backend="flash_atten_npu") - planner = PrefixSharingPlanner(config) - input_ids = [list(range(s)) for s in batch_sizes] - plan = planner.plan(input_ids) - - object.__setattr__(plan, "batch_size", len(batch_sizes)) - object.__setattr__(plan, "original_lengths", batch_sizes) - object.__setattr__(plan, "prefix_lens", prefix_lens) - object.__setattr__(plan, "kept_lengths_q", [b - p for b, p in zip(batch_sizes, prefix_lens)]) - object.__setattr__(plan, "expanded_lengths_kv", list(batch_sizes)) - object.__setattr__(plan, "q_position_offsets", prefix_lens) - object.__setattr__(plan, "kv_position_offsets", [0] * len(batch_sizes)) - - cu_seqlens_q = [0] - cu_seqlens_kv = [0] - max_seqlen_q = 0 - max_seqlen_kv = 0 - for b, p in zip(batch_sizes, prefix_lens): - q_len = b - p - kv_len = b - cu_seqlens_q.append(cu_seqlens_q[-1] + q_len) - cu_seqlens_kv.append(cu_seqlens_kv[-1] + kv_len) - max_seqlen_q = max(max_seqlen_q, q_len) - max_seqlen_kv = max(max_seqlen_kv, kv_len) - - object.__setattr__(plan, "cu_seqlens_q", cu_seqlens_q) - object.__setattr__(plan, "cu_seqlens_kv", cu_seqlens_kv) - object.__setattr__(plan, "max_seqlen_q", max_seqlen_q) - object.__setattr__(plan, "max_seqlen_kv", max_seqlen_kv) - object.__setattr__(plan, "provider_index", [0] * len(batch_sizes)) - object.__setattr__(plan, "is_provider", [p == 0 for p in prefix_lens]) - object.__setattr__(plan, "reuse_specs", ()) - object.__setattr__(plan, "prefix_last_restore", []) - return plan - - -def _make_layout(kept_lengths_q: list[int]) -> PackedBatchLayout: - return PackedBatchLayout.from_valid_lengths(kept_lengths_q) - - -def _random_qkv(total_q: int, total_kv: int, seed: int = 42): - """随机生成 THD Q / 原始(未展开) K/V。原始 K/V 长度 = kept_lengths_q 之和。""" - torch.manual_seed(seed) - q = torch.randn(total_q, NUM_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 - k = torch.randn(total_kv, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 - v = torch.randn(total_kv, NUM_KV_HEADS, HEAD_DIM, dtype=torch.float16, device=DEVICE) * 0.02 - return q, k, v - - -# ═══════════════════════════════════════════════════════════════════════ -# 本实验新增 helper -# ═══════════════════════════════════════════════════════════════════════ -_COMPRESSED_MASK: torch.Tensor | None = None - - -def _compressed_causal_mask() -> torch.Tensor: - """压缩 [2048,2048] 下三角 mask（True=masked），与 baseline get_attention_mask 一致。 - - sparse_mode 2/3/4 共用这一张压缩 mask；区别只在 kernel 锚定方式。 - 构建一次缓存。理论上覆盖任意 seq 长（压缩模式由 actual_seq 重建每段 causal）。 - """ - global _COMPRESSED_MASK - if _COMPRESSED_MASK is None or _COMPRESSED_MASK.device != DEVICE: - _COMPRESSED_MASK = torch.triu( - torch.ones([2048, 2048], dtype=torch.bool, device=DEVICE), diagonal=1 - ) - return _COMPRESSED_MASK - - -def _build_inputs(plan: Any, seed: int = 42): - """构造一次实验所需的全部张量 + ground truth。 - - 返回: - q: [总Q, NUM_HEADS, HEAD_DIM] —— reuser suffix-only - ek, ev: [总KV展开, NUM_KV_HEADS, HEAD_DIM] —— build_kv 展开后的 K/V（reuser 是 prefix+suffix） - ref_out: torch_ref 在相同 (q, ek, ev, plan) 上的输出，作为 ground truth - layout: PackedBatchLayout - """ - layout = _make_layout(plan.kept_lengths_q) - total_q = int(sum(plan.kept_lengths_q)) - # 原始(未展开) K/V 长度 = kept_lengths_q 之和（裁剪后 provider 全长、reuser suffix） - q, k_raw, v_raw = _random_qkv(total_q, total_q, seed=seed) - - store = PrefixAttentionStore() - ref = TorchReferenceBackend() - ek, ev = ref.build_kv( - k_raw, v_raw, store, plan, - packed_batch_layout=layout, layer_id=0, tp_rank=0, - ) - ref_out = ref.attention(q, ek, ev, plan, packed_batch_layout=layout) - store.close() - return q, ek, ev, ref_out, layout - - -def _varlen_tnd_call( - q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, - *, - cu_q: list[int], - cu_kv: list[int], - sparse_mode: int, -) -> torch.Tensor: - """TND varlen npu_fusion_attention 调用。 - - - input_layout="TND" - - actual_seq_qlen/kvlen 用 cu_seqlens（**带前导 0**，对齐 verl util.py:69 的 zeros(batch+1)） - - atten_mask = 压缩 [2048,2048] 下三角 - - scale / keep_prob = 1/√d / 1.0 - - mode 2/3 下 pre/next_tokens 不生效，走默认 - """ - fn = _import_npu_fusion_attention() - result = fn( - q, k, v, - NUM_HEADS, - "TND", - atten_mask=_compressed_causal_mask(), - scale=SCALE, - keep_prob=1.0, - sparse_mode=sparse_mode, - actual_seq_qlen=list(cu_q), - actual_seq_kvlen=list(cu_kv), - ) - return result[0] if isinstance(result, (tuple, list)) else result - - -def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float: - return float((a.float() - b.float()).abs().max().item()) - - -# ═══════════════════════════════════════════════════════════════════════ -# Probe 矩阵。每个返回 dict: {name, pass, detail, out_diff, grad_diff} -# ═══════════════════════════════════════════════════════════════════════ -def probe_a_provider_only_mode2() -> dict: - """A. provider-only + mode 2（baseline sanity）。 - - 两个 provider，无共享，Q==KV。验证 TND varlen + 压缩 mask + mode 2 的接线正确 - （与 baseline 同款），并与 torch_ref 对齐。A 不过 = 环境/接线坏，后续 probe 都不可信。 - """ - plan = _make_plan(batch_sizes=[4, 6], prefix_lens=[0, 0]) - q, ek, ev, ref_out, _ = _build_inputs(plan) - try: - out = _varlen_tnd_call( - q, ek, ev, - cu_q=plan.cu_seqlens_q, cu_kv=plan.cu_seqlens_kv, - sparse_mode=2, - ) - diff = _max_abs_diff(out, ref_out) - ok = diff < _ATOL_FP16 - return {"name": "A", "pass": ok, "detail": f"provider-only mode2 out_diff={diff:.4e}", - "out_diff": diff, "grad_diff": None} - except Exception as e: # noqa: BLE001 - return {"name": "A", "pass": False, "detail": f"exception: {type(e).__name__}: {e}", - "out_diff": None, "grad_diff": None} - - -def probe_b_single_reuser_mode3() -> dict: - """B. 单 reuser + mode 3（核心 probe）。 - - 从 [8,8]/[0,4] 的 batch 里取 reuser(index=1)：Q=suffix(4)、KV=prefix+suffix(8)、 - prefix_len=4。单独喂给 TND varlen + mode 3，与 torch_ref 的 reuser 段输出对齐。 - - B 过 = mode 3 在 varlen 对 Q dict: - """C. 全 batch（provider+reuser）mode 3 单次调用。 - - 一次 npu_fusion_attention 覆盖 provider 段（Q==KV）和 reuser 段（Q dict: - """D. C 的反向（128-tile 约束）。 - - 全 batch mode 3 forward + sum().backward()，对齐 torch_ref 的 Q/K/V 梯度。 - D 过 = varlen 反向通，128-tile 约束不挡路（或被满足）。 - D 崩 tiling → 需要把 max_q/max_kv 补到 128 倍数重试（生产改写时处理）。 - """ - plan = _make_plan(batch_sizes=[8, 8], prefix_lens=[0, 4]) - - def _run(forward_fn): - q, ek, ev, _, _ = _build_inputs(plan) - q = q.clone().detach().requires_grad_(True) - ek = ek.clone().detach().requires_grad_(True) - ev = ev.clone().detach().requires_grad_(True) - out = forward_fn(q, ek, ev) - out.sum().backward() - return {"q": q.grad, "k": ek.grad, "v": ev.grad} - - # FA 反向 - def fa_fwd(qq, kk, vv): - return _varlen_tnd_call(qq, kk, vv, cu_q=plan.cu_seqlens_q, - cu_kv=plan.cu_seqlens_kv, sparse_mode=3) - # ref 反向（torch_ref） - layout = _make_layout(plan.kept_lengths_q) - ref_backend = TorchReferenceBackend() - - def ref_fwd(qq, kk, vv): - return ref_backend.attention(qq, kk, vv, plan, packed_batch_layout=layout) - - try: - grads_fa = _run(fa_fwd) - grads_ref = _run(ref_fwd) - max_grad_diff = 0.0 - for name in ("q", "k", "v"): - d = _max_abs_diff(grads_fa[name], grads_ref[name]) - max_grad_diff = max(max_grad_diff, d) - ok = max_grad_diff < _ATOL_GRAD_FP16 - return {"name": "D", "pass": ok, - "detail": f"full-batch mode3 backward grad_diff={max_grad_diff:.4e}", - "out_diff": None, "grad_diff": max_grad_diff} - except Exception as e: # noqa: BLE001 - return {"name": "D", "pass": False, - "detail": f"backward exception (可能 128-tile): {type(e).__name__}: {e}", - "out_diff": None, "grad_diff": None} - - -def probe_e_long_seq_gt_2048() -> dict: - """E.（可选）mode 3 + 某段 seq>2048。 - - 把 provider 段拉到 >2048，确认压缩 [2048,2048] mask 不限 seq 长（压缩模式由 - actual_seq 重建每段 causal，理论上支持任意 seq）。这是用户关心的点。 - """ - plan = _make_plan(batch_sizes=[2100, 2100], prefix_lens=[0, 100]) - q, ek, ev, ref_out, _ = _build_inputs(plan) - try: - out = _varlen_tnd_call( - q, ek, ev, - cu_q=plan.cu_seqlens_q, cu_kv=plan.cu_seqlens_kv, - sparse_mode=3, - ) - diff = _max_abs_diff(out, ref_out) - ok = diff < _ATOL_FP16 - return {"name": "E", "pass": ok, - "detail": f"seq>2048 mode3 out_diff={diff:.4e}", - "out_diff": diff, "grad_diff": None} - except Exception as e: # noqa: BLE001 - return {"name": "E", "pass": False, "detail": f"exception: {type(e).__name__}: {e}", - "out_diff": None, "grad_diff": None} - - -_PROBES = [ - probe_a_provider_only_mode2, - probe_b_single_reuser_mode3, - probe_c_full_batch_mode3, - probe_d_full_batch_backward, - probe_e_long_seq_gt_2048, -] - - -# ═══════════════════════════════════════════════════════════════════════ -# pytest 包装（每个 probe 一个 test，非 NPU skip） -# ═══════════════════════════════════════════════════════════════════════ -@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") -def test_probe_a_provider_only_mode2(): - r = probe_a_provider_only_mode2() - assert r["pass"], r["detail"] - - -@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") -def test_probe_b_single_reuser_mode3(): - r = probe_b_single_reuser_mode3() - assert r["pass"], r["detail"] - - -@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") -def test_probe_c_full_batch_mode3(): - r = probe_c_full_batch_mode3() - assert r["pass"], r["detail"] - - -@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") -def test_probe_d_full_batch_backward(): - r = probe_d_full_batch_backward() - assert r["pass"], r["detail"] - - -@pytest.mark.skipif(not _HAS_NPU, reason="requires NPU device + mindspeed") -def test_probe_e_long_seq_gt_2048(): - r = probe_e_long_seq_gt_2048() - assert r["pass"], r["detail"] - - -# ═══════════════════════════════════════════════════════════════════════ -# 决策树（__main__ 用） -# ═══════════════════════════════════════════════════════════════════════ -def _decide(results: dict[str, dict]) -> str: - a, b, c, d = results["A"], results["B"], results["C"], results["D"] - if not a["pass"]: - return ("A 失败 → 环境/接线坏（input_layout 字符串/cu_seqlens 格式/scale 等）。" - "先修 A，后续 probe 都不可信。") - if not b["pass"]: - return ("B 失败 → mode 3 在本 CANN 版本的 varlen 下没对 Q None: - if not _HAS_NPU: - print("[skip] 无 NPU 设备或 mindspeed 内核，本实验只能在 NPU 上跑。") - return - print(f"[env] DEVICE={DEVICE} NUM_HEADS={NUM_HEADS} NUM_KV_HEADS={NUM_KV_HEADS} " - f"HEAD_DIM={HEAD_DIM} SCALE={SCALE:.4f}") - print("=" * 70) - results: dict[str, dict] = {} - for probe in _PROBES: - r = probe() - results[r["name"]] = r - tag = "PASS" if r["pass"] else "FAIL" - diff_str = [] - if r["out_diff"] is not None: - diff_str.append(f"out_diff={r['out_diff']:.4e}") - if r["grad_diff"] is not None: - diff_str.append(f"grad_diff={r['grad_diff']:.4e}") - diff_txt = f" [{', '.join(diff_str)}]" if diff_str else "" - print(f"Probe {r['name']}: {tag}{diff_txt}") - print(f" {r['detail']}") - print("=" * 70) - print("[决策] " + _decide(results)) - - -if __name__ == "__main__": - _run_all() diff --git a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py b/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py deleted file mode 100644 index 29abb8fb..00000000 --- a/prefix-sharing/prefix_sharing/backends/flash_atten_npu_tnd.py +++ /dev/null @@ -1,257 +0,0 @@ -"""Ascend NPU Flash Attention backend — TND varlen + sparse_mode=3 (rightDownCausal). - -这是**推荐的 NPU 后端**，与 OFF baseline（MindSpeed ``dot_product_attention``）使用 -相同的 TND varlen 约定，配合 ``sparse_mode=3``。老的 BSH 后端 -(:mod:`prefix_sharing.backends.flash_atten_npu`) 保留作对照/回退，不要删。 - -为什么用 sparse_mode=3（核心） -------------------------------- -``sparse_mode=3`` 即 **rightDownCausal**："以右下顶点划分的下三角"。对一个 Q 比 -KV 短的 segment（Sq < Skv），它把 Q **右对齐到 KV 末端**：query i（局部）见 kv j -iff ``j <= (Skv - Sq) + i``。 - -代到 prefix-sharing 的 reuser（Q=suffix、KV=prefix+suffix、Skv-Sq=prefix_len）： -query i 见 kv j iff ``j <= prefix_len + i`` = **全部 prefix KV 可见 + suffix KV causal**， -正是 reuser 的正确语义。而对 provider（Sq==Skv）退化成 ``j <= i`` = 标准 causal， -与 baseline（``sparse_mode=2`` leftUpCausal）一致。 - -因此**整个 batch（providers + reusers）一次 varlen 调用即可**：kernel 按每个 segment -的 actual_seq 自动判 provider（标准 causal）还是 reuser（右对齐 causal）。reuser 的 -Q 仍是 suffix-only（省算力核心收益不变），mask 用现成的压缩 ``[2048,2048]`` 下三角 -（与 baseline ``get_attention_mask`` 完全一样），**无需自建 mask、无需按 prefix_len -分组拆调用**。 - -为什么不用老 backend 的 BSH + sparse_mode=1 -------------------------------------------- -老 backend 用 BSH + per-sample B1SS mask + ``sparse_mode=1``，实测在 reuser 上结果 -错（mode 1=allMask 对 B1SS 自建 mask 的处理不对）。改 ``sparse_mode=3`` 又崩 -（"attenmask compression requires [2048,2048]"，因为 mode 3 要压缩 ``[2048,2048]``， -不是 B1SS）。TND varlen + 压缩 mask 才是 mode 3 的正确用法。 - -曾经担心 varlen 反向有 128-tile 约束（见老 backend 的 docstring），实测 -（``flash_atten_npu_test.py`` Probe D）已证伪——本 CANN 版本 varlen 反向正常。 - -选用方式 --------- -配置里设 ``backend="flash_atten_npu_tnd"`` 即可指向本后端。 - -Select via config: ``backend="flash_atten_npu_tnd"``. -""" - -from __future__ import annotations - -import importlib -import math -from functools import lru_cache -from typing import Any - -from prefix_sharing.backends.base import BackendCapabilities -from prefix_sharing.backends.flash_atten_base import ( - FlashAttentionMixin, - FlashBackendValidationError, -) -from prefix_sharing.backends.torch_ref import TorchReferenceBackend -from prefix_sharing.core.config import PrefixSharingConfig -from prefix_sharing.core.planner import PrefixSharingPlan - - -_CANDIDATES = [ - ("mindspeed.ops.fusion_attention_v2", "npu_fusion_attention"), - ("mindspeed.ops", "npu_fusion_attention"), -] - - -@lru_cache(maxsize=None) -def _import_npu_fusion_attention(): - last_err = None - for module_name, attr in _CANDIDATES: - try: - module = importlib.import_module(module_name) - return getattr(module, attr) - except ImportError as e: - last_err = e - raise RuntimeError( - "NpuFlashAttentionTndBackend requires MindSpeed (mindspeed.ops). " - "Install MindSpeed matching your CANN version." - ) from last_err - - -def _torch() -> Any: - try: - import torch - except ModuleNotFoundError as exc: - raise RuntimeError("NpuFlashAttentionTndBackend requires PyTorch") from exc - return torch - - -# 压缩 [2048,2048] 下三角 mask（True=masked）按 device 缓存。 -# sparse_mode 2/3/4 共用这张压缩 mask；kernel 拿 actual_seq 重建每段 causal， -# 故不限 seq 长（实测 seq>2048 正常，见 flash_atten_npu_test.py Probe E）。 -_COMPRESSED_MASK: dict[Any, Any] = {} - - -def _compressed_causal_mask(device: Any) -> Any: - torch = _torch() - cached = _COMPRESSED_MASK.get(device) - if cached is None: - cached = torch.triu( - torch.ones([2048, 2048], dtype=torch.bool, device=device), diagonal=1 - ) - _COMPRESSED_MASK[device] = cached - return cached - - -class NpuFlashAttentionTndBackend(FlashAttentionMixin): - """Ascend NPU backend via ``npu_fusion_attention`` (TND varlen, sparse_mode=3). - - 单次 varlen 调用覆盖整个 batch（providers + reusers）： - - - **Provider 段**（Sq==Skv）：``sparse_mode=3`` 退化成标准 causal。 - - **Reuser 段**（Sq None: - self._torch_ref = TorchReferenceBackend() - - def validate(self, config: PrefixSharingConfig, model_config: Any | None = None) -> None: - config.validate(model_config=model_config) - _import_npu_fusion_attention() - - def apply_rope( - self, - query: Any, - key: Any, - prefix_sharing_plan: PrefixSharingPlan, - **kwargs: Any, - ) -> tuple[Any, Any]: - return self._torch_ref.apply_rope(query, key, prefix_sharing_plan, **kwargs) - - def build_kv( - self, - key: Any, - value: Any, - store: Any, - prefix_sharing_plan: PrefixSharingPlan, - *, - packed_batch_layout: Any | None = None, - layer_id: int, - tp_rank: int = 0, - stats: Any | None = None, - ) -> tuple[Any, Any]: - """KV 展开委托给 torch reference（与其它后端一致）。""" - return self._torch_ref.build_kv( - key, - value, - store, - prefix_sharing_plan, - packed_batch_layout=packed_batch_layout, - layer_id=layer_id, - tp_rank=tp_rank, - stats=stats, - ) - - # ------------------------------------------------------------------ - # attention — TND varlen + sparse_mode=3，单次调用全 batch - # ------------------------------------------------------------------ - def attention( - self, - query: Any, - key: Any, - value: Any, - prefix_sharing_plan: PrefixSharingPlan, - **kwargs: Any, - ) -> Any: - """Run prefix-sharing attention via TND-varlen ``npu_fusion_attention``. - - 流程： - 1. ``_prepare_flash_inputs`` 剥 Q 的 TP padding、产出 cu_seqlens_q/kv - （batch+1，带前导 0，取自 plan）。 - 2. 一次 ``npu_fusion_attention``，``input_layout="TND"``、 - ``sparse_mode=3``（rightDownCausal）、压缩 ``[2048,2048]`` mask。 - 3. 必要时把输出按 pad_layout 回填 TP padding，恢复原 Q 形状。 - - K/V 来自 ``build_kv``，已是展开后的 TND（reuser 是 prefix+suffix）， - 无 padding，跟随 ``plan.expanded_lengths_kv``。 - """ - layer_id = kwargs.get("layer_id", "?") - packed_batch_layout = kwargs.get("packed_batch_layout") - if packed_batch_layout is None: - raise FlashBackendValidationError( - "flash_atten_npu_tnd.attention requires packed_batch_layout kwarg." - ) - - print( - f"[PS][backend] flash_atten_npu_tnd attention: " - f"layer={layer_id}, " - f"q_shape={tuple(query.shape)}, k_shape={tuple(key.shape)}, " - f"v_shape={tuple(value.shape)}" - ) - - npu_fusion_attention = _import_npu_fusion_attention() - - # Step 1: 剥 Q 的 TP padding + 取 cu_seqlens_q/kv（plan 语义长度，带前导 0）。 - # _prepare_flash_inputs 返回 8 元组：q, k, v, cu_seqlens_q, cu_seqlens_kv, - # max_seqlen_q, max_seqlen_kv, pad_layout（max_seqlen_* 这里不用）。 - q, k, v, cu_seqlens_q, cu_seqlens_kv, _, _, pad_layout = ( - self._prepare_flash_inputs( - query, - key, - value, - prefix_sharing_plan, - attention_mask=kwargs.get("attention_mask"), - packed_batch_layout=packed_batch_layout, - ) - ) - - num_q_heads = q.shape[1] - head_dim = q.shape[-1] - scale = kwargs.get("softmax_scale") or (1.0 / math.sqrt(head_dim)) - dropout_p = kwargs.get("dropout_p", 0.0) - keep_prob = kwargs.get("keep_prob", 1.0 - dropout_p) - - # Step 2: 单次 varlen 调用，sparse_mode=3（rightDownCausal）。 - # mode 3 下 pre/next_tokens 不生效，走默认；atten_mask 用压缩 [2048,2048]。 - try: - result = npu_fusion_attention( - q, - k, - v, - num_q_heads, - "TND", - atten_mask=_compressed_causal_mask(q.device), - scale=scale, - keep_prob=keep_prob, - sparse_mode=3, - actual_seq_qlen=cu_seqlens_q.tolist(), - actual_seq_kvlen=cu_seqlens_kv.tolist(), - ) - except Exception as exc: - raise FlashBackendValidationError( - f"npu_fusion_attention (TND, sparse_mode=3) failed: " - f"q={tuple(q.shape)}, k={tuple(k.shape)}, v={tuple(v.shape)}, " - f"cu_seqlens_q={cu_seqlens_q.tolist()}, " - f"cu_seqlens_kv={cu_seqlens_kv.tolist()}" - ) from exc - - output = result[0] if isinstance(result, (tuple, list)) else result - - # Step 3: 回填 TP padding，恢复原 Q 形状（TP=1 时 pad_layout 为 None，no-op）。 - if pad_layout is not None: - output = self._repad_output(output, pad_layout) - - return output diff --git a/prefix-sharing/prefix_sharing/core/config.py b/prefix-sharing/prefix_sharing/core/config.py index 6f74959b..7aee42f8 100644 --- a/prefix-sharing/prefix_sharing/core/config.py +++ b/prefix-sharing/prefix_sharing/core/config.py @@ -114,7 +114,7 @@ def validate(self, model_config: Any | None = None, integrate_mode: str | None = return if self.detector != "trie": raise PrefixSharingConfigError("phase 1 supports only detector='trie'") - supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"} + supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu"} if self.backend not in supported_backends: raise PrefixSharingConfigError( f"backend='{self.backend}' is not supported. " @@ -217,10 +217,10 @@ def validate_for_engine( # 基础校验 if self.detector != "trie": raise PrefixSharingConfigError("phase 1 supports only detector='trie'") - if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"}: + if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu"}: raise PrefixSharingConfigError( f"backend='{self.backend}' is not supported. " - f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu, flash_atten_npu_tnd" + f"Supported: torch_ref, flash_atten_gpu, flash_atten_npu" ) if self.boundary_strategy != "prefix_last_restore": raise PrefixSharingConfigError( diff --git a/prefix-sharing/tests/unit_test/test_backend_factory.py b/prefix-sharing/tests/unit_test/test_backend_factory.py index 455d6317..1ef5128d 100644 --- a/prefix-sharing/tests/unit_test/test_backend_factory.py +++ b/prefix-sharing/tests/unit_test/test_backend_factory.py @@ -7,7 +7,6 @@ from prefix_sharing.backends.factory import get_backend_instance from prefix_sharing.backends.flash_atten_gpu import GpuFlashAttentionBackend from prefix_sharing.backends.flash_atten_npu import NpuFlashAttentionBackend -from prefix_sharing.backends.flash_atten_npu_tnd import NpuFlashAttentionTndBackend from prefix_sharing.backends.torch_ref import TorchReferenceBackend from prefix_sharing.core.config import PrefixSharingConfig @@ -34,14 +33,6 @@ def test_factory_flash_atten_npu(): assert backend.capabilities.name == "flash_atten_npu" -def test_factory_flash_atten_npu_tnd(): - config = PrefixSharingConfig(enable_prefix_sharing=True, backend="flash_atten_npu_tnd") - backend = get_backend_instance(config) - assert isinstance(backend, NpuFlashAttentionTndBackend) - assert backend.capabilities.name == "flash_atten_npu_tnd" - assert backend.capabilities.supports_cann - - def test_factory_unknown_backend() -> None: config = PrefixSharingConfig(enable_prefix_sharing=True, backend="unknown") with pytest.raises(ValueError, match="Unknown backend"): @@ -71,7 +62,7 @@ def test_config_validates_backends() -> None: def test_config_accepts_supported_backends(): - for name in ("torch_ref", "flash_atten_gpu", "flash_atten_npu", "flash_atten_npu_tnd"): + for name in ("torch_ref", "flash_atten_gpu", "flash_atten_npu"): cfg = PrefixSharingConfig(enable_prefix_sharing=True, backend=name) cfg.validate() # should not raise \ No newline at end of file