Repository navigation
Conversation
ch-wan
left a comment
ch-wan
left a comment
There was a problem hiding this comment.
Reproduced the issue and checked the fix independently on H200 (torch 2.14.1+cu130, triton 3.8.0), on main f47d895 with this PR's kernel change applied. The kernel file has not changed on main since this PR's base.
- Kernel test.
test_gdn_replayssm_compact_commit.pyagainst unpatched main: 84 failed, 27 passed. With this PR: 111 passed. - End to end. Qwen3.8-27B (BF16) on one H200 with
--kv-cache-dtype fp8_e4m3 --attention-backend flashinfer, MTP (--speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4) and--enable-linear-replayssm-spec. Test-only instrumentation wrote a single NaN into one request's hidden state at one verify step; each run then served 384 chat and shared-prefix requests. Requests with any nonfinite output logprob:
| seed | main | this PR |
|---|---|---|
| 11 | 5 | 1 |
| 12 | 7 | 1 |
| 13 | 5 | 1 |
On main, every extra bad request sits in the same request-pool slot as the injected one, and first turns nonfinite at the first GDN layer after the injection point. The stale ring entries poison each later request that is assigned that slot. With this PR, the NaN stays in the request it was injected into.
So on main, one transient NaN anywhere upstream becomes recurring NaNs for unrelated later requests. That may be worth a line in the description: the fix matters to anyone serving GDN models with --enable-linear-replayssm-spec.
The verify kernel and the decode ring kernel already mask their history loads, so the compact commit is the only place that needed this.
LGTM once CI passes.
Qiaolin-Yu
added
run-ci
|
Required checks failed on b6448d1; this PR cannot be merged until they pass. PR title (run)
|
gdn_replayssm_compact_commit_kernel loaded the whole ring window (MAX_CACHE_LEN rows) regardless of n_history. Reused ring storage can hold non-finite values in slots beyond the tracked history, and zero decay cannot suppress NaN/Inf there: 0*inf contamination poisoned the committed checkpoint. The loads are now masked by history_mask (with the fold math simplified to rely on it), and in the track path rows outside the tracked prefix are explicitly zeroed before the dot products, so a nonfinite update AFTER the tracked prefix can no longer corrupt it. File was byte-identical to the PR base, so the kernel hunk lands verbatim. The PR's test file is 111 CUDA cases (no CPU-exercisable logic); ported with a module-level CUDA skip per the repo convention so CPU CI collects it as one clean skip (verified: ungated copy collects exactly 111 cases). (cherry picked from commit 99a8da4) PR: sgl-project#42688 (open, ported) (cherry picked from commit 5757ec5f59bc8a3957d4157d4752a7ab8a9a8c0e)
Motivation
The compact ReplaySSM checkpoint kernel reads the whole history ring, then gives unused entries zero decay. Reused ring storage can contain NaN or infinity, and multiplying those values by zero still contaminates the checkpoint. A tracked checkpoint can also be corrupted by a nonfinite update that belongs to the active history but lies after the tracked prefix.
For example, a ring with four valid updates and a stale NaN in its fifth entry should produce the same checkpoint as clean unused storage. A checkpoint tracking the first three updates should remain finite when only the fourth update is nonfinite.
Modifications
Mask unused history loads before arithmetic, and mask keys and updates to the tracked prefix before its matrix products. Nonfinite values within the active history still propagate; this prevents unrelated history from contaminating checkpoints rather than sanitizing valid inputs or existing state.
Add synthetic regression tests for all history fields, NaN and positive/negative infinity, ring wraparound, residual modes, shorter tracked prefixes, empty history, and untouched/null slots.
Accuracy Tests
All 111 regression cases passed on a CUDA GPU (no failures or skips), using
pytest -q test/registered/attention/unittests/gdn/test_gdn_replayssm_compact_commit.py. File-scoped pre-commit and Python syntax checks also passed.Speed Tests and Profiling
Not benchmarked for this revision. The change adds load/operand masks within the existing checkpoint kernel and does not add a kernel launch.
CI States
Latest PR Test (Base): 🚫 Run #39342505937
Latest PR Test (Extra): ❌ Run #39342505801
Latest PR Test (AMD ROCm 10): ❌ Run #39342505871