Skip to content

Remove device side validation of QB bounds #3734

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ksivaman wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Remove device side validation of QB bounds #3734

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ba60c50
Remove QB bounds device side validation
ksivaman Oct 23, 2026
82288fa
[pre-commit.ci] auto fixes from pre-commit.com hooks
pre-commit-ci[bot] Oct 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
216 changes: 37 additions & 179 deletions tests/pytorch/test_fused_router.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -661,109 +661,55 @@ def test_qb_topk_uses_logits_device(histogram_mode):


@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
@pytest.mark.parametrize("invalid_bounds", ["equal", "reversed", "nonfinite"])
def test_qb_topk_rejects_invalid_bin_bounds(histogram_mode, invalid_bounds):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
bounds = {
"equal": [1.0, 1.0],
"reversed": [1.0, -1.0],
"nonfinite": [float("nan"), 1.0],
}[invalid_bounds]
bin_bounds = torch.tensor(bounds, device="cuda", dtype=torch.float32)
with pytest.raises(ValueError, match="finite with lower < upper"):
fused_topk_with_score_function(
logits,
4,
False,
None,
None,
None,
"sigmoid",
expert_bias,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)

@pytest.mark.parametrize("use_dense_indices", [False, True])
@pytest.mark.parametrize(
"invalid_bounds",
[[1.0, 1.0], [1.0, -1.0], [float("nan"), 1.0], [float("-inf"), float("inf")]],
ids=["equal", "reversed", "nan", "infinite"],
)
def test_qb_topk_invalid_bin_bounds_stay_in_histogram(
histogram_mode, use_dense_indices, invalid_bounds
):
num_tokens, num_experts, num_bins, topk = 8, 16, 32, 4
logits = torch.randn(num_tokens, num_experts, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(num_experts, device="cuda", dtype=torch.float32)
histogram = torch.zeros(num_experts, num_bins, device="cuda", dtype=torch.int32)
bin_bounds = torch.tensor(invalid_bounds, device="cuda", dtype=torch.float32)
topk_indices = (
torch.empty(num_tokens, topk, device="cuda", dtype=torch.int32)
if use_dense_indices
else None
)

@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
def test_qb_topk_revalidates_updated_bin_bounds(histogram_mode):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
fused_topk_with_score_function(
# bin_bounds values are not validated, so invalid bounds must neither raise nor write outside
# the histogram.
probs, routing_output = fused_topk_with_score_function(
logits,
4,
topk,
False,
None,
None,
None,
"sigmoid",
expert_bias,
topk_indices=topk_indices,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)
bin_bounds.fill_(0.0)
with pytest.raises(ValueError, match="finite with lower < upper"):
fused_topk_with_score_function(
logits,
4,
False,
None,
None,
None,
"sigmoid",
expert_bias,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)


@pytest.mark.parametrize(
"histogram_mode",
[QBHistogramMode.TWO_KERNEL, QBHistogramMode.FUSED_ATOMIC],
)
@pytest.mark.parametrize("use_dense_indices", [False, True])
def test_qb_raw_binding_rejects_invalid_bin_bounds_recoverably(histogram_mode, use_dense_indices):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
topk_indices = (
torch.empty(8, 4, device="cuda", dtype=torch.int32) if use_dense_indices else None
)
invalid_bounds = torch.tensor([1.0, 1.0], device="cuda", dtype=torch.float32)
with pytest.raises(RuntimeError, match="finite with lower < upper"):
tex.fused_topk_with_score_function_qb_fwd(
logits,
4,
None,
expert_bias,
int(RoutingMapFormat.BYTEMAP),
topk_indices,
histogram,
invalid_bounds,
histogram_mode,
)
torch.cuda.synchronize()

# The validation error must not poison the CUDA context.
# Routing does not read the bounds, and every (token, expert) pair is counted in some bin.
valid_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
tex.fused_topk_with_score_function_qb_fwd(
logits,
4,
None,
expert_bias,
int(RoutingMapFormat.BYTEMAP),
topk_indices,
histogram,
valid_bounds,
histogram_mode,
)
torch.cuda.synchronize()
reference = qb_topk_score_function_pytorch(logits, topk, expert_bias, valid_bounds, num_bins)
torch.testing.assert_close(probs, reference["probs"])
if use_dense_indices:
torch.testing.assert_close(
topk_indices_to_routing_map(routing_output, num_experts), reference["routing_map"]
)
else:
torch.testing.assert_close(routing_output, reference["routing_map"])
assert histogram.sum().item() == num_tokens * num_experts


@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
Expand Down Expand Up @@ -848,7 +794,6 @@ def run_iteration():
qb_histogram_mode=histogram_mode,
)
bin_bounds.copy_(next_bounds)
mark_qb_bin_bounds_validated(bin_bounds)
return probs, routing_map

# Match full-iteration capture: an eager warmup ends with a bounds update, and the same
Expand All @@ -871,98 +816,11 @@ def run_iteration():
torch.testing.assert_close(bin_bounds, next_bounds)


@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
def test_qb_topk_eager_call_revalidates_after_captured_bounds_update(histogram_mode):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
def test_mark_qb_bin_bounds_validated_is_deprecated():
bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)
next_bounds = torch.tensor([-0.75, 0.25], device="cuda", dtype=torch.float32)

def run_router():
return fused_topk_with_score_function(
logits,
4,
False,
None,
None,
None,
"sigmoid",
expert_bias,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)

run_router()
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
bin_bounds.copy_(next_bounds)
with pytest.warns(DeprecationWarning, match="mark_qb_bin_bounds_validated is deprecated"):
mark_qb_bin_bounds_validated(bin_bounds)

# Replayed device writes do not advance bin_bounds._version. The capture-time marker must not
# let a subsequent eager router call inherit trust in the replay-produced values.
next_bounds.zero_()
graph.replay()
torch.cuda.synchronize()
with pytest.raises(ValueError, match="finite with lower < upper"):
run_router()


@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
def test_qb_topk_cuda_graph_rejects_unvalidated_bounds(histogram_mode):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)

with pytest.raises(RuntimeError, match="current version must be validated"):
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
fused_topk_with_score_function(
logits,
4,
False,
None,
None,
None,
"sigmoid",
expert_bias,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)


@pytest.mark.parametrize("histogram_mode", ["two_kernel", "fused_atomic"])
def test_qb_topk_cuda_graph_rejects_stale_bounds_validation(histogram_mode):
logits = torch.randn(8, 16, device="cuda", dtype=torch.float32)
expert_bias = torch.zeros(16, device="cuda", dtype=torch.float32)
histogram = torch.zeros(16, 32, device="cuda", dtype=torch.int32)
bin_bounds = torch.tensor([-1.0, 1.0], device="cuda", dtype=torch.float32)

def run_router():
return fused_topk_with_score_function(
logits,
4,
False,
None,
None,
None,
"sigmoid",
expert_bias,
qb_histogram=histogram,
qb_bin_bounds=bin_bounds,
qb_histogram_mode=histogram_mode,
)

run_router()
bin_bounds.zero_()
with pytest.raises(RuntimeError, match="current version must be validated"):
graph = torch.cuda.CUDAGraph()
with torch.cuda.graph(graph):
run_router()


@pytest.mark.parametrize(
"histogram_mode",
Expand Down
Loading
Loading