Remove reverted test (#22058)
This commit is contained in:
Ke Bao
@@ -1,272 +0,0 @@
|
||||
"""Correctness tests for fused_temperature_softmax Triton kernel."""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
from flashinfer.sampling import softmax as flashinfer_softmax
|
||||
|
||||
from sglang.srt.layers.fused_sampling import (
|
||||
fused_temperature_softmax,
|
||||
fused_temperature_softmax_inplace,
|
||||
)
|
||||
from sglang.srt.utils import get_device
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=15,
|
||||
suite="stage-b-test-1-gpu-small",
|
||||
disabled="Test cannot pass in CI due to numerical precision issues",
|
||||
)
|
||||
|
||||
|
||||
def reference_temperature_softmax(logits, temperatures):
|
||||
"""Reference implementation: div + softmax (separate kernels)."""
|
||||
logits = logits.clone()
|
||||
logits.div_(temperatures)
|
||||
return torch.softmax(logits, dim=-1).float()
|
||||
|
||||
|
||||
class TestFusedTemperatureSoftmax(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
torch.set_default_device(get_device())
|
||||
torch.manual_seed(42)
|
||||
|
||||
def _check_close(self, fused, ref, atol=1e-5, rtol=1e-5):
|
||||
"""Assert outputs are close and both are valid probability distributions."""
|
||||
self.assertEqual(fused.shape, ref.shape)
|
||||
# Valid probabilities: non-negative, sum to ~1
|
||||
self.assertTrue((fused >= 0).all(), f"Negative probabilities in fused output")
|
||||
row_sums = fused.sum(dim=-1)
|
||||
torch.testing.assert_close(
|
||||
row_sums,
|
||||
torch.ones_like(row_sums),
|
||||
atol=1e-4,
|
||||
rtol=1e-4,
|
||||
)
|
||||
torch.testing.assert_close(fused, ref, atol=atol, rtol=rtol)
|
||||
|
||||
# --- out-of-place kernel ---
|
||||
|
||||
def test_basic(self):
|
||||
logits = torch.randn(4, 1024, dtype=torch.bfloat16)
|
||||
temps = torch.tensor([0.7, 1.0, 1.5, 2.0], dtype=torch.float32).view(-1, 1)
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_large_vocab(self):
|
||||
logits = torch.randn(8, 128256, dtype=torch.bfloat16)
|
||||
temps = torch.full((8, 1), 0.6, dtype=torch.float32)
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_batch_sizes(self):
|
||||
for bs in [1, 2, 16, 64, 128, 512]:
|
||||
logits = torch.randn(bs, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.rand(bs, 1, dtype=torch.float32) * 1.5 + 0.1
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_temperature_one(self):
|
||||
"""Temperature=1.0 should be equivalent to plain softmax."""
|
||||
logits = torch.randn(16, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.ones(16, 1, dtype=torch.float32)
|
||||
ref = torch.softmax(logits.float(), dim=-1)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_very_low_temperature(self):
|
||||
"""Very low temperature should produce near-one-hot distribution."""
|
||||
logits = torch.randn(4, 1024, dtype=torch.bfloat16)
|
||||
temps = torch.full((4, 1), 0.01, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
# Max probability should be very close to 1.0
|
||||
max_probs = fused.max(dim=-1).values
|
||||
self.assertTrue((max_probs > 0.99).all())
|
||||
|
||||
def test_very_high_temperature(self):
|
||||
"""Very high temperature should produce near-uniform distribution."""
|
||||
logits = torch.randn(4, 1024, dtype=torch.bfloat16)
|
||||
temps = torch.full((4, 1), 100.0, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
uniform = 1.0 / 1024
|
||||
self.assertTrue(
|
||||
(fused - uniform).abs().max() < 0.01,
|
||||
"High temperature should produce near-uniform distribution",
|
||||
)
|
||||
|
||||
def test_fp16_input(self):
|
||||
logits = torch.randn(8, 32000, dtype=torch.float16)
|
||||
temps = torch.rand(8, 1, dtype=torch.float32) * 1.5 + 0.1
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-3, rtol=1e-2)
|
||||
|
||||
def test_fp32_input(self):
|
||||
logits = torch.randn(8, 32000, dtype=torch.float32)
|
||||
temps = torch.rand(8, 1, dtype=torch.float32) + 0.5
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_mixed_temperatures(self):
|
||||
"""Each row has a different temperature."""
|
||||
logits = torch.randn(8, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.tensor(
|
||||
[0.1, 0.5, 0.7, 1.0, 1.2, 1.5, 2.0, 5.0], dtype=torch.float32
|
||||
).view(-1, 1)
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self._check_close(fused, ref, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_empty_batch(self):
|
||||
logits = torch.randn(0, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.ones(0, 1, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self.assertEqual(fused.shape, (0, 32000))
|
||||
|
||||
# --- in-place kernel ---
|
||||
|
||||
def test_inplace_basic(self):
|
||||
logits = torch.randn(8, 32000, dtype=torch.float32)
|
||||
temps = torch.rand(8, 1, dtype=torch.float32) * 1.5 + 0.1
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused_temperature_softmax_inplace(logits, temps)
|
||||
# In-place writes back to logits in the original dtype
|
||||
self._check_close(logits.float(), ref, atol=1e-5, rtol=1e-5)
|
||||
|
||||
def test_inplace_bf16(self):
|
||||
logits = torch.randn(8, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.rand(8, 1, dtype=torch.float32) + 0.5
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused_temperature_softmax_inplace(logits, temps)
|
||||
self._check_close(logits.float(), ref, atol=2e-3, rtol=2e-3)
|
||||
|
||||
def test_inplace_large_vocab(self):
|
||||
logits = torch.randn(4, 128256, dtype=torch.bfloat16)
|
||||
temps = torch.full((4, 1), 0.8, dtype=torch.float32)
|
||||
ref = reference_temperature_softmax(logits, temps)
|
||||
fused_temperature_softmax_inplace(logits, temps)
|
||||
self._check_close(logits.float(), ref, atol=2e-3, rtol=2e-3)
|
||||
|
||||
# --- exact known-value correctness ---
|
||||
|
||||
def test_known_uniform_logits(self):
|
||||
"""Identical logits must produce uniform distribution regardless of temperature."""
|
||||
logits = torch.zeros(2, 5, dtype=torch.float32)
|
||||
temps = torch.tensor([0.5, 2.0], dtype=torch.float32).view(-1, 1)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
expected = torch.full((2, 5), 0.2, dtype=torch.float32, device=fused.device)
|
||||
torch.testing.assert_close(fused, expected, atol=1e-6, rtol=1e-6)
|
||||
|
||||
def test_known_softmax_values(self):
|
||||
"""Verify against hand-computed softmax(logits / T)."""
|
||||
logits = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32)
|
||||
temps = torch.tensor([[1.0]], dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
# softmax([1,2,3]) = exp([1,2,3]) / sum(exp([1,2,3]))
|
||||
e = torch.exp(logits)
|
||||
expected = (e / e.sum(dim=-1, keepdim=True)).to(fused.device)
|
||||
torch.testing.assert_close(fused, expected, atol=1e-6, rtol=1e-6)
|
||||
|
||||
def test_known_softmax_with_temperature(self):
|
||||
"""Verify softmax([1,2,3] / 0.5) against hand computation."""
|
||||
logits = torch.tensor([[1.0, 2.0, 3.0]], dtype=torch.float32)
|
||||
temps = torch.tensor([[0.5]], dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
scaled = logits / 0.5
|
||||
e = torch.exp(scaled)
|
||||
expected = (e / e.sum(dim=-1, keepdim=True)).to(fused.device)
|
||||
torch.testing.assert_close(fused, expected, atol=1e-6, rtol=1e-6)
|
||||
|
||||
# --- argmax preservation ---
|
||||
|
||||
def test_argmax_preserved(self):
|
||||
"""argmax must be invariant to temperature for finite T > 0."""
|
||||
logits = torch.randn(64, 32000, dtype=torch.bfloat16)
|
||||
original_argmax = logits.float().argmax(dim=-1)
|
||||
for t_val in [0.1, 0.5, 1.0, 2.0, 10.0]:
|
||||
temps = torch.full((64, 1), t_val, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
fused_argmax = fused.argmax(dim=-1)
|
||||
self.assertTrue(
|
||||
(original_argmax == fused_argmax).all(),
|
||||
f"argmax changed at temperature={t_val}",
|
||||
)
|
||||
|
||||
# --- numerical stability ---
|
||||
|
||||
def test_large_logits_no_nan(self):
|
||||
"""Extreme logit magnitudes must not produce NaN or Inf."""
|
||||
logits = torch.tensor(
|
||||
[[1e6, -1e6, 0.0], [1e4, 1e4 + 1, 1e4 - 1]], dtype=torch.float32
|
||||
)
|
||||
temps = torch.tensor([[1.0], [0.01]], dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
self.assertFalse(torch.isnan(fused).any(), "NaN in output")
|
||||
self.assertFalse(torch.isinf(fused).any(), "Inf in output")
|
||||
row_sums = fused.sum(dim=-1)
|
||||
torch.testing.assert_close(
|
||||
row_sums,
|
||||
torch.ones_like(row_sums),
|
||||
atol=1e-4,
|
||||
rtol=1e-4,
|
||||
)
|
||||
|
||||
def test_large_logits_inplace_no_nan(self):
|
||||
"""In-place variant: extreme logits must not produce NaN or Inf."""
|
||||
logits = torch.tensor(
|
||||
[[1e6, -1e6, 0.0], [1e4, 1e4 + 1, 1e4 - 1]], dtype=torch.float32
|
||||
)
|
||||
temps = torch.tensor([[1.0], [0.01]], dtype=torch.float32)
|
||||
fused_temperature_softmax_inplace(logits, temps)
|
||||
self.assertFalse(torch.isnan(logits).any(), "NaN in output")
|
||||
self.assertFalse(torch.isinf(logits).any(), "Inf in output")
|
||||
|
||||
# --- comparison with flashinfer.sampling.softmax ---
|
||||
|
||||
def test_vs_flashinfer_basic(self):
|
||||
logits = torch.randn(4, 1024, dtype=torch.bfloat16)
|
||||
temps = torch.tensor([0.7, 1.0, 1.5, 2.0], dtype=torch.float32).view(-1, 1)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
fi = flashinfer_softmax(logits, temperature=temps.view(-1))
|
||||
self._check_close(fused, fi, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_vs_flashinfer_large_vocab(self):
|
||||
logits = torch.randn(8, 128256, dtype=torch.bfloat16)
|
||||
temps = torch.full((8, 1), 0.6, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
fi = flashinfer_softmax(logits, temperature=temps.view(-1))
|
||||
self._check_close(fused, fi, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_vs_flashinfer_batch_sizes(self):
|
||||
for bs in [1, 16, 64, 128, 512]:
|
||||
logits = torch.randn(bs, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.rand(bs, 1, dtype=torch.float32) * 1.5 + 0.1
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
fi = flashinfer_softmax(logits, temperature=temps.view(-1))
|
||||
self._check_close(fused, fi, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_vs_flashinfer_scalar_temperature(self):
|
||||
logits = torch.randn(16, 32000, dtype=torch.bfloat16)
|
||||
temps_2d = torch.full((16, 1), 0.8, dtype=torch.float32)
|
||||
fused = fused_temperature_softmax(logits, temps_2d)
|
||||
fi = flashinfer_softmax(logits, temperature=0.8)
|
||||
self._check_close(fused, fi, atol=1e-4, rtol=1e-3)
|
||||
|
||||
def test_vs_flashinfer_mixed_temperatures(self):
|
||||
logits = torch.randn(8, 32000, dtype=torch.bfloat16)
|
||||
temps = torch.tensor(
|
||||
[0.1, 0.5, 0.7, 1.0, 1.2, 1.5, 2.0, 5.0], dtype=torch.float32
|
||||
).view(-1, 1)
|
||||
fused = fused_temperature_softmax(logits, temps)
|
||||
fi = flashinfer_softmax(logits, temperature=temps.view(-1))
|
||||
self._check_close(fused, fi, atol=1e-4, rtol=1e-3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user