Skip to content

Commit 239e65c

Browse files
drisspgdrisspg
committed
Add dedicated NVFP4 TMA production and compact profiling
1 parent 09b8feb commit 239e65c

8 files changed

Lines changed: 1064 additions & 194 deletions

File tree

‎test/test_mxfp8_tma.py‎

Lines changed: 20 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -265,12 +265,26 @@ def test_mxfp8_tma_gemv_profiles_labeled_regions(grid_scheduler, num_persistent_
265265
grid_scheduler=grid_scheduler,
266266
num_persistent_ctas=num_persistent_ctas,
267267
)
268+
legacy_profile_buffer = torch.zeros(
269+
op.num_profile_units * (1 + 4 * op.max_profile_events_per_cta),
270+
dtype=torch.int64,
271+
device=q_input.device,
272+
)
273+
with pytest.raises(ValueError, match="compact.*exactly"):
274+
op.interface(
275+
q_input,
276+
weight,
277+
input_scale,
278+
weight_scale,
279+
profile_buffer=legacy_profile_buffer,
280+
)
268281

269282
with profile_session(
270283
max_events_per_unit=op.max_profile_events_per_cta,
271284
num_units=(op.num_profile_units, "CTA"),
272285
tag_names=list(MXFP8_TMA_PROFILE_TAGS),
273286
device=q_input.device,
287+
compact=True,
274288
) as (prof, tags):
275289
actual = op.interface(
276290
q_input,
@@ -285,8 +299,12 @@ def test_mxfp8_tma_gemv_profiles_labeled_regions(grid_scheduler, num_persistent_
285299
).bfloat16()
286300
torch.testing.assert_close(actual, expected, atol=1.0, rtol=0.05)
287301
events = decode_events(prof, tags)
288-
assert len(events) == op.num_profile_units * (3 * op.num_k_tiles + 1)
289-
assert {event.tag_name for event in events} == set(MXFP8_TMA_PROFILE_TAGS)
302+
assert len(events) == op.num_profile_units * op.profile_events_per_cta
303+
assert {event.tag_name for event in events} == set(MXFP8_TMA_PROFILE_TAGS) - {
304+
"input_scale_acquire",
305+
"input_scale_copy",
306+
"input_scale_wait",
307+
}
290308
assert {event.unit_id for event in events} == set(range(op.num_profile_units))
291309

292310

‎test/test_nvfp4_tma.py‎

Lines changed: 155 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,46 @@ def test_nvfp4_tma_matches_reference(num_compute_warps):
196196
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
197197

198198

199+
def test_nvfp4_tma_prefetch_input_scales_requires_dedicated_producer():
200+
"""Keep scale production off the consumer path by construction."""
201+
q_input, weight, input_scale, weight_scale, _, _ = make_case(128, 4096)
202+
with pytest.raises(ValueError, match="requires a dedicated producer warp"):
203+
nvfp4_tma_gemv(
204+
q_input,
205+
weight,
206+
input_scale,
207+
weight_scale,
208+
block_n=8,
209+
num_stages=3,
210+
num_compute_warps=4,
211+
prefetch_input_scales=True,
212+
)
213+
214+
215+
@pytest.mark.parametrize(
216+
("dedicated_producer_warp", "prefetch_input_scales"),
217+
[(True, False), (True, True)],
218+
)
219+
def test_nvfp4_tma_prefetch_input_scales_matches_reference(
220+
dedicated_producer_warp, prefetch_input_scales
221+
):
222+
"""Preserve NVFP4 output while independently producing and prefetching scales."""
223+
q_input, weight, input_scale, weight_scale, expected, _ = make_case(128, 4096)
224+
actual = nvfp4_tma_gemv(
225+
q_input,
226+
weight,
227+
input_scale,
228+
weight_scale,
229+
block_n=8,
230+
num_stages=3,
231+
num_compute_warps=4,
232+
dedicated_producer_warp=dedicated_producer_warp,
233+
prefetch_input_scales=prefetch_input_scales,
234+
)
235+
torch.cuda.synchronize()
236+
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
237+
238+
199239
@pytest.mark.parametrize(("block_n", "num_compute_warps"), [(8, 4), (6, 2)])
200240
def test_nvfp4_tma_scale_layout_and_paired_loads(block_n, num_compute_warps):
201241
"""Match random scales across blocked-layout boundaries and odd row ownership."""
@@ -309,11 +349,25 @@ def test_nvfp4_tma_profiles_common_pipeline_regions():
309349
enable_profiling=True,
310350
num_compute_warps=4,
311351
)
352+
legacy_profile_buffer = torch.zeros(
353+
op.num_profile_units * (1 + 4 * op.max_profile_events_per_cta),
354+
dtype=torch.int64,
355+
device="cuda",
356+
)
357+
with pytest.raises(ValueError, match="compact.*exactly"):
358+
op.interface(
359+
q_input,
360+
weight,
361+
input_scale,
362+
weight_scale,
363+
profile_buffer=legacy_profile_buffer,
364+
)
312365
with profile_session(
313366
max_events_per_unit=op.max_profile_events_per_cta,
314367
num_units=(op.num_profile_units, "CTA"),
315368
tag_names=list(NVFP4_TMA_PROFILE_TAGS),
316369
device="cuda",
370+
compact=True,
317371
) as (prof, tags):
318372
actual = op.interface(
319373
q_input,
@@ -324,11 +378,87 @@ def test_nvfp4_tma_profiles_common_pipeline_regions():
324378
)
325379
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
326380
events = decode_events(prof, tags)
327-
assert len(events) == op.num_profile_units * (3 * op.num_k_tiles + 1)
328-
assert {event.tag_name for event in events} == set(NVFP4_TMA_PROFILE_TAGS)
381+
assert len(events) == op.num_profile_units * op.profile_events_per_cta
382+
assert {event.tag_name for event in events} == set(NVFP4_TMA_PROFILE_TAGS) - {
383+
"input_scale_acquire",
384+
"input_scale_copy",
385+
"input_scale_wait",
386+
}
329387

330388

331-
def test_nvfp4_tma_matches_scaled_mm_global_scale_contract():
389+
def test_nvfp4_tma_profiles_sampled_ctas():
390+
"""Record a deterministic CTA subset without changing the launched grid."""
391+
q_input, weight, input_scale, weight_scale, expected, _ = make_case(128, 4096)
392+
op = get_nvfp4_tma_gemv(
393+
128,
394+
4096,
395+
4,
396+
enable_profiling=True,
397+
num_compute_warps=4,
398+
profile_cta_stride=4,
399+
)
400+
with profile_session(
401+
max_events_per_unit=op.max_profile_events_per_cta,
402+
num_units=(op.num_profile_units, "CTA"),
403+
tag_names=list(NVFP4_TMA_PROFILE_TAGS),
404+
device="cuda",
405+
compact=True,
406+
) as (prof, tags):
407+
actual = op.interface(
408+
q_input,
409+
weight,
410+
input_scale,
411+
weight_scale,
412+
profile_buffer=prof.tensor,
413+
)
414+
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
415+
events = decode_events(prof, tags)
416+
assert len(events) == op.num_profile_units * op.profile_events_per_cta
417+
assert {event.unit_id for event in events} == set(range(op.num_profile_units))
418+
419+
420+
def test_nvfp4_tma_profiles_dedicated_persistent_tiles():
421+
"""Profile sampled persistent tiles with dedicated matrix and scale production."""
422+
q_input, weight, input_scale, weight_scale, expected, _ = make_case(256, 4096)
423+
op = get_nvfp4_tma_gemv(
424+
256,
425+
4096,
426+
8,
427+
enable_profiling=True,
428+
num_compute_warps=4,
429+
grid_scheduler=GridScheduler.PERSISTENT,
430+
num_persistent_ctas=4,
431+
profile_cta_stride=2,
432+
dedicated_producer_warp=True,
433+
prefetch_input_scales=True,
434+
)
435+
with profile_session(
436+
max_events_per_unit=op.max_profile_events_per_cta,
437+
num_units=(op.num_profile_units, "CTA"),
438+
tag_names=list(NVFP4_TMA_PROFILE_TAGS),
439+
device="cuda",
440+
compact=True,
441+
) as (prof, tags):
442+
actual = op.interface(
443+
q_input,
444+
weight,
445+
input_scale,
446+
weight_scale,
447+
profile_buffer=prof.tensor,
448+
)
449+
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
450+
events = decode_events(prof, tags)
451+
assert len(events) == op.num_profile_units * op.profile_events_per_cta
452+
assert {event.unit_id for event in events} == set(range(op.num_profile_units))
453+
454+
455+
@pytest.mark.parametrize(
456+
("dedicated_producer_warp", "prefetch_input_scales"),
457+
[(False, False), (True, True)],
458+
)
459+
def test_nvfp4_tma_matches_scaled_mm_global_scale_contract(
460+
dedicated_producer_warp, prefetch_input_scales
461+
):
332462
"""Match F.scaled_mm blockwise and tensorwise NVFP4 scaling semantics."""
333463
from torch.nn.functional import ScalingType, SwizzleType, scaled_mm
334464

@@ -344,6 +474,8 @@ def test_nvfp4_tma_matches_scaled_mm_global_scale_contract():
344474
global_scale_a=input_global_scale,
345475
global_scale_b=weight_global_scale,
346476
num_compute_warps=4,
477+
dedicated_producer_warp=dedicated_producer_warp,
478+
prefetch_input_scales=prefetch_input_scales,
347479
)
348480
expected = scaled_mm(
349481
q_input,
@@ -471,6 +603,26 @@ def test_nvfp4_tma_compact_scale_tma_persistent_reuse():
471603
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
472604

473605

606+
@pytest.mark.parametrize("prefetch_input_scales", [False, True])
607+
def test_nvfp4_tma_dedicated_producer_persistent_reuse(prefetch_input_scales):
608+
"""Keep producer and consumer pipeline phases aligned across persistent tiles."""
609+
q_input, weight, input_scale, weight_scale, expected, _ = make_case(256, 4096)
610+
actual = nvfp4_tma_gemv(
611+
q_input,
612+
weight,
613+
input_scale,
614+
weight_scale,
615+
block_n=8,
616+
num_compute_warps=4,
617+
grid_scheduler=GridScheduler.PERSISTENT,
618+
num_persistent_ctas=4,
619+
dedicated_producer_warp=True,
620+
prefetch_input_scales=prefetch_input_scales,
621+
)
622+
torch.cuda.synchronize()
623+
torch.testing.assert_close(actual, expected, atol=2.0, rtol=0.05)
624+
625+
474626
def test_nvfp4_tma_persistent_cuda_graph():
475627
"""Replay persistent NVFP4 GEMV into caller-owned output."""
476628
q_input, weight, input_scale, weight_scale, expected, _ = make_case(256, 2048)

0 commit comments

Comments
 (0)