|
6 | 6 |
|
7 | 7 | import os |
8 | 8 | import time |
| 9 | +from dataclasses import replace |
9 | 10 | from pathlib import Path |
| 11 | +from types import SimpleNamespace |
10 | 12 |
|
11 | 13 | import pytest |
12 | 14 | import torch |
13 | 15 | from b12x.attention import sparse_mla |
14 | 16 | from cuda.bindings import driver as cuda |
15 | 17 |
|
16 | 18 | from ignition.components.parallel import TensorParallelShard |
| 19 | +from ignition.components.linear import BatchedQuantizedProjection, BatchedUnquantizedProjection |
| 20 | +from ignition.components.sparse_mla import PreparedSparseMLA |
17 | 21 | from ignition.preparation.checkpoint import inspect_indexed_checkpoint |
18 | 22 | from ignition.preparation.compilation import compile_forward |
19 | 23 | from ignition.preparation.execution import GraphSignature, PreparedExecution |
@@ -235,3 +239,86 @@ def forward(): |
235 | 239 | if worker is not None: |
236 | 240 | worker.close() |
237 | 241 | graph.reset() |
| 242 | + |
| 243 | + |
| 244 | +@pytest.mark.parametrize("model_weights", [2], indirect=True) |
| 245 | +@torch.no_grad() |
| 246 | +def test_cp_prefill_query_reuse_matches_separate_values_across_slices(model_weights): |
| 247 | + from b12x.attention._shared.mla.reference import pack_mla_kv_cache_reference |
| 248 | + from ignition.recipes.glm53.dcp import ContextShards, PREFILL_SLICE_ROWS |
| 249 | + |
| 250 | + module, weights, world = model_weights |
| 251 | + device, capacity, blocks, heads = torch.device("cuda:0"), 8192, 32, 64 // world |
| 252 | + torch.manual_seed(53) |
| 253 | + caches = [pack_mla_kv_cache_reference(torch.randn(blocks * 256, 512, device=device)).view(-1, 64, 528) |
| 254 | + for _ in range(world)] |
| 255 | + peers = torch.tensor([cache.data_ptr() for cache in caches], dtype=torch.int64, device=device) |
| 256 | + storage = MLAScratch.allocate(capacity, heads, blocks, "prefill", device, context_ranks=world) |
| 257 | + queries = torch.randn(storage.query.shape, dtype=torch.bfloat16, device=device) * .1 |
| 258 | + arena, operations, graphs, compiled = ScratchArena(device), {}, {}, {} |
| 259 | + kv = weights["kv_b_proj"] |
| 260 | + projection = BatchedUnquantizedProjection if isinstance(kv, torch.Tensor) else BatchedQuantizedProjection |
| 261 | + if isinstance(kv, torch.Tensor): |
| 262 | + values = kv.view(heads, 512, 512) |
| 263 | + absorb_weights, expand_weights = values[:, :256], values[:, 256:] |
| 264 | + else: |
| 265 | + values, scales = kv.weight.values.view(heads, 512, 512), kv.weight.scale_rows.view(heads, 512, 16) |
| 266 | + absorb_weights, expand_weights = (values[:, :256], scales[:, :256]), (values[:, 256:], scales[:, 256:]) |
| 267 | + with RecordedPreparationSession(device=device, checkpoint="glm53-cp-scratch", |
| 268 | + workload_schema="glm-cp-scratch-v1", numerical_policy="default") as session: |
| 269 | + session.register_scratch(arena) |
| 270 | + for rows in (capacity, 2049): |
| 271 | + for alias in (False, True): |
| 272 | + key = rows, alias |
| 273 | + s = storage.rows(rows, "prefill") |
| 274 | + if not alias: |
| 275 | + s = replace(s, expanded=torch.empty_like(s.expanded)) |
| 276 | + metadata = IndexMetadata.allocate(rows, 1, blocks, device) |
| 277 | + metadata.positions.copy_(torch.arange(rows, device=device) + 8) |
| 278 | + metadata.block_table[0].copy_(torch.arange(blocks, device=device)) |
| 279 | + s.index.physical.fill_(-1) |
| 280 | + s.index.physical[:, :4].copy_(torch.arange(4, device=device)) |
| 281 | + s.index.selected_lengths.fill_(4) |
| 282 | + # Isolate the production CP slice loop with real projections, attention and peer-cache reads. |
| 283 | + op = PreparedGlmMLA.__new__(PreparedGlmMLA) |
| 284 | + op.scratch, op.metadata, op.state = s, metadata, SimpleNamespace(peers=peers) |
| 285 | + op.context = SimpleNamespace(shards=ContextShards.on(0, world, device)) |
| 286 | + op.windows, op.slices = (), [] |
| 287 | + caps = sparse_mla.Caps(device=device, num_q_heads=heads, max_q_rows=PREFILL_SLICE_ROWS, max_width=2051, |
| 288 | + softmax_scale=module.scaling, max_batch=PREFILL_SLICE_ROWS, max_kv_rows=s.slab.shape[0], |
| 289 | + kv_dtype=torch.uint8, head_dim=512, v_head_dim=512, page_size=64, |
| 290 | + model_type=sparse_mla.ModelType.GLM_NEXT, cache_record_bytes=528, mode="extend") |
| 291 | + for start in range(0, rows, PREFILL_SLICE_ROWS): |
| 292 | + end = min(start + PREFILL_SLICE_ROWS, rows) |
| 293 | + absorbed, attended = s.absorbed_query[:end - start], s.attended[:end - start] |
| 294 | + name = f"cp{rows}.{alias}.{start}" |
| 295 | + absorb = projection(s.query[start:end].transpose(0, 1), absorb_weights, |
| 296 | + absorbed.transpose(0, 1), major="n", session=session, name=name + ".absorb") |
| 297 | + attention = PreparedSparseMLA(absorbed, s.slab.view(-1, 64, 528), s.index.physical[start:end], |
| 298 | + s.visible[start:end], s.index.selected_lengths[start:end], attended, caps=caps, |
| 299 | + session=session, name=name + ".attention", workspace=arena) |
| 300 | + expand = projection(attended.transpose(0, 1), expand_weights, |
| 301 | + s.expanded[start:end].transpose(0, 1), major="k", session=session, name=name + ".expand") |
| 302 | + op.slices.append((start, end, absorb, attention, expand)) |
| 303 | + operations[key] = op |
| 304 | + session.tune() |
| 305 | + for key, op in operations.items(): |
| 306 | + def forward(op=op): |
| 307 | + op.scratch.query.copy_(queries[:op.scratch.query.shape[0]]) |
| 308 | + op.attend_gathered_context() |
| 309 | + arena.tensor.fill_(255) |
| 310 | + compiled[key] = compile_forward(forward) |
| 311 | + compiled[key]() |
| 312 | + graphs[key] = torch.cuda.CUDAGraph() |
| 313 | + with session.capture(), torch.cuda.graph(graphs[key]): |
| 314 | + compiled[key]() |
| 315 | + session.freeze() |
| 316 | + for rows in (capacity, 2049, capacity): |
| 317 | + graphs[rows, False].replay() |
| 318 | + expected = operations[rows, False].scratch.expanded.clone() |
| 319 | + storage.query.fill_(float("nan")) |
| 320 | + graphs[rows, True].replay() |
| 321 | + torch.cuda.synchronize(device) |
| 322 | + assert torch.equal(operations[rows, True].scratch.expanded, expected) |
| 323 | + for graph in graphs.values(): |
| 324 | + graph.reset() |
0 commit comments