[Bug]: Attention-DP KV metrics copy rank-0 values and mix per-rank counter history #21124

cat > sitecustomize.py <<'PY_OBSERVER'
"""Diagnostic observation only: retain rank-local values before ADP fanout."""
import json
import os
import time
from pathlib import Path

if os.environ.get("ADP_REPRO_DIR"):
    from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor

    original = PyExecutor._update_iter_stats
    counts = {}

    def observed(self, *args, **kwargs):
        stats = original(self, *args, **kwargs)
        rank = int(self.dist.tp_rank)
        count = counts.get(rank, 0)
        if stats is not None and count < 10000:
            counts[rank] = count + 1
            row = {
                "rank": rank,
                "time_ns": time.time_ns(),
                "stats": json.loads(stats.to_json_str()),
            }
            root = Path(os.environ["ADP_REPRO_DIR"])
            with (root / f"local-rank-{rank}.jsonl").open("a") as output:
                output.write(json.dumps(row) + "\n")
        return stats

    PyExecutor._update_iter_stats = observed
PY_OBSERVER

cat > repro.py <<'PY_WORKLOAD'
import asyncio
import hashlib
import json
import os
import time
from pathlib import Path

ROOT = Path(os.environ["ADP_REPRO_DIR"])


async def main():
    import tensorrt_llm
    from tensorrt_llm import LLM, SamplingParams
    from tensorrt_llm.llmapi.llm_args import KvCacheConfig
    from tensorrt_llm.scheduling_params import SchedulingParams
    from tensorrt_llm.metrics import MetricsCollector
    from prometheus_client import generate_latest
    import tensorrt_llm._torch.pyexecutor.adp_iter_stats as adp
    import tensorrt_llm.metrics.collector as collector_module

    manifest = {
        "runtime": tensorrt_llm.__version__,
        "model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        "tp_size": 4,
        "attention_dp": True,
        "all_ranks": os.environ.get("TLLM_METRICS_ALL_RANKS", "0"),
        "source_hashes": {},
    }
    for label, module in (("adp_producer", adp), ("collector", collector_module)):
        manifest["source_hashes"][label] = hashlib.sha256(Path(module.__file__).read_bytes()).hexdigest()
    (ROOT / "manifest.json").write_text(json.dumps(manifest, indent=2))
    print("REPRO_MANIFEST", json.dumps(manifest), flush=True)
    llm = LLM(
        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        backend="pytorch",
        dtype="bfloat16",
        tensor_parallel_size=4,
        enable_attention_dp=True,
        enable_chunked_prefill=True,
        enable_iter_perf_stats=True,
        return_perf_metrics=True,
        max_batch_size=8,
        max_num_tokens=2048,
        max_seq_len=2048,
        cuda_graph_config=None,
        kv_cache_config=KvCacheConfig(
            max_tokens=8192,
            enable_block_reuse=True,
            use_kv_cache_manager_v2=True,
        ),
    )
    collector = MetricsCollector({"model_name": "adp-repro", "engine_type": "trtllm"})
    manifest["workload_started_ns"] = time.time_ns()
    (ROOT / "manifest.json").write_text(json.dumps(manifest, indent=2))
    stopped = asyncio.Event()
    observed_ranks = set()
    row_count = 0

    async def poll():
        nonlocal row_count
        with (ROOT / "exported-rows.jsonl").open("w") as output:
            while not stopped.is_set():
                rows = await asyncio.to_thread(llm.get_stats, 0.2)
                for row in rows:
                    observed_ranks.add(int(row.get("attentionDpRank", row.get("rank", 0))))
                    output.write(json.dumps(row) + "\n")
                    row_count += 1
                    collector.log_iteration_stats(row)
                output.flush()
                await asyncio.sleep(0.02)

    async def one(rank, variant, length, output_tokens):
        prompt = [100 + rank * 31 + variant] * length
        handle = llm.generate_async(
            inputs=prompt,
            sampling_params=SamplingParams(max_tokens=output_tokens, temperature=0.0, ignore_eos=True),
            streaming=True,
            scheduling_params=SchedulingParams(attention_dp_rank=rank, attention_dp_relax=False),
            cache_salt=f"rank-{rank}-variant-{variant}",
        )
        terminal = None
        async for result in handle:
            if result.finished:
                terminal = result
        assert terminal is not None, (rank, variant, "missing terminal response")
        assert not terminal.error, str(terminal.error)
        print("REQUEST_COMPLETE", rank, variant, length, output_tokens, flush=True)

    task = asyncio.create_task(poll())
    success = False
    try:
        # Establish real work on every rank, with distinct prompt lengths.
        await asyncio.gather(*(one(rank, 0, 64 * (rank + 1), 32) for rank in range(4)))
        # Keep rank 0 busy while all other ranks are idle.
        await one(0, 1, 640, 128)
        # Different cache footprints and reuse activity on rank 1.
        await one(1, 2, 1024, 64)
        await one(1, 2, 1024, 64)
        # Simultaneously exercise distinct load and retained-cache sizes.
        await asyncio.gather(one(0, 3, 128, 128), one(2, 4, 768, 32))
        await asyncio.sleep(1)
        assert observed_ranks == {0, 1, 2, 3}, observed_ranks
        (ROOT / "metrics.prom").write_bytes(generate_latest())
        success = True
    finally:
        stopped.set()
        try:
            await asyncio.wait_for(task, 15)
        finally:
            await asyncio.to_thread(llm.shutdown)
            summary = {"workload_passed": success, "observed_ranks": sorted(observed_ranks), "rows": row_count}
            (ROOT / "workload-summary.json").write_text(json.dumps(summary, indent=2))
            print("WORKLOAD_SUMMARY", json.dumps(summary), flush=True)


if __name__ == "__main__":
    asyncio.run(main())
PY_WORKLOAD

cat > analyze.py <<'PY_ANALYZE'
"""Require real multi-rank evidence; do not infer correctness from identical rows."""
import json
import sys
from collections import defaultdict
from pathlib import Path

root = Path(sys.argv[1])
manifest = json.loads((root / "manifest.json").read_text())
workload = json.loads((root / "workload-summary.json").read_text())
assert workload["workload_passed"], workload
assert workload["observed_ranks"] == [0, 1, 2, 3], workload
assert manifest["source_hashes"]["adp_producer"] == "816c440e0e95bafb159ad669c8e30ad17cf8c8d33dbe9052f544e64ea06c8efe"

rows = [json.loads(line) for line in (root / "exported-rows.jsonl").read_text().splitlines()]
local = defaultdict(dict)
for rank in range(4):
    path = root / f"local-rank-{rank}.jsonl"
    assert path.exists(), f"Missing actual measurements for rank {rank}"
    for line in path.read_text().splitlines():
        measurement = json.loads(line)
        if measurement["time_ns"] < manifest["workload_started_ns"]:
            continue
        stats = measurement["stats"]
        assert "iter" in stats, stats.keys()
        local[rank][stats["iter"]] = stats
    assert local[rank], f"No workload measurements for rank {rank}"

evidence = []
matched = 0
for row in rows:
    rank = int(row.get("attentionDpRank", row.get("rank", 0)))
    iteration = row.get("iter")
    own = local[rank].get(iteration)
    zero = local[0].get(iteration)
    if rank == 0 or own is None or zero is None:
        continue
    if not all("kvCacheStats" in stats for stats in (row, own, zero)):
        continue
    matched += 1
    own_kv, zero_kv, exported_kv = (stats["kvCacheStats"] for stats in (own, zero, row))
    for field in ("usedNumBlocks", "freeNumBlocks", "reusedBlocks", "missedBlocks"):
        if field not in own_kv or field not in zero_kv or field not in exported_kv:
            continue
        if own_kv[field] != zero_kv[field] and exported_kv[field] == zero_kv[field]:
            evidence.append({
                "iter": iteration,
                "rank": rank,
                "field": field,
                "actual_rank_value": own_kv[field],
                "actual_rank0_value": zero_kv[field],
                "exported_rank_value": exported_kv[field],
            })

result = {
    "workload": workload,
    "matched_nonzero_rank_rows": matched,
    "rank0_copy_confirmed": bool(evidence),
    "mismatched_field_observations": len(evidence),
    "examples": evidence[:20],
    "limits": [
        "Direct LLM statistics plus native collector; not a full Dynamo/HTTP/OpenEngine serving comparison.",
        "Diagnostic observation adds local file writes; no performance claim.",
    ],
}
assert matched > 0, "No aligned local/exported observations; result is inconclusive"
(root / "analysis.json").write_text(json.dumps(result, indent=2))
print(json.dumps(result, indent=2))
PY_ANALYZE

cat > replay_collector.py <<'PY_REPLAY'
"""Replay real rank-local captures through the unmodified native collector.

This isolates collector behavior; it is not a claim that default fanout exports
these truthful local rows today.
"""
import importlib
import importlib.util
import json
import sys
import types
from pathlib import Path

from prometheus_client import REGISTRY

source = Path(importlib.util.find_spec("tensorrt_llm").submodule_search_locations[0]) / "metrics"
package = types.ModuleType("native_metrics_replay")
package.__path__ = [str(source)]
sys.modules[package.__name__] = package
module = importlib.import_module("native_metrics_replay.collector")
root = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/repro/output")
rows = []
for rank in range(4):
    for line in (root / f"local-rank-{rank}.jsonl").read_text().splitlines():
        capture = json.loads(line)
        row = capture["stats"]
        row["attentionDpRank"] = rank
        rows.append((row["iter"], rank, row))
rows.sort(key=lambda entry: (entry[0], entry[1]))
labels = {"model_name": "rank-local-replay", "engine_type": "trtllm"}
collector = module.MetricsCollector(labels)
previous = {"reusedBlocks": {}, "missedBlocks": {}}
reference = {"reusedBlocks": 0, "missedBlocks": 0}
for _, rank, row in rows:
    collector.log_iteration_stats(row)
    for field in reference:
        value = row.get("kvCacheStats", {}).get(field)
        if value is not None:
            reference[field] += max(0, value - previous[field].get(rank, 0))
            previous[field][rank] = value
result = {
    "input": "Real per-rank measurements captured before default ADP fanout",
    "order": "iteration then rank",
    "rows": len(rows),
    "reused_blocks": {
        "collector": REGISTRY.get_sample_value("trtllm_kv_cache_reused_blocks_total", labels),
        "per_rank_reference": reference["reusedBlocks"],
    },
    "missed_blocks": {
        "collector": REGISTRY.get_sample_value("trtllm_kv_cache_missed_blocks_total", labels),
        "per_rank_reference": reference["missedBlocks"],
    },
    "limit": "Controlled replay isolates collector rank mixing, not default live exporter output.",
}
assert result["reused_blocks"]["collector"] is not None
assert result["missed_blocks"]["collector"] is not None
print(json.dumps(result, indent=2))
PY_REPLAY