Skip to content

Commit 589b6de

Browse files
ani300ani300
committed
Fix PR comments and format
Signed-off-by: Antoni Viros i Martin <aviros@ibm.com>
1 parent db3437e commit 589b6de

2 files changed

Lines changed: 53 additions & 48 deletions

File tree

‎aiu_fms_testing_utils/testing/validation.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -269,6 +269,7 @@ def extract_validation_information(
269269
attention_specific_kwargs = {}
270270
if "paged" in extra_kwargs.get("attn_name", "sdpa"):
271271
from aiu_fms_testing_utils.utils.paged import generate
272+
272273
attention_specific_kwargs["prefill_chunk_size"] = prefill_chunk_size
273274
else:
274275
# TODO: Add a unified generation dependent on attn_type

‎aiu_fms_testing_utils/utils/paged.py‎

Lines changed: 52 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -120,8 +120,9 @@ def generate(
120120

121121
BLOCK_SIZE = 64
122122
if prefill_chunk_size > 0:
123-
assert prefill_chunk_size % BLOCK_SIZE == 0, "Chunk size must be a multiple of the page size"
124-
CHUNK_SIZE = prefill_chunk_size
123+
assert prefill_chunk_size % BLOCK_SIZE == 0, (
124+
"Chunk size must be a multiple of the page size"
125+
)
125126

126127
# these variables are guaranteed to be set in another location (inference.py, test_decoders.py, etc.)
127128
# if we set these variables here, we run the risk of warming up and generating with different sizes
@@ -268,36 +269,36 @@ def generate(
268269
# remove extra pads from the input_ids, slot_mapping, position_ids, mask to account for empty pages
269270
# each input should be padded to its smallest multiple of BLOCK_SIZE (64)
270271
# we need to clone these tensors to ensure the pointer offset is 0
271-
input_ids_i = input_ids[seq_i][-current_tkv:].unsqueeze(0).clone()
272-
slot_mapping_i = (
272+
input_ids_seq = input_ids[seq_i][-current_tkv:].unsqueeze(0).clone()
273+
slot_mapping_seq = (
273274
torch.tensor(slot_mapping[seq_i][-current_tkv:], dtype=torch.int64)
274275
.unsqueeze(0)
275276
.clone()
276277
)
277-
position_ids_i = (
278+
position_ids_seq = (
278279
kwargs["position_ids"][seq_i][-current_tkv:].unsqueeze(0).clone()
279280
)
280281

281282
# This view will result in a discontiguous tensor (creates a new graph during compile)
282283
# For this reason, we must explicitly make contiguous
283-
mask_i = (
284+
mask_seq = (
284285
kwargs["mask"][seq_i][:, -current_tkv:, -current_tkv:]
285286
.unsqueeze(0)
286287
.contiguous()
287288
)
288289

289290
# batch static
290-
torch._dynamo.mark_static(input_ids_i, 0)
291-
torch._dynamo.mark_static(slot_mapping_i, 0)
292-
torch._dynamo.mark_static(position_ids_i, 0)
293-
torch._dynamo.mark_static(mask_i, 0)
291+
torch._dynamo.mark_static(input_ids_seq, 0)
292+
torch._dynamo.mark_static(slot_mapping_seq, 0)
293+
torch._dynamo.mark_static(position_ids_seq, 0)
294+
torch._dynamo.mark_static(mask_seq, 0)
294295

295296
# seq dynamic
296-
torch._dynamo.mark_dynamic(input_ids_i, 1)
297-
torch._dynamo.mark_dynamic(slot_mapping_i, 1)
298-
torch._dynamo.mark_dynamic(position_ids_i, 1)
299-
torch._dynamo.mark_dynamic(mask_i, 2)
300-
torch._dynamo.mark_dynamic(mask_i, 3)
297+
torch._dynamo.mark_dynamic(input_ids_seq, 1)
298+
torch._dynamo.mark_dynamic(slot_mapping_seq, 1)
299+
torch._dynamo.mark_dynamic(position_ids_seq, 1)
300+
torch._dynamo.mark_dynamic(mask_seq, 2)
301+
torch._dynamo.mark_dynamic(mask_seq, 3)
301302

302303
# FP8 per-sentence scale handling
303304
if "fp8" in kwargs["attn_name"]:
@@ -308,24 +309,24 @@ def generate(
308309
last_n_tokens = kwargs.get("last_n_tokens", 0)
309310

310311
if prefill_chunk_size > 0:
311-
left_padded_prompt_mask_ij = None
312+
left_padded_prompt_mask_seq_chunk = None
312313
# Chunked prefill
313-
for chunk_j in range(math.ceil(current_tkv / CHUNK_SIZE)):
314-
chunk_start = -current_tkv + chunk_j * CHUNK_SIZE
314+
for chunk_j in range(math.ceil(current_tkv / prefill_chunk_size)):
315+
chunk_start = -current_tkv + chunk_j * prefill_chunk_size
315316
chunk_end = -current_tkv + min(
316-
(chunk_j + 1) * CHUNK_SIZE, current_tkv
317+
(chunk_j + 1) * prefill_chunk_size, current_tkv
317318
)
318319

319320
ids_length = input_ids[seq_i].shape[0]
320-
input_ids_ij = (
321+
input_ids_seq_chunk = (
321322
input_ids[seq_i][
322323
chunk_start + ids_length : chunk_end + ids_length
323324
]
324325
.unsqueeze(0)
325326
.clone()
326327
)
327328
slots_length = len(slot_mapping[seq_i])
328-
slot_mapping_ij = (
329+
slot_mapping_seq_chunk = (
329330
torch.tensor(
330331
slot_mapping[seq_i][
331332
chunk_start + slots_length : chunk_end
@@ -337,7 +338,7 @@ def generate(
337338
.clone()
338339
)
339340
pids_length = kwargs["position_ids"][seq_i].shape[0]
340-
position_ids_ij = (
341+
position_ids_seq_chunk = (
341342
kwargs["position_ids"][seq_i][
342343
chunk_start + pids_length : chunk_end + pids_length
343344
]
@@ -347,19 +348,21 @@ def generate(
347348

348349
# This view will result in a discontiguous tensor (creates a new graph during compile)
349350
# For this reason, we must explicitly make contiguous
350-
if left_padded_prompt_mask_ij is None:
351-
left_padded_prompt_mask_ij = (position_ids_ij == 0).sum(
352-
dim=1
353-
) - 1
354-
current_tkv_mask_ij = torch.min(
355-
torch.tensor((chunk_j + 1) * CHUNK_SIZE, dtype=torch.int64),
351+
if left_padded_prompt_mask_seq_chunk is None:
352+
left_padded_prompt_mask_seq_chunk = (
353+
position_ids_seq_chunk == 0
354+
).sum(dim=1) - 1
355+
current_tkv_mask_seq_chunk = torch.min(
356+
torch.tensor(
357+
(chunk_j + 1) * prefill_chunk_size, dtype=torch.int64
358+
),
356359
current_tkv,
357360
).unsqueeze(0)
358361

359362
table_length = len(block_table[seq_i])
360363
block_start = -current_tkv // BLOCK_SIZE + table_length
361364
block_end = chunk_end // BLOCK_SIZE + table_length
362-
block_table_ij = torch.tensor(
365+
block_table_seq_chunk = torch.tensor(
363366
block_table[seq_i][block_start:block_end], dtype=torch.int64
364367
).unsqueeze(0)
365368

@@ -368,26 +371,28 @@ def generate(
368371
"last_n_tokens": kwargs["last_n_tokens"],
369372
"past_key_value_states": current_kv_cache,
370373
"use_cache": kwargs["use_cache"],
371-
"position_ids": position_ids_ij,
372-
"left_padded_prompt_mask": left_padded_prompt_mask_ij,
373-
"current_tkv_mask": current_tkv_mask_ij,
374-
"block_table": block_table_ij,
375-
"slot_mapping": slot_mapping_ij,
374+
"position_ids": position_ids_seq_chunk,
375+
"left_padded_prompt_mask": left_padded_prompt_mask_seq_chunk,
376+
"current_tkv_mask": current_tkv_mask_seq_chunk,
377+
"block_table": block_table_seq_chunk,
378+
"slot_mapping": slot_mapping_seq_chunk,
376379
}
377380

378381
# batch static
379-
torch._dynamo.mark_static(input_ids_ij, 0)
380-
torch._dynamo.mark_static(slot_mapping_ij, 0)
381-
torch._dynamo.mark_static(position_ids_ij, 0)
382-
torch._dynamo.mark_static(block_table_ij, 0)
382+
torch._dynamo.mark_static(input_ids_seq_chunk, 0)
383+
torch._dynamo.mark_static(slot_mapping_seq_chunk, 0)
384+
torch._dynamo.mark_static(position_ids_seq_chunk, 0)
385+
torch._dynamo.mark_static(block_table_seq_chunk, 0)
383386

384387
# seq dynamic
385-
torch._dynamo.mark_dynamic(input_ids_ij, 1)
386-
torch._dynamo.mark_dynamic(slot_mapping_ij, 1)
387-
torch._dynamo.mark_dynamic(position_ids_ij, 1)
388-
torch._dynamo.mark_dynamic(block_table_ij, 1)
388+
torch._dynamo.mark_dynamic(input_ids_seq_chunk, 1)
389+
torch._dynamo.mark_dynamic(slot_mapping_seq_chunk, 1)
390+
torch._dynamo.mark_dynamic(position_ids_seq_chunk, 1)
391+
torch._dynamo.mark_dynamic(block_table_seq_chunk, 1)
389392

390-
logits, current_kv_cache = model(input_ids_ij, **chunked_kwargs)
393+
logits, current_kv_cache = model(
394+
input_ids_seq_chunk, **chunked_kwargs
395+
)
391396

392397
# only last token must be handled here to properly stack the tensors
393398
logits = logits[:, -1, :]
@@ -396,10 +401,10 @@ def generate(
396401

397402
else:
398403
output, current_kv_cache = model(
399-
input_ids_i,
400-
slot_mapping=slot_mapping_i,
401-
position_ids=position_ids_i,
402-
mask=mask_i,
404+
input_ids_seq,
405+
slot_mapping=slot_mapping_seq,
406+
position_ids=position_ids_seq,
407+
mask=mask_seq,
403408
past_key_value_states=current_kv_cache,
404409
use_cache=kwargs["use_cache"],
405410
last_n_tokens=last_n_tokens,
@@ -409,7 +414,6 @@ def generate(
409414
# only last token must be handled here to properly stack the tensors
410415
output = output[:, -1, :]
411416

412-
413417
# TODO: Figure out how to do this cleanly
414418
if "fp8" in kwargs["attn_name"]:
415419
for layer_idx, (t1, t2) in enumerate(current_kv_cache):

0 commit comments

Comments
 (0)