📊 feat: distinguish compute density and communication density in DensityObserver

- Add record_comm_density() call in select_blocks to track CPU block selection
- Add get_per_layer_comm_density() method for detailed analysis
- Update print_summary() to show both densities and H2D savings ratio
- Set DensityObserver mode (offload/gpu_only) in test_ruler.py
- Update get_summary() to return both density types

Key insight: Comm density can be 100% even when compute density is ~37%
because sparse BSA blocks are distributed across all CPU blocks.
Since CPU block granularity is 32x coarser (4096 vs 128 tokens),
any() aggregation across heads/Q-blocks results in all CPU blocks being needed.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
This commit is contained in:
Zijie Tian
2026-02-05 01:43:17 +08:00
parent 1ea5afd886
commit 51bd678335
3 changed files with 35 additions and 4 deletions

View File

@@ -266,14 +266,31 @@ class DensityObserver(Observer):
return 0.0
return sum(all_densities) / len(all_densities)
@classmethod
def get_per_layer_comm_density(cls) -> Dict[int, float]:
"""
获取每层的 communication density (CPU block 粒度)。
Returns:
Dict[layer_id, avg_comm_density]
"""
result = {}
for layer_id, densities in cls._layer_comm_densities.items():
if densities:
result[layer_id] = sum(densities) / len(densities)
return result
@classmethod
def get_summary(cls) -> dict:
"""返回统计摘要"""
per_layer = cls.get_per_layer_density()
per_layer_comm = cls.get_per_layer_comm_density()
return {
"mode": cls._mode,
"overall_density": cls.get_overall_density(),
"per_layer_density": per_layer,
"overall_compute_density": cls.get_overall_density(),
"overall_comm_density": cls.get_overall_comm_density(),
"per_layer_compute_density": per_layer,
"per_layer_comm_density": per_layer_comm,
"num_layers": len(per_layer),
"last_mask_shape": {
"q_blocks": cls._last_q_blocks,
@@ -301,7 +318,9 @@ class DensityObserver(Observer):
print(f"[DensityObserver] Mode: {cls._mode}")
print(f" Compute density: {overall:.4f} (min: {min_density:.4f} @ layer {min_layer})")
if overall_comm > 0:
print(f" Comm density: {overall_comm:.4f}")
# Offload mode: show both densities with explanation
print(f" Comm density: {overall_comm:.4f} (CPU block granularity)")
print(f" Savings ratio: {1 - overall_comm:.1%} H2D transfer reduction")
print(f" Num layers: {len(per_layer)}")
# 输出 layer 0 的 density 用于对比
if 0 in per_layer: