Repository navigation
Add out-of-tree DFlash extension points - #38740
Conversation
fe2fe4b to
45eb8f3
Compare
alexnails
left a comment
alexnails
left a comment
There was a problem hiding this comment.
I have not really scoped the other parts of spec dec yet, what does it look to lift eagle and other spec dec methods?
|
@alexnails This PR is explicitly targeting adding a DFlash OOT, I have not looked at other spec dec method yet. |
| handle_speculative_decoding(args) | ||
|
|
||
| self.assertEqual(args.speculative_draft_attention_backend, "custom_backend") | ||
| self.assertTrue(any("got 'typo'" in message for message in logs.output)) |
There was a problem hiding this comment.
This assertion does not match the new warning message: the implementation logs does not support attention_backend 'typo', without got. This test should fail once it is executed. Could we update the assertion or the warning text?
| PlatformWithoutDefault(), | ||
| ), | ||
| patch( | ||
| "sglang.srt.arg_groups.overrides.attention_backends_of", |
There was a problem hiding this comment.
speculative_hook imports attention_backends_of directly, so patching the symbol in overrides will not replace the binding used by the code under test. Could we patch sglang.srt.arg_groups.speculative_hook.attention_backends_of instead?
| return "triton" if get_platform().is_hip else "flashinfer" | ||
|
|
||
| def is_supported_backend(backend: str) -> bool: | ||
| return backend in supported_draft_backends or ( |
There was a problem hiding this comment.
For an out-of-tree platform, should this capability hook be authoritative rather than additive? Currently, any backend in supported_draft_backends is accepted even if supports_speculative_draft_attention_backend() returns False, and trtllm_mha bypasses the hook entirely. This may defer an unsupported-backend failure until initialization. Could we branch on is_out_of_tree() first and add tests for explicit built-in backend values on an OOT platform?
There was a problem hiding this comment.
Thank you! that makes sense, fixed!
| def get_fallback_backend() -> str: | ||
| if current_platform.is_out_of_tree(): | ||
| try: | ||
| return current_platform.get_default_attention_backend() |
There was a problem hiding this comment.
Should the result of get_default_attention_backend() also be validated with supports_speculative_draft_attention_backend()? A platform default attention backend may not necessarily be supported by the DFlash draft worker. Alternatively, a dedicated get_default_speculative_draft_attention_backend() hook might make this contract clearer.
e4cd96a to
076a625
Compare
| @@ -0,0 +1,190 @@ | |||
| import unittest | |||
There was a problem hiding this comment.
The test file is too long, can we clean up this?
There was a problem hiding this comment.
Did a clean up, lint, and rebase~
076a625 to
a02e333
Compare
…points # Conflicts: # python/sglang/srt/platforms/interface.py
a02e333 to
c8dce21
Compare
3 participants
Motivation
Out-of-tree platforms currently cannot use DFlash without patching SGLang
internals. DFlash support, draft attention backend selection, speculative KV-
cache location assignment, and graph-capture behavior are hard-coded for
built-in platforms.
Modifications
SRTPlatform.backends.
support_cuda_graph()when initializing the DFlash draft worker.and cache-location dispatch.
Existing behavior for built-in platforms remains unchanged.
Accuracy Tests
This changes platform integration and dispatch rather than model computation.
Added focused unit tests covering:
Local runtime tests were not run because the checkout does not currently have
torchorpytestinstalled. Python syntax validation passed.Speed Tests and Profiling
Not run. No model kernels or numerical operations are changed. Built-in
platforms retain their existing implementations; the new dispatch path is
used only by out-of-tree platforms.
Checklist
docs.sglang.io/developer_guide/contribution_guide.html#format-code-with-pre-
commit).
docs.sglang.io/developer_guide/contribution_guide.html#run-and-add-unit-
tests).
extension points with docstrings.
computation is unchanged.
CI States
Latest PR Test (Base): ❌ Run #35443345709
Latest PR Test (Extra): ❌ Run #35443345603
Latest PR Test (AMD ROCm 10): ❌ Run #35443345692