[DSA] Trim top-k v2 output modes and tighten its PDL waits (#35041)
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
DarkSharpness
co-authored by
Claude Opus 5
Claude Opus 5
parent
1c82955861
commit
746418a1ec
@@ -141,28 +141,41 @@ def _reference(scores, seq_lens, k):
|
||||
return ref
|
||||
|
||||
|
||||
def _plan(seq_lens):
|
||||
"""Plan, then break stream adjacency with the transform launch.
|
||||
|
||||
The transform kernel prefetches the plan metadata BEFORE its PDL wait, which
|
||||
is only legal while the plan kernel is not the immediately preceding kernel
|
||||
in the stream -- in production the plan is built during per-forward metadata
|
||||
prep, a whole model forward earlier. Launching the two back to back would
|
||||
read the plan through a programmatic dependency that guarantees no memory
|
||||
visibility, so mirror the production separation instead.
|
||||
"""
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
torch.cuda.synchronize()
|
||||
return metadata
|
||||
|
||||
|
||||
def _run(scores, seq_lens, page_table, inv_cpu, k):
|
||||
batch = scores.shape[0]
|
||||
metadata = _plan(seq_lens)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata)
|
||||
torch.cuda.synchronize()
|
||||
out_cpu = out.cpu().tolist()
|
||||
return [_invert(out_cpu[i], inv_cpu[i]) for i in range(batch)]
|
||||
|
||||
|
||||
def _run_raw(scores, seq_lens, page_table, k):
|
||||
"""Run the kernel and return its optional raw (pre-transform) top-k index
|
||||
output per row, dropping -1 padding -- the selected positions themselves,
|
||||
NOT the page-table transform of them."""
|
||||
def _run_raw(scores, seq_lens, k):
|
||||
"""Run with no page table and return the selected indices per row, dropping
|
||||
-1 padding -- the selected positions themselves, NOT a page transform."""
|
||||
batch = scores.shape[0]
|
||||
metadata = _plan(seq_lens)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
raw = torch.full((batch, k), -1, dtype=torch.int32, device=scores.device)
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
|
||||
topk_transform_512_v2(scores, seq_lens, None, out, PAGE_SIZE, metadata)
|
||||
torch.cuda.synchronize()
|
||||
raw_cpu = raw.cpu().tolist()
|
||||
return [[v for v in raw_cpu[i] if v != -1] for i in range(batch)]
|
||||
out_cpu = out.cpu().tolist()
|
||||
return [[v for v in out_cpu[i] if v != -1] for i in range(batch)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_mode", ["identity", "perm"])
|
||||
@@ -229,68 +242,25 @@ def test_topk_v2_ragged(batch: int, shape: str, k: int, per_row_pt: bool) -> Non
|
||||
_assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, lengths.cpu(), k)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("page_mode", ["identity", "perm"])
|
||||
@pytest.mark.parametrize(
|
||||
"batch,seq",
|
||||
[
|
||||
(8, 256), # trivial
|
||||
(8, 4096), # register
|
||||
(4, 131072), # fused small-batch cluster
|
||||
(64, 131072), # persistent cluster + main<3> epilogue
|
||||
(256, 131072), # non-cluster streaming
|
||||
],
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_topk_v2_raw_indices(batch: int, seq: int, page_mode: str) -> None:
|
||||
"""The optional raw-index output must be the pre-transform position of each
|
||||
transformed output slot (out[j] == page_to_indices(raw[j])), and -1 aligns."""
|
||||
k = 512
|
||||
torch.manual_seed(batch * 131 + seq)
|
||||
device = "cuda"
|
||||
width = (seq + 3) & ~3
|
||||
scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
|
||||
seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
|
||||
num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
page_table, inv_cpu = _make_page_table(batch, num_pages, page_mode, device)
|
||||
out = torch.full((batch, k), -1, dtype=torch.int32, device=device)
|
||||
raw = torch.full((batch, k), -1, dtype=torch.int32, device=device)
|
||||
|
||||
metadata = plan_topk_v2(seq_lens)
|
||||
topk_transform_512_v2(scores, seq_lens, page_table, out, PAGE_SIZE, metadata, raw)
|
||||
torch.cuda.synchronize()
|
||||
|
||||
out_cpu, raw_cpu = out.cpu().tolist(), raw.cpu().tolist()
|
||||
for i in range(batch):
|
||||
for j in range(k):
|
||||
o, r = out_cpu[i][j], raw_cpu[i][j]
|
||||
if o == -1:
|
||||
assert r == -1, f"b={i} j={j}: out=-1 but raw={r}"
|
||||
else:
|
||||
inv = (int(inv_cpu[i][o >> PAGE_BITS]) << PAGE_BITS) | (o & PAGE_MASK)
|
||||
assert r == inv, f"b={i} j={j}: raw={r} != inverse(out)={inv}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("k", [512, 1024, 2048])
|
||||
@pytest.mark.parametrize("batch,seq", FIXED_CONFIGS)
|
||||
@torch.inference_mode()
|
||||
def test_topk_v2_output_indices(batch: int, seq: int, k: int) -> None:
|
||||
"""Validate the raw (pre-transform) index output DIRECTLY against torch.topk.
|
||||
"""Validate the selected indices DIRECTLY against torch.topk.
|
||||
|
||||
Unlike ``test_topk_v2`` -- which checks the page-transformed output and inverts
|
||||
it through the page table -- this exercises the selected indices themselves, so
|
||||
it isolates the top-k selection from the page-table transform. A permuted page
|
||||
table is used so raw != out, catching any bug that leaks transformed page
|
||||
indices into the raw buffer. Covers every dispatch template/boundary.
|
||||
Runs the no-page-table mode, so the output is the selected positions
|
||||
themselves. Unlike ``test_topk_v2`` -- which checks the page-transformed
|
||||
output and inverts it through the page table -- this isolates the top-k
|
||||
selection from the transform, and it is the only coverage of that mode.
|
||||
Covers every dispatch template/boundary.
|
||||
"""
|
||||
torch.manual_seed(batch * 100003 + seq * 7 + k + 1)
|
||||
device = "cuda"
|
||||
width = (seq + 3) & ~3
|
||||
scores = torch.randn(batch, width, dtype=torch.float32, device=device)[:, :seq]
|
||||
seq_lens = torch.full((batch,), seq, dtype=torch.int32, device=device)
|
||||
num_pages = (seq + PAGE_SIZE - 1) // PAGE_SIZE
|
||||
page_table, _ = _make_page_table(batch, num_pages, "perm", device)
|
||||
|
||||
our_raw = _run_raw(scores, seq_lens, page_table, k)
|
||||
our_raw = _run_raw(scores, seq_lens, k)
|
||||
ref_raw = _reference(scores, seq_lens, k)
|
||||
_assert_topk_close(scores.cpu(), ref_raw, our_raw, batch, seq_lens.cpu(), k)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user