Core changes that were missing from the original PR:
1. unsloth/models/llama.py:
- Add _ensure_cache_is_dynamic() to convert tuple/list KV caches
to DynamicCache for transformers v5 compatibility
- Add _slice_position_ids() to handle position_ids slicing when
input_ids is trimmed to uncached tokens
- Fix unsloth_fast_generate: skip setting cache_implementation
when user provides past_key_values (avoids ValueError on
transformers >= 4.57)
- Fix _fast_prepare_inputs_for_generation: when past_key_values
covers fewer tokens than input_ids, keep only the uncached
portion instead of always slicing to last token
- Fix CausalLM_fast_forward: add input_ids.shape[1] == 1 guard
so multi-token prefill with external KV cache falls through
to the regular model forward path
2. unsloth/models/mistral.py:
- Same input_ids.shape[1] == 1 guard and multi-token fallback
as CausalLM_fast_forward
3. tests/test_past_kv_models.py:
- Extract model loading into _load_model() with proper SkipTest
handling (only skips on loading errors, not on test failures)
- Remove blanket try/except that was masking real failures
- Remove exact output match assertion (4-bit quantization with
different computation paths can produce slightly different but
equally valid outputs)
4. tests/test_past_kv_utils.py:
- Fix comment to accurately describe the inlined functions
Tested: Llama-3.2-1B-Instruct, Qwen3-0.6B, gemma-2-2b-it
All 4 integration tests pass, 14 unit tests pass, benchmark runs,
and standard training (Gemma2, Llama) is not regressed.
- Add examples/kv_cache_multiturn_benchmark.py: realistic 15-turn tech
support conversation (~1500 history tokens) with scaling test at 4, 8,
12, and 16 messages to show speedup increasing with history length
- Add tests/test_past_kv_models.py: GPU integration tests for Llama,
Qwen3, and Gemma2 past_key_values generation plus tuple-KV v5 compat
- Add tests/test_past_kv_utils.py: CPU-only unit tests for
_ensure_cache_is_dynamic and _slice_position_ids utilities
- Add _ensure_cache_is_dynamic to convert tuple/list KV caches to DynamicCache
before transformers v5's _get_cache rejects them
- Wrap generate() via fix_prepare_inputs_for_generation for all model types
- Fix Qwen3 RoPE: index cos/sin by position_ids (both branches were identical)
- Fix Gemma2 softcapping attention: handle 4D masks and Q_len != K_len during
prefill with past_key_values
- Add _slice_position_ids shared utility, replace inline duplication in
PeftModel_fast_forward, MistralForCausalLM_fast_forward, and
CausalLM_fast_forward (covers Llama/Qwen3/Gemma2)
- Remove redundant seq_len assignment in _fast_prepare_inputs_for_generation
- Add unit tests (test_past_kv_utils.py) and GPU integration tests
(test_past_kv_models.py) for Llama, Qwen3, and Gemma2
* Update CODEOWNERS for studio and cli
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
* packing optimziation with cache to reduce D2H copy
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* cache per device to avoid race condition for multi-gpu
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* add cache freeing up func
---------
Co-authored-by: ruixiangw <ruixiangw@nvidia.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: ruixiang <wangruixiang07@outlook.com>
* Rebuild Studio branch on top of main
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Fix security and code quality issues for Studio PR #4237
- Validate models_dir query param against allowed directory roots
to prevent path traversal in /api/models/local endpoint
- Replace string startswith() with Path.is_relative_to() for
frontend path traversal check in serve_frontend
- Sanitize SSE error messages to not leak exception details to
clients (4 locations in inference.py)
- Bind port-discovery socket to 127.0.0.1 instead of all interfaces
in llama_cpp backend
- Import datasets_root and resolve_output_dir in embedding training
function to fix NameError and use managed output directory
- Remove stale .gitignore entries for package-lock.json and test
directories so tests can be tracked in version control
- Add venv-reexecution logic to ui CLI command matching the studio
command behavior
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Move models_dir path validation before try/except block
The HTTPException(403) was inside the try/except Exception handler,
so it would be caught and re-raised as a 500. Moving the validation
before the try block ensures the 403 is returned directly and also
makes the control flow clearer for static analysis (path is validated
before any filesystem operations).
* Use os.path.realpath + startswith for models_dir validation
CodeQL py/path-injection does not recognize Path.is_relative_to() as
a sanitizer. Switched to os.path.realpath + str.startswith which is
a recognized sanitizer pattern in CodeQL's taint analysis. The
startswith check uses root_str + os.sep to prevent prefix collisions
(e.g. /app/models_evil matching /app/models).
* Never pass user input to Path constructor in models_dir validation
CodeQL traces taint through Path(resolved) even after a startswith
barrier guard. Fix: the user-supplied models_dir is only used as a
string for comparison against allowed roots. The Path object passed
to _scan_models_dir comes from the trusted allowed_roots list, not
from user input. This fully breaks the taint chain.
---------
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
- Workers now compute backend_path and venv_t5 locally via Path(__file__)
- Moved .venv_t5 to ~/.unsloth/studio/.venv_t5
- Added ensure_studio_directories() call on server startup
- Expanded CLI studio command into sub-app with setup subcommand
Tier 1 check-format was picking images.zip over testmini.parquet,
causing wrong columns (image/label) and broken VLM mapping.
Also log first VLM conversion failure instead of swallowing silently.
