File size: 4,149 Bytes
c92e699 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 | # SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
#
# SPDX-License-Identifier: Apache-2.0
from pathlib import Path
from types import SimpleNamespace
import torch
from models.tt_transformers.tt import attention as attention_module
class _StopAfterBiasConstruction(Exception):
pass
class _FakeDecoderOptimizations:
def get_tensor_dtype(self, **_):
return None
def get_math_fidelity(self, **_):
return None
def test_galaxy_qkv_bias_uses_weight_matching_2d_sharding_and_survives_column_reduce(monkeypatch, expect_error):
mesh_device = object()
mesh_mapper_calls = []
as_tensor_calls = []
monkeypatch.setattr(attention_module.ttnn, "ShardTensorToMesh", lambda *args, **kwargs: object())
monkeypatch.setattr(attention_module.ttnn, "ReplicateTensorToMesh", lambda *args, **kwargs: object())
monkeypatch.setattr(attention_module.ttnn, "from_torch", lambda *args, **kwargs: object())
def shard_tensor_2d_mesh(device, *, dims, mesh_shape):
mapper = object()
mesh_mapper_calls.append((device, dims, mesh_shape, mapper))
return mapper
def as_tensor(tensor, **kwargs):
as_tensor_calls.append((tensor, kwargs))
return SimpleNamespace(shape=(1, tensor.shape[-1]))
monkeypatch.setattr(attention_module.ttnn, "ShardTensor2dMesh", shard_tensor_2d_mesh)
monkeypatch.setattr(attention_module.ttnn, "as_tensor", as_tensor)
monkeypatch.setattr(
attention_module.ttnn,
"reshape",
lambda *args, **kwargs: (_ for _ in ()).throw(_StopAfterBiasConstruction("bias captured")),
)
configuration = SimpleNamespace(
num_devices=32,
dim=8192,
n_heads=64,
head_dim=128,
max_seq_len=32768,
max_batch_size=32,
n_kv_heads=8,
min_kv_prefill_shard_seqlen=0,
ccl_dtype=None,
MAX_QKV_MM_SEQ_LEN=2048,
tile_size=32,
tile_padded_batch_rows=32,
rms_norm_add_unit_offset=False,
use_qk_fused=False,
use_hf_rope=False,
arch_name="wormhole_b0",
max_grid_size=None,
compute_kernel_config_hifi2=None,
compute_kernel_config_hifi2_fp16=None,
compute_kernel_config_hifi4=None,
layer_types=None,
cluster_shape=[8, 4],
is_multichip=True,
dummy_weights=False,
qkv_size=10240,
get_model_config=lambda: {},
ccl_topology=lambda: None,
get_state_dict_prefix=lambda *_: "layers.0.attention",
)
args = SimpleNamespace(decoders_optimizations=_FakeDecoderOptimizations())
state_dict = {
"layers.0.attention.wq.bias": torch.arange(16),
"layers.0.attention.wk.bias": torch.arange(100, 108),
"layers.0.attention.wv.bias": torch.arange(200, 208),
}
with expect_error(_StopAfterBiasConstruction, "bias captured"):
attention_module.Attention(
mesh_device=mesh_device,
tt_ccl=None,
args=args,
state_dict=state_dict,
weight_cache_path=Path("cache"),
layer_num=0,
dtype=None,
transformation_mats=None,
configuration=configuration,
)
assert len(mesh_mapper_calls) == 1
assert mesh_mapper_calls[0][:3] == (mesh_device, (-1, None), [8, 4])
qkv_bias, tensor_kwargs = as_tensor_calls[0]
assert tensor_kwargs["mesh_mapper"] is mesh_mapper_calls[0][3]
assert tensor_kwargs["cache_file_name"].name.endswith("wqkv_bias_prefill_sharded_2d_col_reduce_4")
# Each row shard is replicated over all four mesh columns and added before
# the column all-reduce. Their sum must restore one copy of the model bias,
# rather than the four copies produced by an unscaled replicated bias.
expected_device_shards = [
torch.tensor([2 * i, 2 * i + 1, 100 + i, 200 + i], dtype=qkv_bias.dtype) for i in range(8)
]
reduced_device_shards = [shard * configuration.cluster_shape[1] for shard in torch.chunk(qkv_bias, 8)]
assert all(torch.equal(actual, expected) for actual, expected in zip(reduced_device_shards, expected_device_shards))
|