Run M.O.G.-SEC-27B-1M-CTX on Modal: fp8 KV cache + deployment - #1
Open
Jackson57279 wants to merge 2 commits into
Open
Jackson57279 wants to merge 2 commits into
Jackson57279 wants to merge 2 commits into
Conversation
Working-tree changes that predate the Modal + fp8-KV work and are unrelated to it, committed separately so the following commit reviews cleanly. Introduces the accelerator abstraction (freetoken/accelerator.py, hardware.py), the eager PyTorch XPU path for dense Llama, the torch-native attention backend and the XPU benchmark, plus their tests. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…oyment Adds `--kv-cache-dtype` so the paged KV cache can be stored narrower than the compute dtype, and a Modal deployment for M.O.G.-SEC-27B-1M-CTX-NVFP4 built around it. Why: single-stream decode is memory-bandwidth bound. Each token re-reads every active weight plus the entire KV cache, so tok/s ~= bandwidth / bytes-per-token. For the 27B hybrid at 1M context that is 91.6 GB/token with a bf16 cache -- 9.2 TB/s to hit 100 tok/s, which no single GPU has. An fp8 cache halves the KV half of that (65.5 -> 32.8 GB), bringing a single B200 from ~61 to ~95-108 tok/s and putting the target in reach at $6.25/hr. Engine changes: - EngineConfig.kv_cache_dtype (None = follow --dtype, so existing behaviour is bit-identical) with a resolved `kv_dtype` property used as the single source of truth by both the pool allocation and the byte budget, so capacity planning and the real allocation cannot disagree. - spec_kv_bytes_per_token budgets off the KV dtype rather than the compute dtype. - Context carries `compute_dtype`, letting the FlashInfer backend plan `q_data_type` at compute width while `kv_data_type` follows the pool -- queries are never quantized. - MHAKVCache.store_kv narrows K/V on write: store_cache is a templated raw byte-copy and cannot convert. The torch fallback aliases fp8 through uint8 because index_copy_ has no CPU kernel for float8. - BackendInfo.supports_quantized_kv gates the feature. A narrow cache on a backend that assumes q.dtype == cache.dtype is rejected at config time instead of failing inside a kernel launch, and `--attention-backend auto` skips such backends. Measured on L4/Qwen3-0.6B, same 18.44 GiB KV budget: bf16 stores 172,637 tokens at 114,688 B/token; fp8_e5m2 stores 345,271 at 57,344 -- exactly 2x the context, with a greedy continuation byte-identical to bf16 (271/271 chars). fp8_e4m3 is NOT recommended and warns: the pool stores K/V by a straight cast with no per-tensor scale and e4m3 saturates at +/-448, so activations clip and generation degenerates (observed collapsing into repeated punctuation after 14 chars). e5m2 spends the same byte on exponent range instead of mantissa and survives an unscaled cast; being the same width, e4m3 buys no extra bandwidth anyway. Also restores an availability guard in _adjust_config: the accelerator refactor replaced `torch.cuda.get_device_name(0) if torch.cuda.is_available() else None` with an unconditional call, which raised out of a driver probe on CPU-only hosts and broke config-time runs (and the test suite off-GPU). deploy/ contains the Modal app (CPU-built image, volume-cached weights and JIT artifacts, CPU-only download/validate gates before any GPU spend) and capacity.py, the bytes-per-token cost model whose numbers back the claims above. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3f81bbac96
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
Comment on lines
+336
to
+337
| t_prefill = run(1) # prefill + 1 decode step | ||
| t_full = run(gen) # prefill + gen decode steps |
There was a problem hiding this comment.
Both timings reuse the same LLM, whose default radix cache retains the first run's prompt. Consequently run(gen) receives a prefix-cache hit while run(1) performs the full prefill, so subtracting them also subtracts the 1M-token prefill cost and can produce a negative duration/nan instead of decode throughput. Disable/clear prefix caching between runs, reverse the measurement design, or time decode steps directly.
Useful? React with 👍 / 👎.
| modal run deploy/modal_app.py::validate # CPU-only config gate | ||
| modal run deploy/modal_app.py::smoke # GPU: load + generate | ||
| modal run deploy/modal_app.py::bench | ||
| modal serve deploy/modal_app.py # OpenAI-compatible endpoint |
There was a problem hiding this comment.
Running the documented modal serve deploy/modal_app.py command cannot expose an OpenAI-compatible API because this module only defines function/local entrypoints and there is no @modal.web_endpoint/@modal.asgi_app endpoint anywhere in the repository. Users can run smoke tests and benchmarks, but the new deployment cannot actually serve requests as advertised.
Useful? React with 👍 / 👎.
Comment on lines
+134
to
+136
| ("fp8 KV+head", dict(kv_bits=8, lm_head_bits=8)), | ||
| ("fp8 all", dict(kv_bits=8, lm_head_bits=8, linear_attn_bits=8)), | ||
| ("fp4 KV/fp8 all", dict(kv_bits=4, lm_head_bits=8, linear_attn_bits=8)), |
There was a problem hiding this comment.
These configurations are fed into the “Cheapest options clearing 100 tok/s” selection even though the checkpoint constants above identify the head and GDN weights as BF16 and this change implements only FP8 KV storage; the CLI does not support FP4 KV at all. As a result, the default report recommends a single B200 using fp4 KV/fp8 all at 157 tok/s, while the actually deployable FP8-KV configuration is reported at only 95 tok/s. Keep hypothetical modes out of winners or implement the corresponding quantization paths before presenting them as runnable options.
Useful? React with 👍 / 👎.
![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
22 issues found across 44 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="python/freetoken/core.py">
<violation number="1" location="python/freetoken/core.py:173">
P2: This new field changes the positional `Context` constructor and breaks callers that pass `moe_offload_cache` or `linear_state_pool` positionally. Make `compute_dtype` keyword-only (or append it after the existing positional fields) so those callers keep their previous bindings.</violation>
</file>
<file name="python/freetoken/engine/config.py">
<violation number="1" location="python/freetoken/engine/config.py:27">
P2: Adding these fields before the existing optional constructor fields breaks existing positional `EngineConfig(...)` callers: the old fourth argument now binds to `kv_cache_dtype`, and every later option shifts as well. Make the new options keyword-only or place them so existing positional parameters retain their original positions.</violation>
</file>
<file name="python/freetoken/attention/torch_native.py">
<violation number="1" location="python/freetoken/attention/torch_native.py:30">
P2: When the torch backend serves a full-attention model with sinks, this branch aborts every forward even though the backend is advertised as supporting full attention and consuming `AttentionSpec`. Implement sink handling in the portable attention path, or reject/avoid this backend for sink-bearing models during config validation.</violation>
</file>
<file name="python/freetoken/kvcache/mha_pool.py">
<violation number="1" location="python/freetoken/kvcache/mha_pool.py:139">
P2: When `k` already matches the cache dtype but `v` does not, this guard skips the cast and passes mixed-width tensors to `store_cache`. Check both inputs, or cast `v` independently, before the raw byte-copy.</violation>
</file>
<file name="python/freetoken/benchmark/xpu.py">
<violation number="1" location="python/freetoken/benchmark/xpu.py:134">
P3: The text output drops the dtype and workload geometry, so results from different `--memory-mib`, `--prefill-tokens`, or `--decode-context` runs cannot be interpreted from the captured report. Include the configuration fields in the human-readable output, as the JSON branch already does.</violation>
</file>
<file name="tests/moe/test_offload.py">
<violation number="1" location="tests/moe/test_offload.py:416">
P1: The fake accelerator_runtime returned here only defines get_device_name, but _adjust_config calls _rt.is_available() on it (engine.py:1317) before reading gpu_name. Because this test's config has moe_backend='auto' and is_moe=True, that branch is reached and the test fails with AttributeError. Add is_available=lambda: True to the SimpleNamespace.</violation>
</file>
<file name="deploy/capacity.py">
<violation number="1" location="deploy/capacity.py:115">
P2: `footprint_gb()` counts only one GDN state slot when deciding whether a GPU fits. Count the configured physical state-pool slots, including hybrid-radix snapshot slots, before declaring a deployment feasible.</violation>
<violation number="2" location="deploy/capacity.py:136">
P2: The plan recommends an unsupported `fp4 KV` configuration as the cheapest 100-token/s deployment. Remove this row from deployable results, or mark it hypothetical and exclude it from `winners` until a 4-bit KV backend exists.</violation>
</file>
<file name="python/freetoken/hardware.py">
<violation number="1" location="python/freetoken/hardware.py:38">
P3: For a non-B60 XPU, the report says the device model is unvalidated but still prescribes Arc B-series Level Zero settings. Emit this recommendation only when `is_b60` is true, or replace it with generic unsupported-device guidance.</violation>
<violation number="2" location="python/freetoken/hardware.py:45">
P2: When CUDA is selected, this branch prints `Status: dense Llama eager path available`, which is the XPU/B60 validation message. Use CUDA-specific or accelerator-neutral status text for CUDA reports.</violation>
<violation number="3" location="python/freetoken/hardware.py:70">
P1: On CUDA-only or CPU-only PyTorch builds that do not expose `torch.xpu`, `ft hardware` raises `AttributeError` before reporting accelerator availability. Guard the optional XPU runtime and make `inspect_hardware` treat a missing XPU as unavailable.</violation>
</file>
<file name="deploy/modal_app.py">
<violation number="1" location="deploy/modal_app.py:25">
P1: The documented `modal serve` command cannot expose an OpenAI-compatible endpoint because this app registers no web or ASGI endpoint. Add a Modal web endpoint that starts the FreeToken server, or remove this serving claim.</violation>
<violation number="2" location="deploy/modal_app.py:219">
P2: When model configuration validation fails, `validate` returns normally and reports success to the caller. Re-raise the exception after logging so this preflight actually blocks GPU deployment.</violation>
<violation number="3" location="deploy/modal_app.py:233">
P1: Both test wrappers interpolate caller-controlled `target` into `shell=True`, so a target such as `tests; <command>` executes arbitrary commands in the remote container. Pass pytest arguments as a list with `shell=False` in both functions.</violation>
<violation number="4" location="deploy/modal_app.py:337">
P2: Clear the prefix cache or use independent cache states before differencing these runs. The second call reuses the first call's cached prompt, so `t_full - t_prefill` does not isolate decode time and can yield a negative or `nan` throughput result.</violation>
<violation number="5" location="deploy/modal_app.py:340">
P2: The decode-rate figure this PR headlines is inflated because run(1) is the first forward pass on a cold engine. t_prefill therefore includes JIT compilation / CUDA-graph-capture / kernel-cache warmup, so decode_s = t_full - t_prefill over-subtracts and decode_tps = (gen-1)/decode_s comes out higher than the true steady-state rate. The docstring promises "what comes out is the steady-state interactive rate a user would feel", but the first measured prefill is not steady-state. Warm up the engine before timing (e.g. a discard generate or a min-of-N prefill) so the differencing removes real prefill time, not warmup.</violation>
<violation number="6" location="deploy/modal_app.py:407">
P2: The no-argument verification runs the known-clipping `fp8_e4m3` mode instead of the recommended `fp8_e5m2`, so it reports degraded generation by default. Change the default to `fp8_e5m2`.</violation>
</file>
<file name="tests/kvcache/test_kv_cache_dtype.py">
<violation number="1" location="tests/kvcache/test_kv_cache_dtype.py:105">
P2: test_store_kv_casts_into_a_narrower_pool and test_fp8_pool_allocates_half_the_bytes build MHAKVCache without initializing TP info, so get_tp_info() raises RuntimeError when this file runs standalone or before another test has set the global. The rest of tests/kvcache guards this with an _init_tp() that calls set_tp_info(rank=0, size=1) when try_get_tp_info() is None. Add that guard before both instantiations.</violation>
</file>
<file name="python/freetoken/models/llama/config.py">
<violation number="1" location="python/freetoken/models/llama/config.py:8">
P3: `_checkpoint_quantization` re-implements the extraction already in `detect_expert_quant` (models/config.py) instead of reusing it, and the two now diverge on edge cases: this one returns `"quantized"` when a quantization_config has no method/algo, and drops the fp4→`"nvfp4"` normalization that `detect_expert_quant` performs. Since `checkpoint_quantization` is only read as a `!= "none"` gate (accelerator.py XPU check), the string difference is currently inert, but the duplicated logic is a maintenance hazard. Reuse the shared helper (e.g. `detect_expert_quant` / a small wrapper) so the two detection paths stay consistent.</violation>
</file>
<file name="tests/test_xpu_llama_eager.py">
<violation number="1" location="tests/test_xpu_llama_eager.py:85">
P2: The final two asserts (`freetoken.kernel.triton.norm`/`rope` not in `sys.modules`) cannot hold on the CPU host this test targets: constructing `LlamaForCausalLM` earlier imports both modules via the `else` branches of `RMSNormFused.__init__` and `RotaryEmbedding.__init__` (taken whenever `torch.version.xpu` is unset and flashinfer is absent). The test will fail on a bare CPU image, contradicting the claimed pass, or the asserts are vacuous if flashinfer is present. Drop the two `sys.modules` asserts or restructure so construction does not defeat them.</violation>
</file>
<file name="python/freetoken/kernel/torch_ops.py">
<violation number="1" location="python/freetoken/kernel/torch_ops.py:13">
P2: In the portable `rmsnorm`, `x.float()` is materialized three times per call (once in `normalized`, twice inside `square().mean()` after `.float()` is applied again) plus `weight.float()` each call. On the XPU eager path this runs per layer per token and is bandwidth-bound, so these full-tensor dtype conversions add avoidable copy overhead. Convert to float once (or compute in the native dtype with `torch.rsqrt(x.square().mean(-1, keepdim=True) + eps)`) and reuse the result.</violation>
</file>
<file name="python/freetoken/engine/engine.py">
<violation number="1" location="python/freetoken/engine/engine.py:176">
P2: `_kv_is_quantized` treats any `kv_cache_dtype != dtype` as 'quantized', not just a narrower cache. A same-width cache (e.g. `--kv-cache-dtype bfloat16` with `--dtype float16`) is therefore forced onto the fi-only backend or rejected at config time, even though the cache is not narrower and `store_kv` already converts on write. Compare byte widths instead of dtype equality so only an actually-narrower cache triggers the `supports_quantized_kv` gate.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
|
|
||
| monkeypatch.setattr( | ||
| "freetoken.engine.engine.accelerator_runtime", | ||
| lambda _kind: SimpleNamespace(get_device_name=lambda _device: "test gpu"), |
There was a problem hiding this comment.
P1: The fake accelerator_runtime returned here only defines get_device_name, but _adjust_config calls _rt.is_available() on it (engine.py:1317) before reading gpu_name. Because this test's config has moe_backend='auto' and is_moe=True, that branch is reached and the test fails with AttributeError. Add is_available=lambda: True to the SimpleNamespace.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At tests/moe/test_offload.py, line 416:
<comment>The fake accelerator_runtime returned here only defines get_device_name, but _adjust_config calls _rt.is_available() on it (engine.py:1317) before reading gpu_name. Because this test's config has moe_backend='auto' and is_moe=True, that branch is reached and the test fails with AttributeError. Add is_available=lambda: True to the SimpleNamespace.</comment>
<file context>
@@ -404,18 +404,23 @@ def test_lru_gpu_cache_assigns_unique_slots_for_large_miss_batch():
+ monkeypatch.setattr(
+ "freetoken.engine.engine.accelerator_runtime",
+ lambda _kind: SimpleNamespace(get_device_name=lambda _device: "test gpu"),
+ )
+
</file context>
Suggested change
| lambda _kind: SimpleNamespace(get_device_name=lambda _device: "test gpu"), | |
| lambda _kind: SimpleNamespace(is_available=lambda: True, get_device_name=lambda _device: "test gpu"), |
| "--accelerator", choices=("auto", "cuda", "xpu"), default="auto" | ||
| ) | ||
| args = parser.parse_args(argv) | ||
| report = inspect_hardware(args.accelerator, torch.cuda, torch.xpu) |
There was a problem hiding this comment.
P1: On CUDA-only or CPU-only PyTorch builds that do not expose torch.xpu, ft hardware raises AttributeError before reporting accelerator availability. Guard the optional XPU runtime and make inspect_hardware treat a missing XPU as unavailable.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/freetoken/hardware.py, line 70:
<comment>On CUDA-only or CPU-only PyTorch builds that do not expose `torch.xpu`, `ft hardware` raises `AttributeError` before reporting accelerator availability. Guard the optional XPU runtime and make `inspect_hardware` treat a missing XPU as unavailable.</comment>
<file context>
@@ -0,0 +1,72 @@
+ "--accelerator", choices=("auto", "cuda", "xpu"), default="auto"
+ )
+ args = parser.parse_args(argv)
+ report = inspect_hardware(args.accelerator, torch.cuda, torch.xpu)
+ print(format_hardware_report(report))
+ return 0
</file context>
| modal run deploy/modal_app.py::validate # CPU-only config gate | ||
| modal run deploy/modal_app.py::smoke # GPU: load + generate | ||
| modal run deploy/modal_app.py::bench | ||
| modal serve deploy/modal_app.py # OpenAI-compatible endpoint |
There was a problem hiding this comment.
P1: The documented modal serve command cannot expose an OpenAI-compatible endpoint because this app registers no web or ASGI endpoint. Add a Modal web endpoint that starts the FreeToken server, or remove this serving claim.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At deploy/modal_app.py, line 25:
<comment>The documented `modal serve` command cannot expose an OpenAI-compatible endpoint because this app registers no web or ASGI endpoint. Add a Modal web endpoint that starts the FreeToken server, or remove this serving claim.</comment>
<file context>
@@ -0,0 +1,429 @@
+ modal run deploy/modal_app.py::validate # CPU-only config gate
+ modal run deploy/modal_app.py::smoke # GPU: load + generate
+ modal run deploy/modal_app.py::bench
+ modal serve deploy/modal_app.py # OpenAI-compatible endpoint
+"""
+
</file context>
| import subprocess | ||
|
|
||
| rc = subprocess.run( | ||
| f"cd /opt/freetoken && python -m pytest -q {target}", |
There was a problem hiding this comment.
P1: Both test wrappers interpolate caller-controlled target into shell=True, so a target such as tests; <command> executes arbitrary commands in the remote container. Pass pytest arguments as a list with shell=False in both functions.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At deploy/modal_app.py, line 233:
<comment>Both test wrappers interpolate caller-controlled `target` into `shell=True`, so a target such as `tests; <command>` executes arbitrary commands in the remote container. Pass pytest arguments as a list with `shell=False` in both functions.</comment>
<file context>
@@ -0,0 +1,429 @@
+ import subprocess
+
+ rc = subprocess.run(
+ f"cd /opt/freetoken && python -m pytest -q {target}",
+ shell=True, check=False,
+ ).returncode
</file context>
| # Model compute dtype. Distinct from `kv_cache.dtype`, which is the KV *storage* | ||
| # width and may be narrower (fp8); attention backends plan queries at this dtype | ||
| # while reading the cache at the pool's. | ||
| compute_dtype: torch.dtype = torch.bfloat16 |
There was a problem hiding this comment.
P2: This new field changes the positional Context constructor and breaks callers that pass moe_offload_cache or linear_state_pool positionally. Make compute_dtype keyword-only (or append it after the existing positional fields) so those callers keep their previous bindings.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/freetoken/core.py, line 173:
<comment>This new field changes the positional `Context` constructor and breaks callers that pass `moe_offload_cache` or `linear_state_pool` positionally. Make `compute_dtype` keyword-only (or append it after the existing positional fields) so those callers keep their previous bindings.</comment>
<file context>
@@ -167,6 +167,10 @@ def padded_size(self) -> int:
+ # Model compute dtype. Distinct from `kv_cache.dtype`, which is the KV *storage*
+ # width and may be narrower (fp8); attention backends plan queries at this dtype
+ # while reading the cache at the pool's.
+ compute_dtype: torch.dtype = torch.bfloat16
# NOTE: this table always treat page_size = 1
page_table: torch.Tensor = field(init=False)
</file context>
Suggested change
| compute_dtype: torch.dtype = torch.bfloat16 | |
| compute_dtype: torch.dtype = field(default=torch.bfloat16, kw_only=True) |
| t_full = run(gen) # prefill + gen decode steps | ||
|
|
||
| decode_s = t_full - t_prefill | ||
| decode_tps = (gen - 1) / decode_s if decode_s > 0 else float("nan") |
There was a problem hiding this comment.
P2: The decode-rate figure this PR headlines is inflated because run(1) is the first forward pass on a cold engine. t_prefill therefore includes JIT compilation / CUDA-graph-capture / kernel-cache warmup, so decode_s = t_full - t_prefill over-subtracts and decode_tps = (gen-1)/decode_s comes out higher than the true steady-state rate. The docstring promises "what comes out is the steady-state interactive rate a user would feel", but the first measured prefill is not steady-state. Warm up the engine before timing (e.g. a discard generate or a min-of-N prefill) so the differencing removes real prefill time, not warmup.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At deploy/modal_app.py, line 340:
<comment>The decode-rate figure this PR headlines is inflated because run(1) is the first forward pass on a cold engine. t_prefill therefore includes JIT compilation / CUDA-graph-capture / kernel-cache warmup, so decode_s = t_full - t_prefill over-subtracts and decode_tps = (gen-1)/decode_s comes out higher than the true steady-state rate. The docstring promises "what comes out is the steady-state interactive rate a user would feel", but the first measured prefill is not steady-state. Warm up the engine before timing (e.g. a discard generate or a min-of-N prefill) so the differencing removes real prefill time, not warmup.</comment>
<file context>
@@ -0,0 +1,429 @@
+ t_full = run(gen) # prefill + gen decode steps
+
+ decode_s = t_full - t_prefill
+ decode_tps = (gen - 1) / decode_s if decode_s > 0 else float("nan")
+ prefill_tps = ctx / t_prefill
+
</file context>
| return time.time() - t | ||
|
|
||
| t_prefill = run(1) # prefill + 1 decode step | ||
| t_full = run(gen) # prefill + gen decode steps |
There was a problem hiding this comment.
P2: Clear the prefix cache or use independent cache states before differencing these runs. The second call reuses the first call's cached prompt, so t_full - t_prefill does not isolate decode time and can yield a negative or nan throughput result.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At deploy/modal_app.py, line 337:
<comment>Clear the prefix cache or use independent cache states before differencing these runs. The second call reuses the first call's cached prompt, so `t_full - t_prefill` does not isolate decode time and can yield a negative or `nan` throughput result.</comment>
<file context>
@@ -0,0 +1,429 @@
+ return time.time() - t
+
+ t_prefill = run(1) # prefill + 1 decode step
+ t_full = run(gen) # prefill + gen decode steps
+
+ decode_s = t_full - t_prefill
</file context>
Comment on lines
+134
to
+144
| return "\n".join( | ||
| ( | ||
| f"Device: {result.device}", | ||
| f"PyTorch: {result.torch_version}", | ||
| f"Level Zero V2: {result.level_zero_v2}", | ||
| f"Level Zero immediate command lists: {result.queue_mode}", | ||
| f"Memory copy: {result.memory_copy_gbps:.2f} GB/s", | ||
| f"Prefill SDPA: {result.prefill_tokens_per_second:.1f} tokens/s", | ||
| f"Decode SDPA: {result.decode_steps_per_second:.1f} steps/s", | ||
| ) | ||
| ) |
There was a problem hiding this comment.
P3: The text output drops the dtype and workload geometry, so results from different --memory-mib, --prefill-tokens, or --decode-context runs cannot be interpreted from the captured report. Include the configuration fields in the human-readable output, as the JSON branch already does.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/freetoken/benchmark/xpu.py, line 134:
<comment>The text output drops the dtype and workload geometry, so results from different `--memory-mib`, `--prefill-tokens`, or `--decode-context` runs cannot be interpreted from the captured report. Include the configuration fields in the human-readable output, as the JSON branch already does.</comment>
<file context>
@@ -0,0 +1,178 @@
+def format_result(result: XpuBenchmarkResult, *, as_json: bool) -> str:
+ if as_json:
+ return json.dumps(asdict(result), sort_keys=True)
+ return "\n".join(
+ (
+ f"Device: {result.device}",
</file context>
Suggested change
| return "\n".join( | |
| ( | |
| f"Device: {result.device}", | |
| f"PyTorch: {result.torch_version}", | |
| f"Level Zero V2: {result.level_zero_v2}", | |
| f"Level Zero immediate command lists: {result.queue_mode}", | |
| f"Memory copy: {result.memory_copy_gbps:.2f} GB/s", | |
| f"Prefill SDPA: {result.prefill_tokens_per_second:.1f} tokens/s", | |
| f"Decode SDPA: {result.decode_steps_per_second:.1f} steps/s", | |
| ) | |
| ) | |
| return "\n".join( | |
| ( | |
| f"Device: {result.device}", | |
| f"PyTorch: {result.torch_version}", | |
| f"Level Zero V2: {result.level_zero_v2}", | |
| f"Level Zero immediate command lists: {result.queue_mode}", | |
| f"Dtype: {result.dtype}", | |
| f"Memory: {result.memory_mib} MiB", | |
| f"Prefill: {result.prefill_tokens} tokens", | |
| f"Decode context: {result.decode_context} tokens", | |
| f"Attention: {result.heads} heads x {result.head_dim} head dim", | |
| f"Warmup: {result.warmup}, iterations: {result.iterations}", | |
| f"Memory copy: {result.memory_copy_gbps:.2f} GB/s", | |
| f"Prefill SDPA: {result.prefill_tokens_per_second:.1f} tokens/s", | |
| f"Decode SDPA: {result.decode_steps_per_second:.1f} steps/s", | |
| ) | |
| ) |
| recommendations = ( | ||
| "Enable Resizable BAR in firmware.", | ||
| "Use current Intel compute drivers and the PyTorch XPU wheel.", | ||
| "Keep the Arc B-series Level Zero V2 default; benchmark the legacy adapter only if needed.", |
There was a problem hiding this comment.
P3: For a non-B60 XPU, the report says the device model is unvalidated but still prescribes Arc B-series Level Zero settings. Emit this recommendation only when is_b60 is true, or replace it with generic unsupported-device guidance.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/freetoken/hardware.py, line 38:
<comment>For a non-B60 XPU, the report says the device model is unvalidated but still prescribes Arc B-series Level Zero settings. Emit this recommendation only when `is_b60` is true, or replace it with generic unsupported-device guidance.</comment>
<file context>
@@ -0,0 +1,72 @@
+ recommendations = (
+ "Enable Resizable BAR in firmware.",
+ "Use current Intel compute drivers and the PyTorch XPU wheel.",
+ "Keep the Arc B-series Level Zero V2 default; benchmark the legacy adapter only if needed.",
+ )
+ return HardwareReport(kind, name, is_b60, recommendations)
</file context>
Comment on lines
+8
to
+18
| def _checkpoint_quantization(hf_config: Any) -> str: | ||
| quantization = getattr(hf_config, "quantization_config", None) | ||
| if quantization is None: | ||
| return "none" | ||
| if isinstance(quantization, dict): | ||
| method = quantization.get("quant_method") or quantization.get("quant_algo") | ||
| else: | ||
| method = getattr(quantization, "quant_method", None) or getattr( | ||
| quantization, "quant_algo", None | ||
| ) | ||
| return str(method or "quantized").lower() |
There was a problem hiding this comment.
P3: _checkpoint_quantization re-implements the extraction already in detect_expert_quant (models/config.py) instead of reusing it, and the two now diverge on edge cases: this one returns "quantized" when a quantization_config has no method/algo, and drops the fp4→"nvfp4" normalization that detect_expert_quant performs. Since checkpoint_quantization is only read as a != "none" gate (accelerator.py XPU check), the string difference is currently inert, but the duplicated logic is a maintenance hazard. Reuse the shared helper (e.g. detect_expert_quant / a small wrapper) so the two detection paths stay consistent.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At python/freetoken/models/llama/config.py, line 8:
<comment>`_checkpoint_quantization` re-implements the extraction already in `detect_expert_quant` (models/config.py) instead of reusing it, and the two now diverge on edge cases: this one returns `"quantized"` when a quantization_config has no method/algo, and drops the fp4→`"nvfp4"` normalization that `detect_expert_quant` performs. Since `checkpoint_quantization` is only read as a `!= "none"` gate (accelerator.py XPU check), the string difference is currently inert, but the duplicated logic is a maintenance hazard. Reuse the shared helper (e.g. `detect_expert_quant` / a small wrapper) so the two detection paths stay consistent.</comment>
<file context>
@@ -5,6 +5,19 @@
from freetoken.models.config import ModelConfig, RotaryConfig
+def _checkpoint_quantization(hf_config: Any) -> str:
+ quantization = getattr(hf_config, "quantization_config", None)
+ if quantization is None:
</file context>
Suggested change
| def _checkpoint_quantization(hf_config: Any) -> str: | |
| quantization = getattr(hf_config, "quantization_config", None) | |
| if quantization is None: | |
| return "none" | |
| if isinstance(quantization, dict): | |
| method = quantization.get("quant_method") or quantization.get("quant_algo") | |
| else: | |
| method = getattr(quantization, "quant_method", None) or getattr( | |
| quantization, "quant_algo", None | |
| ) | |
| return str(method or "quantized").lower() | |
| from freetoken.models.config import ModelConfig, RotaryConfig, detect_expert_quant | |
| def _checkpoint_quantization(hf_config: Any) -> str: | |
| return detect_expert_quant(hf_config) | |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Goal
Serve
Blackfrost-Research/M.O.G.-SEC-27B-1M-CTX-NVFP4with FreeToken on Modal at 100 decode tok/s with a 1M-token context, cheaply.
What I found
FreeToken already supports this architecture —
Qwen3_5ForConditionalGenerationisregistered and maps onto the existing
qwen3_5_moepackage's dense path. No model codewas needed. Verified against the real checkpoint on Modal:
The real problem was bandwidth, not support. Decode re-reads every active weight plus
the whole KV cache each step, so
tok/s ≈ bandwidth / bytes-per-token. The model isdense (no MoE sparsity to exploit), and at 1M context with a bf16 cache that is
91.6 GB/token — 9.2 TB/s to hit 100 tok/s, which no single GPU has.
The hybrid architecture is what makes it tractable at all: only 16 of 64 layers carry a
KV cache (
full_attention_interval: 4); the other 48 GatedDeltaNet layers hold a fixed155 MB state independent of context length.
The lever:
--kv-cache-dtypeThe KV cache was hard-wired to the compute dtype (
engine.pypassedself.dtypestraight to the pool). Decoupling it and storing fp8 halves the KV half of the budget:
A single B200 at $6.25/hr (~$16/M output tokens) is the cheapest configuration that
clears 100 tok/s, and only with an fp8 cache.
deploy/capacity.pyregenerates thistable across every Modal SKU.
Measured, not just modelled
modal run deploy/modal_app.py::verify_fp8_kvon an L4 with Qwen3-0.6B — same 18.44 GiBKV budget both runs:
fp8_e5m2fp8_e4m3Exactly 2× the context in the same memory, and e5m2's greedy continuation was identical
to bf16 — not merely close.
fp8_e4m3is not recommended and now warns. The pool stores K/V by a straight castwith no per-tensor scale, and e4m3 saturates at ±448, so activations clip and generation
collapses into repeated punctuation. e5m2 spends the same byte on exponent range instead
of mantissa and survives an unscaled cast. Both are one byte, so e4m3 buys no extra
bandwidth — it would only win with calibrated
k_scale/v_scaleplumbed to the kernel,which is not implemented.
Changes
EngineConfig.kv_cache_dtype(None= follow--dtype, so existing behaviour isbit-identical) with a resolved
kv_dtypeused as the single source of truth by boththe pool allocation and the byte budget — planning and allocation cannot disagree.
Context.compute_dtypelets the FlashInfer backend planq_data_typeat compute widthwhile
kv_data_typefollows the pool. Queries are never quantized.MHAKVCache.store_kvnarrows K/V on write (store_cacheis a raw byte-copy and cannotconvert); the torch fallback aliases fp8 through uint8 because
index_copy_has no CPUkernel for float8.
BackendInfo.supports_quantized_kvgates the feature: an unsupporting backend isrejected at config time rather than failing inside a kernel launch, and
--attention-backend autoskips those backends.torch.cuda.get_device_name(0) if torch.cuda.is_available() else Nonewith anunconditional call, which raised out of a driver probe on CPU-only hosts — breaking
config-time runs and the test suite off-GPU. Guard restored.
deploy/— Modal app (CPU-built image, volume-cached weights + JIT artifacts,CPU-only download/validate gates before any GPU spend) and
capacity.py, the costmodel behind the numbers above.
docs/modal-deployment.md.Testing
337 passed, 6 skipped, 0 failedon the Modal image (CPU), plus a GPU pass. Newtests/kvcache/test_kv_cache_dtype.pycovers dtype resolution, budget halving, thenarrowing store for both fp8 dtypes, the backend capability gate, and the e5m2-vs-e4m3
range rationale.
The Modal account has no payment method, so B200/H200/H100/A100/L40S are all blocked —
only T4, L4 and A10 are available, none of which can hold a 29.4 GB checkpoint. The
weights are downloaded to the volume and the config path is validated, but
smokeandbenchhave not run against the 27B, so the 100 tok/s figure is modelled, notmeasured. Add a payment method and:
The fp8-KV win itself is measured (table above) — just on a small model on an L4.
🤖 Generated with Claude Code
Summary by cubic
Adds fp8 KV cache support and a Modal deployment for M.O.G.-SEC-27B-1M-CTX, targeting 100 decode tok/s at 1M-token context. Storing the KV cache at fp8_e5m2 halves its bandwidth cost, the main decode bottleneck, and keeps generation byte-identical to bf16 on tested models.
Changes
--kv-cache-dtype(defaults to--dtype) and decouples the cache pool, byte budget, and FlashInfer query planning;fp8_e5m2doubles context capacity for the same memory.deploy/with a Modal app and capacity model; verified on L4/Qwen3-0.6B, not yet run end-to-end on the 27B because the Modal account lacks a payment method for large GPUs.setup.pyand the engine with an accelerator abstraction (CUDA/XPU) and an eager PyTorch path for dense Llama on Intel Arc, plus a hardware probe command.Limitations
fp8_e4m3is not recommended; without per-tensor scales it clips activations and collapses generation.Written for commit 3f81bba. Summary will update on new commits.